diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..531534f --- /dev/null +++ b/.env.example @@ -0,0 +1,29 @@ +# ========================================== +# YiGraph 环境变量配置模板 +# 复制此文件为 .env 并填入实际值: +# cp .env.example .env +# ========================================== + +# ---------- LLM API Keys ---------- +# 百炼 DashScope API Key（用于通义系列模型） +DASHSCOPE_API_KEY=your_dashscope_key_here + +# OpenAI 兼容 API Key（用于 GPT 系列模型及兼容接口） +OPENAI_API_KEY=your_openai_key_here + +# ---------- 数据库密码 ---------- +# Neo4j 图数据库密码 +NEO4J_PASSWORD=your_neo4j_password + +# NebulaGraph 图数据库密码 +NEBULA_PASSWORD=your_nebula_password + +# ---------- Web 安全 ---------- +# Flask 调试模式（生产环境必须设为 0） +FLASK_DEBUG=0 + +# CORS 允许的来源地址（前端部署域名） +CORS_ORIGINS=http://localhost:5089 + +# WebSocket 认证 Token +YIGRAPH_WS_TOKEN=your_websocket_token diff --git a/.gitignore b/.gitignore index 1e2b585..5a3101d 100644 --- a/.gitignore +++ b/.gitignore @@ -147,3 +147,13 @@ docs/_build # database *.db + +# ========================================== +# 含敏感凭证的配置文件（禁止提交真实密钥） +# ========================================== +web/frontend/models.json +aag/expert_search_engine/data_process/openai_extractor/openai_config.json +aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json + +# 示例模板（允许提交） +!*.example diff --git a/aag/api/services/chat_service.py b/aag/api/services/chat_service.py index 93df5ae..f025b8c 100644 --- a/aag/api/services/chat_service.py +++ b/aag/api/services/chat_service.py @@ -5,11 +5,6 @@ import logging import asyncio from typing import Dict, Any, Optional, Callable, List -import sys -import os - -# Add project root to path -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))) from .engine_service import EngineService @@ -36,14 +31,14 @@ class ChatService: def __init__(self): """Initialize chat service.""" self.engine_service = EngineService.get_instance() - + async def process_message( self, message: str, model: Optional[str] = None, dataset: Optional[str] = None, mode: str = "normal", - expert_mode: bool = False + expert_mode: bool = False, ) -> Dict[str, Any]: """ Process a user message (sync-style API). @@ -68,26 +63,19 @@ async def process_message( logger.info(f"Select dataset: {dataset}") engine.specific_dataset(dataset) - logger.info(f"Process message: mode={mode}, model={model}, message={message[:50]}...") + logger.info( + f"Process message: mode={mode}, model={model}, message={message[:50]}..." + ) result = await engine.run(message, mode=mode) - return { - "success": True, - "result": result, - "mode": mode - } + return {"success": True, "result": result, "mode": mode} except Exception as e: logger.error(f"Process message failed: {str(e)}") - return { - "success": False, - "error": str(e), - "mode": mode - } - + return {"success": False, "error": str(e), "mode": mode} + async def process_dag_modification( - self, - modification_request: str + self, modification_request: str ) -> Dict[str, Any]: """ Handle DAG modification request (expert mode). @@ -104,17 +92,11 @@ async def process_dag_modification( result = await engine.expert_modify_dag(modification_request) - return { - "success": True, - **result - } + return {"success": True, **result} except Exception as e: logger.error(f"DAG modification failed: {str(e)}") - return { - "success": False, - "error": str(e) - } - + return {"success": False, "error": str(e)} + async def start_expert_analysis(self) -> Dict[str, Any]: """ Start expert-mode analysis. @@ -128,17 +110,11 @@ async def start_expert_analysis(self) -> Dict[str, Any]: result = await engine.expert_start_analysis() - return { - "success": True, - "result": result - } + return {"success": True, "result": result} except Exception as e: logger.error(f"Expert analysis failed: {str(e)}") - return { - "success": False, - "error": str(e) - } - + return {"success": False, "error": str(e)} + async def process_streaming_chat( self, message: str, @@ -147,7 +123,7 @@ async def process_streaming_chat( dataset_type: Optional[str] = None, mode: str = "normal", expert_mode: bool = False, - callback: Optional[Callable[[Dict[str, Any]], None]] = None + callback: Optional[Callable[[Dict[str, Any]], None]] = None, ): """ Stream chat (for WebSocket). @@ -168,11 +144,13 @@ async def process_streaming_chat( engine = self.engine_service.get_engine() if callback: - callback({ - "type": "thinking", - "contentType": "text", - "content": "Analyzing your question..." - }) + callback( + { + "type": "thinking", + "contentType": "text", + "content": "Analyzing your question...", + } + ) if dataset: logger.info(f"Select dataset: {dataset}, type: {dataset_type}") @@ -182,71 +160,79 @@ async def process_streaming_chat( result = await engine.run(message, mode=mode, callback=callback) - if mode in {"interact", "expert"} and isinstance(result, dict) and "dag_info" in result: + if ( + mode in {"interact", "expert"} + and isinstance(result, dict) + and "dag_info" in result + ): dag_content = self._convert_dag_to_frontend_format(result) if callback: - callback({ - "type": "result", - "contentType": "dag", - "content": dag_content - }) - elif mode in {"interact", "expert"} and isinstance(result, dict) and "error" in result: + callback( + {"type": "result", "contentType": "dag", "content": dag_content} + ) + elif ( + mode in {"interact", "expert"} + and isinstance(result, dict) + and "error" in result + ): if callback: - callback({ - "type": "result", - "contentType": "text", - "content": CHAT_FRIENDLY_ERROR_MSG - }) + callback( + { + "type": "result", + "contentType": "text", + "content": CHAT_FRIENDLY_ERROR_MSG, + } + ) elif mode == "normal" and isinstance(result, dict) and "dag_info" in result: # Normal mode: result has DAG info — send DAG first - dag_content = self._convert_dag_to_frontend_format({ - "dag_info": result.get("dag_info", {}) - }) + dag_content = self._convert_dag_to_frontend_format( + {"dag_info": result.get("dag_info", {})} + ) if callback: - callback({ - "type": "result", - "contentType": "dag", - "content": dag_content - }) + callback( + {"type": "result", "contentType": "dag", "content": dag_content} + ) result_text = result.get("analysis_result", "") if callback and result_text: - paragraphs = [p.strip() for p in result_text.split('\n') if p.strip()] + paragraphs = [ + p.strip() for p in result_text.split("\n") if p.strip() + ] for para in paragraphs: - callback({ - "type": "result", - "contentType": "text", - "content": para - }) + callback( + {"type": "result", "contentType": "text", "content": para} + ) else: result_text = str(result) if result else "No result." if _is_scheduler_error_text(result_text): result_text = CHAT_FRIENDLY_ERROR_MSG if callback: - paragraphs = [p.strip() for p in result_text.split('\n') if p.strip()] + paragraphs = [ + p.strip() for p in result_text.split("\n") if p.strip() + ] for para in paragraphs: - callback({ - "type": "result", - "contentType": "text", - "content": para - }) + callback( + {"type": "result", "contentType": "text", "content": para} + ) if callback: - callback({ - "type": "stream_end" - }) - + callback({"type": "stream_end"}) + except Exception as e: logger.error(f"Stream chat failed: {str(e)}", exc_info=True) if callback: - callback({ - "type": "result", - "contentType": "text", - "content": CHAT_FRIENDLY_ERROR_MSG - }) + callback( + { + "type": "result", + "contentType": "text", + "content": CHAT_FRIENDLY_ERROR_MSG, + } + ) callback({"type": "stream_end"}) return - - def _convert_dag_to_frontend_format(self, dag_result: Dict[str, Any]) -> Dict[str, Any]: + + def _convert_dag_to_frontend_format( + self, dag_result: Dict[str, Any] + ) -> Dict[str, Any]: """ Convert DAG result to frontend format. @@ -277,11 +263,13 @@ def _convert_dag_to_frontend_format(self, dag_result: Dict[str, Any]) -> Dict[st else: tasktype_str = "Unknown" - nodes.append({ - "id": str(step_id), - "label": step_info.get("question", ""), - "tasktype": tasktype_str - }) + nodes.append( + { + "id": str(step_id), + "label": step_info.get("question", ""), + "tasktype": tasktype_str, + } + ) edges_info = dag_info.get("edges", []) if edges_info: @@ -289,26 +277,22 @@ def _convert_dag_to_frontend_format(self, dag_result: Dict[str, Any]) -> Dict[st else: topological_order = dag_info.get("topological_order", []) for i in range(len(topological_order) - 1): - edges.append({ - "from": str(topological_order[i]), - "to": str(topological_order[i + 1]) - }) - - return { - "nodes": nodes, - "edges": edges - } + edges.append( + { + "from": str(topological_order[i]), + "to": str(topological_order[i + 1]), + } + ) + + return {"nodes": nodes, "edges": edges} except Exception as e: logger.error(f"Convert DAG format failed: {str(e)}") - return { - "nodes": [], - "edges": [] - } - + return {"nodes": [], "edges": []} + def process_dag_confirmation( self, dag_confirm: str, - callback: Optional[Callable[[Dict[str, Any]], None]] = None + callback: Optional[Callable[[Dict[str, Any]], None]] = None, ) -> Dict[str, Any]: """ Handle DAG confirmation (yes/no). @@ -324,12 +308,11 @@ def process_dag_confirmation( return self.start_expert_analysis() else: if callback: - callback({ - "type": "result", - "contentType": "text", - "content": "DAG rejected. Please modify and resubmit." - }) - return { - "success": True, - "message": "DAG rejected" - } + callback( + { + "type": "result", + "contentType": "text", + "content": "DAG rejected. Please modify and resubmit.", + } + ) + return {"success": True, "message": "DAG rejected"} diff --git a/aag/api/services/dataset_service.py b/aag/api/services/dataset_service.py index 87df29a..75df664 100644 --- a/aag/api/services/dataset_service.py +++ b/aag/api/services/dataset_service.py @@ -6,29 +6,24 @@ import logging import asyncio import json -import sys -import os from typing import Dict, Any, List, Optional -# 添加项目路径 -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))) - -from aag.api.DocumentAPI import DocumentAPIServer +from ..DocumentAPI import DocumentAPIServer logger = logging.getLogger(__name__) class CollectingSocket: """用于收集DocumentAPI响应的Socket类""" - + def __init__(self, msgs, collector): self.msgs = msgs self.collector = collector - + async def send(self, msg): self.collector.add_response(msg) logger.debug(f"[DatasetService] Response: {msg}") - + async def __aiter__(self): for m in self.msgs: yield m @@ -39,51 +34,49 @@ class DatasetService: 数据集管理服务 - 封装DocumentAPI 提供知识库的创建、删除、查询等功能 """ - + def __init__(self): """初始化数据集服务""" self.document_api = DocumentAPIServer() - self._response_cache = {} # 用于存储异步响应结果 - + # 缓存策略：当前通过 ResponseCollector 实时收集响应， + # 如需跨请求缓存结果，可在此初始化持久化缓存（如 LRU Cache）。 + async def create_knowledge_base( - self, - file_path: str, - graph_name: str, - db_name: Optional[str] = None + self, file_path: str, graph_name: str, db_name: Optional[str] = None ) -> Dict[str, Any]: """ 创建知识库（从文本文件） - + Args: file_path: 文件路径 graph_name: 图名称 db_name: 数据库名称（可选） - + Returns: 创建结果字典 """ try: logger.info(f"创建知识库: graph_name={graph_name}, file_path={file_path}") - + message = { "action": "create_kb", "file_path": file_path, - "graph_name": graph_name + "graph_name": graph_name, } - + if db_name: message["db_name"] = db_name - + # 创建响应收集器 response_collector = ResponseCollector() dummy_socket = CollectingSocket([json.dumps(message)], response_collector) - + # 执行处理 await self.document_api.handler(dummy_socket) - + # 获取响应 responses = response_collector.get_responses() - + # 解析响应 for response in responses: try: @@ -91,66 +84,52 @@ async def create_knowledge_base( if data.get("type") == "data" and data.get("contentType") == "json": content = data.get("content", {}) if content.get("success"): - return { - "success": True, - "data": content.get("data", {}) - } + return {"success": True, "data": content.get("data", {})} elif data.get("type") == "error": return { "success": False, - "error": data.get("content", "未知错误") + "error": data.get("content", "未知错误"), } except json.JSONDecodeError: continue - - return { - "success": True, - "message": "知识库创建请求已提交" - } - + + return {"success": True, "message": "知识库创建请求已提交"} + except Exception as e: logger.error(f"创建知识库失败: {str(e)}") - return { - "success": False, - "error": str(e) - } - + return {"success": False, "error": str(e)} + async def delete_knowledge_base( - self, - graph_name: str, - db_name: Optional[str] = None + self, graph_name: str, db_name: Optional[str] = None ) -> Dict[str, Any]: """ 删除知识库 - + Args: graph_name: 图名称 db_name: 数据库名称（可选） - + Returns: 删除结果字典 """ try: logger.info(f"删除知识库: graph_name={graph_name}") - - message = { - "action": "delete_kb", - "graph_name": graph_name - } - + + message = {"action": "delete_kb", "graph_name": graph_name} + if db_name: message["db_name"] = db_name - + # 创建响应收集器 response_collector = ResponseCollector() dummy_socket = CollectingSocket([json.dumps(message)], response_collector) - + # 执行处理 await self.document_api.handler(dummy_socket) - + # 获取响应 responses = response_collector.get_responses() - + # 解析响应 for response in responses: try: @@ -158,61 +137,46 @@ async def delete_knowledge_base( if data.get("type") == "data" and data.get("contentType") == "json": content = data.get("content", {}) if content.get("success"): - return { - "success": True, - "data": content.get("data", {}) - } + return {"success": True, "data": content.get("data", {})} elif data.get("type") == "error": return { "success": False, - "error": data.get("content", "未知错误") + "error": data.get("content", "未知错误"), } except json.JSONDecodeError: continue - - return { - "success": True, - "message": "知识库删除请求已提交" - } - + + return {"success": True, "message": "知识库删除请求已提交"} + except Exception as e: logger.error(f"删除知识库失败: {str(e)}") - return { - "success": False, - "error": str(e) - } - - async def get_triplets( - self, - graph_name: str - ) -> Dict[str, Any]: + return {"success": False, "error": str(e)} + + async def get_triplets(self, graph_name: str) -> Dict[str, Any]: """ 获取知识库的三元组 - + Args: graph_name: 图名称 - + Returns: 三元组列表 """ try: logger.info(f"获取三元组: graph_name={graph_name}") - - message = { - "action": "get_triplets", - "graph_name": graph_name - } - + + message = {"action": "get_triplets", "graph_name": graph_name} + # 创建响应收集器 response_collector = ResponseCollector() dummy_socket = CollectingSocket([json.dumps(message)], response_collector) - + # 执行处理 await self.document_api.handler(dummy_socket) - + # 获取响应 responses = response_collector.get_responses() - + # 解析响应 for response in responses: try: @@ -220,30 +184,21 @@ async def get_triplets( if data.get("type") == "data" and data.get("contentType") == "json": content = data.get("content", {}) if content.get("success"): - return { - "success": True, - "data": content.get("data", []) - } + return {"success": True, "data": content.get("data", [])} elif data.get("type") == "error": return { "success": False, - "error": data.get("content", "未知错误") + "error": data.get("content", "未知错误"), } except json.JSONDecodeError: continue - - return { - "success": False, - "error": "未获取到响应" - } - + + return {"success": False, "error": "未获取到响应"} + except Exception as e: logger.error(f"获取三元组失败: {str(e)}") - return { - "success": False, - "error": str(e) - } - + return {"success": False, "error": str(e)} + async def create_kb_from_graph( self, graph_name: str, @@ -254,11 +209,11 @@ async def create_kb_from_graph( vertex_file: Optional[str] = None, vertex_id_field: Optional[str] = None, vertex_name_field: Optional[str] = None, - weight_field: Optional[str] = None + weight_field: Optional[str] = None, ) -> Dict[str, Any]: """ 从图文件创建知识库 - + Args: graph_name: 图名称 edge_file: 边文件路径 @@ -269,40 +224,40 @@ async def create_kb_from_graph( vertex_id_field: 顶点ID字段名（可选） vertex_name_field: 顶点名称字段名（可选） weight_field: 权重字段名（可选） - + Returns: 创建结果字典 """ try: logger.info(f"从图文件创建知识库: graph_name={graph_name}") - + message = { "action": "upload_graph", "graph_name": graph_name, "edge_file": edge_file, "source_field": source_field, "target_field": target_field, - "relation_field": relation_field + "relation_field": relation_field, } - + if vertex_file: message["vertex_file"] = vertex_file message["vertex_id_field"] = vertex_id_field message["vertex_name_field"] = vertex_name_field - + if weight_field: message["weight_field"] = weight_field - + # 创建响应收集器 response_collector = ResponseCollector() dummy_socket = CollectingSocket([json.dumps(message)], response_collector) - + # 执行处理 await self.document_api.handler(dummy_socket) - + # 获取响应 responses = response_collector.get_responses() - + # 解析响应 for response in responses: try: @@ -310,46 +265,36 @@ async def create_kb_from_graph( if data.get("type") == "data" and data.get("contentType") == "json": content = data.get("content", {}) if content.get("success"): - return { - "success": True, - "data": content.get("data", {}) - } + return {"success": True, "data": content.get("data", {})} elif data.get("type") == "error": return { "success": False, - "error": data.get("content", "未知错误") + "error": data.get("content", "未知错误"), } except json.JSONDecodeError: continue - - return { - "success": True, - "message": "知识库创建请求已提交" - } - + + return {"success": True, "message": "知识库创建请求已提交"} + except Exception as e: logger.error(f"从图文件创建知识库失败: {str(e)}") - return { - "success": False, - "error": str(e) - } + return {"success": False, "error": str(e)} class ResponseCollector: """用于收集异步响应的辅助类""" - + def __init__(self): self.responses = [] - + def add_response(self, response: str): """添加响应""" self.responses.append(response) - + def get_responses(self) -> List[str]: """获取所有响应""" return self.responses - + def clear(self): """清空响应""" self.responses = [] - diff --git a/aag/api/services/engine_service.py b/aag/api/services/engine_service.py index 42c1e90..a564639 100644 --- a/aag/api/services/engine_service.py +++ b/aag/api/services/engine_service.py @@ -5,9 +5,14 @@ import logging import sys import os +import threading from typing import Optional -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))) +sys.path.append( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ) +) from aag.engine.aag_engine import AAGEngine from aag.config.engine_config import load_config_from_yaml @@ -20,31 +25,39 @@ class EngineService: """ AAG engine management service (singleton); manages engine lifecycle. """ - - _instance: Optional['EngineService'] = None + + _instance: Optional["EngineService"] = None _engine: Optional[AAGEngine] = None _initialized: bool = False - + _lock: threading.Lock = threading.Lock() # 线程安全锁，保护单例创建 + def __init__(self): """Private constructor; use get_instance() to obtain the instance.""" if EngineService._instance is not None: - raise RuntimeError("EngineService is a singleton. Use get_instance() to get the instance.") + raise RuntimeError( + "EngineService is a singleton. Use get_instance() to get the instance." + ) self._config_path = DEFAULT_CONFIG_PATH - + @classmethod - def get_instance(cls) -> 'EngineService': - """Return the engine service singleton.""" + def get_instance(cls) -> "EngineService": + """返回引擎服务单例（线程安全）。""" if cls._instance is None: - cls._instance = cls() + with cls._lock: + # 双重检查：锁内再次确认，防止竞态条件 + if cls._instance is None: + cls._instance = cls() return cls._instance - + def get_engine(self) -> AAGEngine: """Get or initialize the AAG engine.""" if self._engine is None or not self._initialized: try: logger.info(f"Initializing AAG engine; config: {self._config_path}") if not self._config_path.exists(): - raise FileNotFoundError(f"Config file not found: {self._config_path}") + raise FileNotFoundError( + f"Config file not found: {self._config_path}" + ) config = load_config_from_yaml(str(self._config_path)) self._engine = AAGEngine(config) self._initialized = True @@ -52,18 +65,19 @@ def get_engine(self) -> AAGEngine: except Exception as e: logger.error(f"AAG engine initialization failed: {str(e)}") raise - + return self._engine - + def set_config_path(self, config_path: str): """Set config file path (call before initializing).""" from pathlib import Path + self._config_path = Path(config_path) if self._initialized: logger.warning("Config path changed; engine will need re-initialization") self._initialized = False self._engine = None - + async def shutdown(self): """Shut down the engine and release resources.""" if self._engine is not None: @@ -76,14 +90,13 @@ async def shutdown(self): except Exception as e: logger.error(f"Error shutting down AAG engine: {str(e)}") raise - + def is_initialized(self) -> bool: """Return True if the engine is initialized.""" return self._initialized and self._engine is not None - + def reset(self): """Reset the engine (for tests or re-initialization).""" self._engine = None self._initialized = False logger.info("Engine service reset") - diff --git a/aag/computing_engine/code_executor.py b/aag/computing_engine/code_executor.py index 009fb8e..0f82590 100644 --- a/aag/computing_engine/code_executor.py +++ b/aag/computing_engine/code_executor.py @@ -13,54 +13,84 @@ logger = logging.getLogger(__name__) +# 安全内置函数白名单（用于 exec() 沙箱，禁止 __import__/open/eval/exec 等危险函数） +SAFE_BUILTINS = { + "len": len, + "sorted": sorted, + "reversed": reversed, + "enumerate": enumerate, + "zip": zip, + "map": map, + "filter": filter, + "sum": sum, + "min": min, + "max": max, + "abs": abs, + "round": round, + "int": int, + "float": float, + "str": str, + "bool": bool, + "list": list, + "dict": dict, + "set": set, + "tuple": tuple, + "range": range, + "any": any, + "all": all, + "isinstance": isinstance, + "type": type, + "print": print, +} + + class DynamicCodeExecutor: - def __init__(self, timeout: int = 120, auto_install: bool = True): self.timeout = timeout self.auto_install = auto_install self.installed_packages = set() - + def extract_imports(self, code: str) -> list: """提取代码中的 import 语句""" imports = [] - pattern1 = r'^\s*import\s+([a-zA-Z0-9_]+)' - pattern2 = r'^\s*from\s+([a-zA-Z0-9_]+)\s+import' - - for line in code.split('\n'): + pattern1 = r"^\s*import\s+([a-zA-Z0-9_]+)" + pattern2 = r"^\s*from\s+([a-zA-Z0-9_]+)\s+import" + + for line in code.split("\n"): match1 = re.match(pattern1, line) match2 = re.match(pattern2, line) if match1: imports.append(match1.group(1)) elif match2: imports.append(match2.group(1)) - + return imports - + def install_package(self, package_name: str) -> bool: """安装 Python 包""" if package_name in self.installed_packages: logger.info(f"📦 {package_name} 已安装，跳过") return True - + # 包名映射 package_mapping = { - 'cv2': 'opencv-python', - 'sklearn': 'scikit-learn', - 'PIL': 'Pillow', - 'yaml': 'pyyaml', + "cv2": "opencv-python", + "sklearn": "scikit-learn", + "PIL": "Pillow", + "yaml": "pyyaml", } - + install_name = package_mapping.get(package_name, package_name) - + try: logger.info(f"📥 正在安装 {install_name}...") result = subprocess.run( - [sys.executable, '-m', 'pip', 'install', install_name], + [sys.executable, "-m", "pip", "install", install_name], capture_output=True, text=True, - timeout=self.timeout + timeout=self.timeout, ) - + if result.returncode == 0: logger.info(f"✅ {install_name} 安装成功") self.installed_packages.add(package_name) @@ -68,29 +98,44 @@ def install_package(self, package_name: str) -> bool: else: logger.error(f"❌ {install_name} 安装失败: {result.stderr}") return False - + except Exception as e: logger.error(f"❌ 安装 {install_name} 时出错: {e}") return False - + def check_and_install_dependencies(self, code: str) -> bool: """检查并安装代码依赖""" if not self.auto_install: return True - + imports = self.extract_imports(code) logger.info(f"🔍 检测到导入: {imports}") - + # 标准库列表（无需安装） - stdlib = {'os', 'sys', 're', 'json', 'math', 'datetime', - 'time', 'random', 'itertools', 'functools', - 'collections', 'typing', 'pathlib', 'copy', - 'statistics', 'decimal', 'fractions'} - + stdlib = { + "os", + "sys", + "re", + "json", + "math", + "datetime", + "time", + "random", + "itertools", + "functools", + "collections", + "typing", + "pathlib", + "copy", + "statistics", + "decimal", + "fractions", + } + for package in imports: if package in stdlib: continue - + try: importlib.import_module(package) logger.info(f"✅ {package} 已存在") @@ -98,24 +143,31 @@ def check_and_install_dependencies(self, code: str) -> bool: logger.warning(f"⚠️ {package} 未安装，尝试安装...") if not self.install_package(package): return False - + return True - - def execute(self, code: str, data: Any, global_graph: Optional[GraphData] = None, is_numeric_analysis: bool = False) -> Any: + + def execute( + self, + code: str, + data: Any, + global_graph: Optional[GraphData] = None, + is_numeric_analysis: bool = False, + ) -> Any: """在独立命名空间中执行代码""" # 检查并安装依赖 if not self.check_and_install_dependencies(code): raise RuntimeError("依赖安装失败") - + # 如果 data 是序列化后的图字典，将其转换回 NetworkX 图对象 # 这样后处理代码就可以正常使用 data.nodes() 和 data.edges() 等方法 if isinstance(data, dict) and data.get("type") == "graph": try: import networkx as nx + graph_type = data.get("graph_type", "DiGraph") nodes = data.get("nodes", []) edges = data.get("edges", []) - + # 根据图类型创建对应的图对象 if graph_type == "DiGraph": restored_graph = nx.DiGraph() @@ -128,7 +180,7 @@ def execute(self, code: str, data: Any, global_graph: Optional[GraphData] = None else: # 默认使用 DiGraph restored_graph = nx.DiGraph() - + # 添加节点和边 restored_graph.add_nodes_from(nodes) for edge in edges: @@ -136,34 +188,40 @@ def execute(self, code: str, data: Any, global_graph: Optional[GraphData] = None dst = edge.get("dst") if src and dst: # 复制边属性（排除 src 和 dst） - edge_attrs = {k: v for k, v in edge.items() if k not in ("src", "dst")} + edge_attrs = { + k: v for k, v in edge.items() if k not in ("src", "dst") + } restored_graph.add_edge(src, dst, **edge_attrs) - - logger.info(f"🔄 已将序列化的图字典转换回 NetworkX {graph_type} 对象 (节点数: {len(nodes)}, 边数: {len(edges)})") + + logger.info( + f"🔄 已将序列化的图字典转换回 NetworkX {graph_type} 对象 (节点数: {len(nodes)}, 边数: {len(edges)})" + ) data = restored_graph except Exception as e: - logger.warning(f"⚠️ 无法将序列化的图字典转换回 NetworkX 图对象: {e}，将使用原始字典") + logger.warning( + f"⚠️ 无法将序列化的图字典转换回 NetworkX 图对象: {e}，将使用原始字典" + ) - # 创建独立命名空间 + # 创建独立命名空间（使用安全 builtins 白名单） namespace = { "data": data, - "global_graph": global_graph, - "__builtins__": __builtins__, # 保留内置函数 + "global_graph": global_graph, + "__builtins__": SAFE_BUILTINS, } - + try: # 执行代码 exec(code, namespace) - + # 检查 process 函数 if "process" not in namespace: raise ValueError("后处理代码必须定义 'process(data)' 函数") - + process_func = namespace["process"] - + if not callable(process_func): raise ValueError("'process' 必须是一个函数") - + # 根据 is_numeric_analysis 决定参数传递方式 if is_numeric_analysis: # 数值分析场景：data 是包含多个字段的字典，使用 **data 解包传递 @@ -181,12 +239,11 @@ def execute(self, code: str, data: Any, global_graph: Optional[GraphData] = None else: # 后处理场景：data 是算法结果（可能是字典），直接传递整个 data result = process_func(data) - + logger.info(f"✅ 后处理执行成功") - + return result - + except Exception as e: logger.error(f"❌ 后处理失败: {e}", exc_info=True) raise RuntimeError(f"后处理代码执行错误: {e}") - diff --git a/aag/computing_engine/computing_engine.py b/aag/computing_engine/computing_engine.py index 03bccd2..e90adcb 100644 --- a/aag/computing_engine/computing_engine.py +++ b/aag/computing_engine/computing_engine.py @@ -7,9 +7,12 @@ from aag.utils.path_utils import DEFAULT_CONFIG_SERVER_PATH from aag.computing_engine.mcp_client import GraphMCPClient -from aag.computing_engine.code_executor import DynamicCodeExecutor +from aag.computing_engine.code_executor import DynamicCodeExecutor, SAFE_BUILTINS from aag.expert_search_engine.database.datatype import GraphData -from aag.computing_engine.graph_query.nl_query_engine import NaturalLanguageQueryEngine, LLMInterface +from aag.computing_engine.graph_query.nl_query_engine import ( + NaturalLanguageQueryEngine, + LLMInterface, +) from aag.computing_engine.graph_query.graph_query import Neo4jGraphClient, Neo4jConfig @@ -27,8 +30,8 @@ class ComputingEngine: def __init__(self, config_path: str = DEFAULT_CONFIG_SERVER_PATH): self.config_path = Path(config_path) self.clients: Dict[str, GraphMCPClient] = {} - self.engine_supported_algorithms = {} # algorithm_type -> engine_name - self.algorithm_tool_mapping = {} # algorithm_name -> tool_name + self.engine_supported_algorithms = {} # algorithm_type -> engine_name + self.algorithm_tool_mapping = {} # algorithm_name -> tool_name self.parameter_modules = {} self._initialized = False self.code_executor = DynamicCodeExecutor(timeout=120, auto_install=True) @@ -36,16 +39,27 @@ def __init__(self, config_path: str = DEFAULT_CONFIG_SERVER_PATH): self.neo4j_config: Optional[Dict[str, Any]] = None self.reasoner = None # set in initialize_graph_query_engine + # 引擎名白名单：防止 YAML 配置注入导致的任意模块加载 + ALLOWED_ENGINES = {"networkx", "pyg"} + async def initialize(self): """Load config and connect to all MCP servers.""" config = self._load_config() self.engine_supported_algorithms = config.get("engine_supported_algorithms", {}) - self._parse_algorithm_tool_mapping(config.get("engine_supported_algorithms", {})) - + self._parse_algorithm_tool_mapping( + config.get("engine_supported_algorithms", {}) + ) + for engine_name, server in config.get("servers", {}).items(): + # 白名单校验：仅允许加载已知引擎模块 + if engine_name not in self.ALLOWED_ENGINES: + logger.error( + f"❌ Engine '{engine_name}' not in ALLOWED_ENGINES whitelist; skipped" + ) + continue client = GraphMCPClient( server_command=server.get("command", "python"), - server_args=server.get("args", []) + server_args=server.get("args", []), ) ok = await client.connect() if ok: @@ -55,17 +69,20 @@ async def initialize(self): logger.warning(f"⚠️ Engine '{engine_name}' failed to connect") try: - module_path = f"aag.computing_engine.{engine_name}_server.parameter_utils" + module_path = ( + f"aag.computing_engine.{engine_name}_server.parameter_utils" + ) module = importlib.import_module(module_path) self.parameter_modules[engine_name] = module logger.info(f"🧩 Loaded parameter module: {module_path}") except ImportError as e: - logger.warning(f"⚠️ No parameter_utils.py found for engine '{engine_name}' ({e})") + logger.warning( + f"⚠️ No parameter_utils.py found for engine '{engine_name}' ({e})" + ) self._initialized = True logger.info(f"✅ Loaded {len(self.clients)} computing engines") # logger.info(f"🧠 Annotation modules loaded: {list(self.parameter_modules.keys())}") - def _load_config(self) -> dict: """Load engine definitions from config file.""" if not self.config_path.exists(): @@ -80,10 +97,10 @@ def _parse_algorithm_tool_mapping(self, engine_algorithms: dict): bfs: - tool: bfs_edges """ - + for engine_name, algorithms in engine_algorithms.items(): engine_mapping = {} - + for algo_name, tool_configs in algorithms.items(): if isinstance(tool_configs, list) and len(tool_configs) > 0: tool_config = tool_configs[0] @@ -91,11 +108,15 @@ def _parse_algorithm_tool_mapping(self, engine_algorithms: dict): if tool_base_name and isinstance(tool_base_name, str): engine_mapping[algo_name] = f"run_{tool_base_name}" else: - logger.warning(f"⚠️ Invalid tool name for algorithm '{algo_name}' in engine '{engine_name}': {tool_base_name}") - + logger.warning( + f"⚠️ Invalid tool name for algorithm '{algo_name}' in engine '{engine_name}': {tool_base_name}" + ) + self.algorithm_tool_mapping[engine_name] = engine_mapping - - logger.info(f"✅ Parsed {sum(len(m) for m in self.algorithm_tool_mapping.values())} algorithm-tool mappings") + + logger.info( + f"✅ Parsed {sum(len(m) for m in self.algorithm_tool_mapping.values())} algorithm-tool mappings" + ) logger.debug(f"Algorithm-tool mapping: {self.algorithm_tool_mapping}") def _resolve_engine(self, algo_name: str) -> str: @@ -105,7 +126,9 @@ def _resolve_engine(self, algo_name: str) -> str: if isinstance(algorithms, dict): if algo_name in algorithms: return engine - logger.warning(f"⚠️ Algorithm '{algo_name}' not found in config; fallback to 'networkx'") + logger.error( + f"❌ Algorithm '{algo_name}' not found in config; fallback to 'networkx'" + ) return "networkx" def _resolve_tool_name(self, algo_name: str, engine_name: str) -> str: @@ -125,34 +148,55 @@ def _resolve_tool_name(self, algo_name: str, engine_name: str) -> str: engine_mapping = self.algorithm_tool_mapping.get(engine_name, {}) if algo_name in engine_mapping: return engine_mapping[algo_name] - - logger.error(f"❌ Cannot resolve tool name for algorithm '{algo_name}' in engine '{engine_name}'") - raise ValueError(f"Algorithm '{algo_name}' has no tool mapping in engine '{engine_name}'. " - f"Please check config_servers.yaml") - async def run_algorithm(self, algo_name: str, parameters: Dict[str, Any], post_processing_code: Optional[str] = None, global_graph: Optional[GraphData] = None) -> Dict[str, Any]: + logger.error( + f"❌ Cannot resolve tool name for algorithm '{algo_name}' in engine '{engine_name}'" + ) + raise ValueError( + f"Algorithm '{algo_name}' has no tool mapping in engine '{engine_name}'. " + f"Please check config_servers.yaml" + ) + + async def run_algorithm( + self, + algo_name: str, + parameters: Dict[str, Any], + post_processing_code: Optional[str] = None, + global_graph: Optional[GraphData] = None, + ) -> Dict[str, Any]: """Run the specified algorithm.""" try: engine_name = self._resolve_engine(algo_name) if engine_name not in self.clients: - return {"success": False, "error": f"Engine '{engine_name}' not connected"} + return { + "success": False, + "error": f"Engine '{engine_name}' not connected", + } client = self.clients[engine_name] tool_name = self._resolve_tool_name(algo_name, engine_name) - + prepared_params = parameters or {} if self._should_normalize(engine_name, tool_name): - prepared_params = self._normalize_parameters(engine_name, tool_name, prepared_params) - - logger.info(f"🚀 Running '{algo_name}' (tool: {tool_name}) on engine [{engine_name}]") - result = await client.call_tool(tool_name, prepared_params, post_processing_code, global_graph) + prepared_params = self._normalize_parameters( + engine_name, tool_name, prepared_params + ) + + logger.info( + f"🚀 Running '{algo_name}' (tool: {tool_name}) on engine [{engine_name}]" + ) + result = await client.call_tool( + tool_name, prepared_params, post_processing_code, global_graph + ) return result except Exception as e: logger.error(f"❌ Algorithm '{algo_name}' failed: {e}") return {"success": False, "error": str(e)} - async def get_algorithm_description(self, algo_name: str) -> tuple[str, Optional[dict]]: + async def get_algorithm_description( + self, algo_name: str + ) -> tuple[str, Optional[dict]]: """Get tool description (input/output schema) for the given algorithm.""" engine = self._resolve_engine(algo_name) client = self.clients.get(engine) @@ -174,7 +218,7 @@ async def get_algorithm_description(self, algo_name: str) -> tuple[str, Optional f"```json\n{json.dumps(annotated_input, indent=2, ensure_ascii=False)}\n```\n\n" f"**Output Structure** *(data received by your `process(data)` in the `result` field)*:\n" f"```json\n{json.dumps(output_schema, indent=2, ensure_ascii=False)}\n```\n" - f"{'-'*50}\n" + f"{'-' * 50}\n" ) tool_metadata = { @@ -202,9 +246,10 @@ def _annotate_schema(self, input_schema: Dict, engine_name: str) -> Dict: try: return annotate_func(input_schema) except Exception as e: - logger.error(f"❌ Error running annotate_schema() for '{engine_name}': {e}") + logger.error( + f"❌ Error running annotate_schema() for '{engine_name}': {e}" + ) return input_schema - def _should_normalize(self, engine_name: str, tool_name: str) -> bool: module = self.parameter_modules.get(engine_name) @@ -215,7 +260,9 @@ def _should_normalize(self, engine_name: str, tool_name: str) -> bool: return guard(tool_name) return hasattr(module, "normalize_parameters") - def _normalize_parameters(self, engine_name: str, tool_name: str, parameters: Dict[str, Any]) -> Dict[str, Any]: + def _normalize_parameters( + self, engine_name: str, tool_name: str, parameters: Dict[str, Any] + ) -> Dict[str, Any]: module = self.parameter_modules.get(engine_name) if not module: return parameters @@ -226,12 +273,21 @@ def _normalize_parameters(self, engine_name: str, tool_name: str, parameters: Di client = self.clients.get(engine_name) tool_info = client.available_tools.get(tool_name) if client else None try: - return normalize_fn(tool_name, dict(parameters or {}), tool_info, logger=logger) + return normalize_fn( + tool_name, dict(parameters or {}), tool_info, logger=logger + ) except Exception as exc: logger.warning("⚠️ Parameter normalization failed (%s): %s", tool_name, exc) return parameters - def execute_code(self, code: str, data: Any, global_graph: Optional[GraphData] = None, fallback_to_direct_exec: bool = True, is_numeric_analysis: bool = False) -> Any: + def execute_code( + self, + code: str, + data: Any, + global_graph: Optional[GraphData] = None, + fallback_to_direct_exec: bool = True, + is_numeric_analysis: bool = False, + ) -> Any: """ Execute dynamically generated code (numeric analysis, post-processing, etc.). @@ -248,34 +304,54 @@ def execute_code(self, code: str, data: Any, global_graph: Optional[GraphData] = RuntimeError: If execution fails. """ try: - value = self.code_executor.execute(code, data, global_graph=global_graph, is_numeric_analysis=is_numeric_analysis) + value = self.code_executor.execute( + code, + data, + global_graph=global_graph, + is_numeric_analysis=is_numeric_analysis, + ) return { "algorithm": "numeric_analysis_code", "success": True, "result": value, "error": None, "summary": "Numeric analysis code executed successfully.", - } + } except (ValueError, AttributeError) as e: if not fallback_to_direct_exec: - return {"success": False, "result": None, "error": str(e), "summary": "Numeric analysis code execution failed."} + return { + "success": False, + "result": None, + "error": str(e), + "summary": "Numeric analysis code execution failed.", + } try: if isinstance(data, dict): - namespace = {**data, "__builtins__": __builtins__} + namespace = {**data, "__builtins__": SAFE_BUILTINS} else: - namespace = {"data": data, "__builtins__": __builtins__} + namespace = {"data": data, "__builtins__": SAFE_BUILTINS} exec(code, namespace) return { - "algorithm": "numeric_analysis_code", - "success": True, - "result": namespace.get("result", data), - "error": None, - "summary": "Numeric analysis code executed successfully via direct exec fallback.", + "algorithm": "numeric_analysis_code", + "success": True, + "result": namespace.get("result", data), + "error": None, + "summary": "Numeric analysis code executed successfully via direct exec fallback.", } except Exception as fallback_error: - return {"success": False, "result": None, "error": str(fallback_error), "summary": "Numeric analysis code execution failed."} + return { + "success": False, + "result": None, + "error": str(fallback_error), + "summary": "Numeric analysis code execution failed.", + } except Exception as e: - return {"success": False, "result": None, "error": str(e), "summary": "Numeric analysis code execution failed."} + return { + "success": False, + "result": None, + "error": str(e), + "summary": "Numeric analysis code execution failed.", + } def initialize_graph_query_engine(self, neo4j_config: Dict[str, Any], reasoner): """ @@ -299,7 +375,7 @@ def initialize_graph_query_engine(self, neo4j_config: Dict[str, Any], reasoner): config = Neo4jConfig( uri=neo4j_config.get("uri", "bolt://localhost:7687"), user=neo4j_config.get("user", "neo4j"), - password=neo4j_config.get("password", "") + password=neo4j_config.get("password", ""), ) logger.info(f"📝 Creating Neo4jGraphClient, uri={config.uri}") db_client = Neo4jGraphClient(config) @@ -310,7 +386,10 @@ def initialize_graph_query_engine(self, neo4j_config: Dict[str, Any], reasoner): self.nl_query_engine.initialize() logger.info("✓ NaturalLanguageQueryEngine initialized in ComputingEngine") except Exception as e: - logger.error(f"✗ NaturalLanguageQueryEngine initialization failed: {e}", exc_info=True) + logger.error( + f"✗ NaturalLanguageQueryEngine initialization failed: {e}", + exc_info=True, + ) self.nl_query_engine = None def execute_graph_query(self, query: str) -> Dict[str, Any]: @@ -326,7 +405,7 @@ def execute_graph_query(self, query: str) -> Dict[str, Any]: if not self.nl_query_engine: return { "success": False, - "error": "Graph query engine not initialized; check Neo4j config." + "error": "Graph query engine not initialized; check Neo4j config.", } try: result = self.nl_query_engine.ask(query) @@ -335,7 +414,7 @@ def execute_graph_query(self, query: str) -> Dict[str, Any]: logger.error(f"Graph query execution failed: {e}", exc_info=True) return { "success": False, - "error": f"Graph query execution failed: {str(e)}" + "error": f"Graph query execution failed: {str(e)}", } async def shutdown(self): diff --git a/aag/computing_engine/graph_query/graph_query.py b/aag/computing_engine/graph_query/graph_query.py index dcbdd1c..fe20f8d 100644 --- a/aag/computing_engine/graph_query/graph_query.py +++ b/aag/computing_engine/graph_query/graph_query.py @@ -5,11 +5,13 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple from neo4j import GraphDatabase from neo4j.exceptions import ClientError, DatabaseError, TransientError +import os import re import time JsonDict = Dict[str, Any] + @dataclass class Neo4jConfig: uri: str @@ -33,25 +35,48 @@ class Neo4jGraphClient: # Neo4j 保留字（部分） RESERVED_KEYWORDS = { - 'MATCH', 'RETURN', 'WHERE', 'CREATE', 'DELETE', 'SET', - 'MERGE', 'WITH', 'UNWIND', 'CASE', 'WHEN', 'THEN', - 'ELSE', 'END', 'ORDER', 'BY', 'SKIP', 'LIMIT', 'AS', - 'AND', 'OR', 'NOT', 'IN', 'IS', 'NULL', 'TRUE', 'FALSE' + "MATCH", + "RETURN", + "WHERE", + "CREATE", + "DELETE", + "SET", + "MERGE", + "WITH", + "UNWIND", + "CASE", + "WHEN", + "THEN", + "ELSE", + "END", + "ORDER", + "BY", + "SKIP", + "LIMIT", + "AS", + "AND", + "OR", + "NOT", + "IN", + "IS", + "NULL", + "TRUE", + "FALSE", } def __init__(self, config: Neo4jConfig): """ 初始化 Neo4j 客户端 - + Args: config: Neo4j 连接配置 """ self._driver = GraphDatabase.driver( - config.uri, + config.uri, auth=(config.user, config.password), max_connection_lifetime=3600, # 连接最大存活时间 1小时 - max_connection_pool_size=50, # 连接池大小 - connection_acquisition_timeout=60.0 # 获取连接超时 + max_connection_pool_size=50, # 连接池大小 + connection_acquisition_timeout=60.0, # 获取连接超时 ) self._db = config.database @@ -76,30 +101,30 @@ def run( *, read: bool = True, max_retries: int = 3, - show_query: bool = True + show_query: bool = True, ) -> List[JsonDict]: """ 执行 Cypher 查询（带自动重试） - + Args: cypher: Cypher 查询语句 params: 查询参数 read: 是否为读操作（True=读，False=写） max_retries: 最大重试次数 show_query: 是否打印填充参数后的查询 - + Returns: 查询结果列表 - + Raises: RuntimeError: 查询失败 """ params = params or {} - + # 可选：打印填充参数后的查询（用于调试） if show_query: self._print_filled_query(cypher, params) - + # 重试逻辑 for attempt in range(max_retries): try: @@ -112,7 +137,7 @@ def run( return session.write_transaction( lambda tx: [r.data() for r in tx.run(cypher, params)] ) - + except TransientError as e: # 临时性错误：重试 if attempt == max_retries - 1: @@ -120,47 +145,52 @@ def run( f"Neo4j TransientError after {max_retries} retries: {e}\n" f"Cypher: {cypher}\nParams: {params}" ) from e - - wait_time = 2 ** attempt # 指数退避 - print(f"⚠️ TransientError, retrying in {wait_time}s ({attempt + 1}/{max_retries})...") + + wait_time = 2**attempt # 指数退避 + print( + f"⚠️ TransientError, retrying in {wait_time}s ({attempt + 1}/{max_retries})..." + ) time.sleep(wait_time) - + except (ClientError, DatabaseError) as e: # 客户端错误或数据库错误：不重试 raise RuntimeError( - f"Neo4j Error: {e}\n" - f"Cypher: {cypher}\n" - f"Params: {params}" + f"Neo4j Error: {e}\nCypher: {cypher}\nParams: {params}" ) from e def _print_filled_query(self, cypher: str, params: Dict) -> None: """ 打印填充参数后的查询（用于调试） - + 注意：这个方法只用于显示，实际执行时 Neo4j 驱动会正确处理参数 """ - print("\n" + "="*80) + print("\n" + "=" * 80) print("📝 执行的 Cypher 查询语句（参数已填充）:") - print("-"*80) - + print("-" * 80) + filled_cypher = cypher if params: import json + # 按参数名长度倒序排序，避免子串替换问题 - sorted_params = sorted(params.items(), key=lambda x: len(x[0]), reverse=True) - + sorted_params = sorted( + params.items(), key=lambda x: len(x[0]), reverse=True + ) + for key, value in sorted_params: # ⚠️ CRITICAL: 跳过元组类型的参数（这些是 DSL 内部格式，不应该出现在最终 Cypher 中） # 元组格式如 (">", 500000) 应该已经在构建 WHERE 子句时被转换为 Cypher 操作符 if isinstance(value, tuple): # 这种情况不应该发生，如果发生了说明有 bug - print(f"⚠️ WARNING: 参数 ${key} 是元组格式 {value}，这不应该出现在最终查询中！") + print( + f"⚠️ WARNING: 参数 ${key} 是元组格式 {value}，这不应该出现在最终查询中！" + ) continue - + # 根据值类型格式化 if isinstance(value, str): # 转义单引号 - formatted_value = f"'{value.replace(chr(39), chr(39)+chr(39))}'" + formatted_value = f"'{value.replace(chr(39), chr(39) + chr(39))}'" elif isinstance(value, (int, float)): formatted_value = str(value) elif isinstance(value, bool): @@ -176,22 +206,22 @@ def _print_filled_query(self, cypher: str, params: Dict) -> None: formatted_value = json.dumps(value, ensure_ascii=False) else: formatted_value = str(value) - + # 使用正则确保完整匹配（避免 $id 替换 $id2 的问题） filled_cypher = re.sub( - r'\$' + re.escape(key) + r'\b', # \b 确保单词边界 + r"\$" + re.escape(key) + r"\b", # \b 确保单词边界 formatted_value, - filled_cypher + filled_cypher, ) - + print(filled_cypher) - print("="*80 + "\n") + print("=" * 80 + "\n") # ===== Schema 获取 ===== def get_schema(self) -> Dict: """ 获取图数据库 Schema 信息（增强版） - + Returns: { "node_labels": { @@ -209,63 +239,72 @@ def get_schema(self) -> Dict: "patterns": [pattern_strings] } """ - schema = { - "node_labels": {}, - "relationship_types": {}, - "patterns": [] - } - + schema = {"node_labels": {}, "relationship_types": {}, "patterns": []} + # 1. 获取所有节点标签及其属性（包含示例值） - labels_result = self.run("CALL db.labels() YIELD label RETURN label", show_query=False) + labels_result = self.run( + "CALL db.labels() YIELD label RETURN label", show_query=False + ) valid_labels = [item["label"] for item in labels_result] - + for label in valid_labels: if not label or not self._is_valid_identifier(label): continue - + # 获取该标签的属性和示例值 - props_result = self.run(f""" + props_result = self.run( + f""" MATCH (n:`{label}`) WITH n LIMIT 1 UNWIND keys(n) AS key RETURN key, n[key] AS sample_value ORDER BY key - """, show_query=False) - + """, + show_query=False, + ) + if props_result: schema["node_labels"][label] = { "properties": [p["key"] for p in props_result], - "sample_values": {p["key"]: p["sample_value"] for p in props_result} + "sample_values": { + p["key"]: p["sample_value"] for p in props_result + }, } - + # 2. 获取所有关系类型及其属性（包含示例值） rels_result = self.run( "CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType", - show_query=False + show_query=False, ) valid_rels = [item["relationshipType"] for item in rels_result] - + for rel_type in valid_rels: if not rel_type or not self._is_valid_identifier(rel_type): continue - + # 获取该关系的属性和示例值 - props_result = self.run(f""" + props_result = self.run( + f""" MATCH ()-[r:`{rel_type}`]->() WITH r LIMIT 1 UNWIND keys(r) AS key RETURN key, r[key] AS sample_value ORDER BY key - """, show_query=False) - + """, + show_query=False, + ) + if props_result: schema["relationship_types"][rel_type] = { "properties": [p["key"] for p in props_result], - "sample_values": {p["key"]: p["sample_value"] for p in props_result} + "sample_values": { + p["key"]: p["sample_value"] for p in props_result + }, } - + # 3. 获取关系模式 - patterns = self.run(""" + patterns = self.run( + """ MATCH (a)-[r]->(b) WITH labels(a)[0] AS start_label, type(r) AS rel_type, @@ -275,26 +314,28 @@ def get_schema(self) -> Dict: AND end_label IS NOT NULL RETURN DISTINCT start_label, rel_type, end_label LIMIT 100 - """, show_query=False) - + """, + show_query=False, + ) + schema["patterns"] = [ f"({p['start_label']})-[:{p['rel_type']}]->({p['end_label']})" for p in patterns ] - + return schema # ===== 验证方法 ===== @staticmethod def _is_valid_identifier(name: str) -> bool: """快速检查标识符是否有效（字母数字下划线）""" - return bool(re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', name)) + return bool(re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", name)) @classmethod def _sanitize_label(cls, label: Optional[str]) -> str: """ 严格验证节点标签 - + 规则： - 必须以字母开头 - 只能包含字母、数字、下划线 @@ -303,26 +344,28 @@ def _sanitize_label(cls, label: Optional[str]) -> str: """ if not label: return "" - + # 长度检查 if len(label) > 255: raise ValueError(f"Label too long (max 255): {label}") - + # 格式检查 - if not re.match(r'^[A-Za-z][A-Za-z0-9_]*$', label): - raise ValueError(f"Invalid label format (must start with letter, contain only alphanumeric and underscore): {label}") - + if not re.match(r"^[A-Za-z][A-Za-z0-9_]*$", label): + raise ValueError( + f"Invalid label format (must start with letter, contain only alphanumeric and underscore): {label}" + ) + # 保留字检查 if label.upper() in cls.RESERVED_KEYWORDS: raise ValueError(f"Reserved keyword cannot be used as label: {label}") - + return label @classmethod def _sanitize_rel_type(cls, rel_type: Optional[str]) -> str: """ 严格验证关系类型 - + 规则： - 通常使用大写字母和下划线（如 FOLLOWS、HAS_FRIEND） - 也允许小写和驼峰（兼容性） @@ -330,24 +373,26 @@ def _sanitize_rel_type(cls, rel_type: Optional[str]) -> str: """ if not rel_type: return "" - + if len(rel_type) > 255: raise ValueError(f"Relationship type too long (max 255): {rel_type}") - + # 格式检查（允许大小写字母、数字、下划线） - if not re.match(r'^[A-Za-z][A-Za-z0-9_]*$', rel_type): + if not re.match(r"^[A-Za-z][A-Za-z0-9_]*$", rel_type): raise ValueError(f"Invalid relationship type format: {rel_type}") - + if rel_type.upper() in cls.RESERVED_KEYWORDS: - raise ValueError(f"Reserved keyword cannot be used as relationship type: {rel_type}") - + raise ValueError( + f"Reserved keyword cannot be used as relationship type: {rel_type}" + ) + return rel_type @staticmethod def _sanitize_property_key(key: str) -> str: """ 验证属性键 - + 规则： - 必须以字母开头 - 只能包含字母、数字、下划线 @@ -355,39 +400,36 @@ def _sanitize_property_key(key: str) -> str: """ if not key: raise ValueError("Property key cannot be empty") - + if len(key) > 255: raise ValueError(f"Property key too long (max 255): {key}") - - if not re.match(r'^[a-zA-Z][a-zA-Z0-9_]*$', key): + + if not re.match(r"^[a-zA-Z][a-zA-Z0-9_]*$", key): raise ValueError(f"Invalid property key format: {key}") - + return key # ========================================================= # 1. 根据 ID / 唯一键查节点 # ========================================================= def get_node_by_internal_id( - self, - internal_id: int, - *, - return_props: bool = True + self, internal_id: int, *, return_props: bool = True ) -> Optional[JsonDict]: """ 使用 Neo4j 内部 id(n) 查询 - + 注意：内部 id 在导入/删除后可能变化，不建议作为业务主键 - + Args: internal_id: Neo4j 内部节点 ID return_props: 是否只返回节点属性 - + Returns: 节点数据或 None """ cypher = "MATCH (n) WHERE id(n) = $id RETURN n AS node" res = self.run(cypher, {"id": internal_id}) - + if not res: return None return res[0]["node"] if return_props else res[0] @@ -398,31 +440,31 @@ def get_node_by_unique_key( key: str, value: Any, *, - return_fields: Optional[List[str]] = None + return_fields: Optional[List[str]] = None, ) -> Optional[JsonDict]: """ 根据 label + 唯一键属性查询节点 - + 示例： # 返回整个节点 get_node_by_unique_key("User", "userId", "u123") - + # 只返回指定字段 get_node_by_unique_key("Account", "node_key", "Collins Steven", return_fields=["acct_id", "acct_stat", "acct_open_date"]) - + Args: label: 节点标签 key: 属性键（如 userId） value: 属性值（如 u123） return_fields: 要返回的字段列表（None=返回整个节点） - + Returns: 节点数据或 None """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + # 构建RETURN子句 if return_fields: # 返回指定字段 @@ -434,13 +476,13 @@ def get_node_by_unique_key( else: # 返回整个节点 return_clause = "RETURN n AS node" - + cypher = f"MATCH (n:`{label}` {{`{key}`: $value}}) {return_clause} LIMIT 1" res = self.run(cypher, {"value": value}) - + if not res: return None - + # 如果返回指定字段，直接返回结果字典；否则返回node if return_fields: return res[0] @@ -456,21 +498,21 @@ def filter_nodes_by_properties( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 根据属性条件筛选多个节点（支持多条件AND组合） - + 功能： - 支持多个属性条件的AND组合筛选 - 支持指定返回字段（而非返回整个节点） - 支持排序和限制结果数量 - + 应用场景： 1. 筛选查询：查找满足特定条件的所有节点 2. 条件过滤：根据多个属性值组合筛选 3. 字段投影：只返回需要的字段，减少数据传输 - + 示例： # 查询居住在US、州为VT的客户姓名和城市 results = client.filter_nodes_by_properties( @@ -478,7 +520,7 @@ def filter_nodes_by_properties( {"country": "US", "state": "VT"}, return_fields=["last_name", "first_name", "city"] ) - + # 查询账户币种是USD且账户状态为A的客户 results = client.filter_nodes_by_properties( "Account", @@ -487,21 +529,21 @@ def filter_nodes_by_properties( order_by="last_name", limit=10 ) - + # ⚠️ 范围条件查询：initial_deposit 大于 500000 的账户 results = client.filter_nodes_by_properties( "Account", {"initial_deposit": (">", 500000)}, return_fields=["acct_id", "initial_deposit"] ) - + # 布尔值查询：prior_sar_count 为 true 的账户 results = client.filter_nodes_by_properties( "Account", {"prior_sar_count": true}, # 注意：Python 中是 True，但会自动转换为 Neo4j 的 true return_fields=["acct_id"] ) - + Args: label: 节点标签 conditions: 属性条件字典，支持两种格式： @@ -513,39 +555,49 @@ def filter_nodes_by_properties( order_by: 排序字段（如 "last_name"） order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回数量（None=不限制） - + Returns: [ {"last_name": "Smith", "first_name": "John", "city": "Burlington"}, {"last_name": "Doe", "first_name": "Jane", "city": "Montpelier"}, ... ] - + 注意： - 这是多节点筛选查询，不是单节点精确查找 - 所有条件使用AND逻辑组合 - 如果需要OR逻辑，请使用filter_query方法 """ label = self._sanitize_label(label) - + if not conditions: raise ValueError("conditions cannot be empty") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + # 构建WHERE子句 where_parts = [] params = {} for i, (key, value) in enumerate(conditions.items()): key = self._sanitize_property_key(key) param_name = f"cond_{i}" - + # ⚠️ 修复2: 支持范围条件 (operator, value) 元组格式或 [operator, value] 数组格式 - if (isinstance(value, (tuple, list)) and len(value) == 2): + if isinstance(value, (tuple, list)) and len(value) == 2: operator, actual_value = value # 验证操作符 - valid_operators = ["=", ">", "<", ">=", "<=", "!=", "IN", "CONTAINS", "STARTS WITH"] + valid_operators = [ + "=", + ">", + "<", + ">=", + "<=", + "!=", + "IN", + "CONTAINS", + "STARTS WITH", + ] if operator.upper() in ["IN", "CONTAINS"]: where_parts.append(f"a.`{key}` {operator.upper()} ${param_name}") elif operator.upper() == "STARTS WITH": @@ -559,9 +611,9 @@ def filter_nodes_by_properties( # 简单等值条件 where_parts.append(f"a.`{key}` = ${param_name}") params[param_name] = value - + where_clause = "WHERE " + " AND ".join(where_parts) - + # 构建RETURN子句 if return_fields: # 返回指定字段 @@ -573,19 +625,19 @@ def filter_nodes_by_properties( else: # 返回整个节点 return_clause = "RETURN a AS node" - + # 可选的ORDER BY子句 order_clause = "" if order_by: order_by = self._sanitize_property_key(order_by) order_clause = f"ORDER BY a.`{order_by}` {order_direction}" - + # 可选的LIMIT子句 limit_clause = "" if limit is not None: limit_clause = "LIMIT $limit" params["limit"] = limit - + # 构建完整查询 cypher = f""" MATCH (a:`{label}`) @@ -594,7 +646,7 @@ def filter_nodes_by_properties( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # add gjq @@ -610,22 +662,22 @@ def filter_relationships( aggregate_field: Optional[str] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 基于关系属性条件筛选关系（支持聚合统计） - + 功能： - 根据关系属性条件筛选关系 - 支持指定起点和终点节点类型 - 支持返回指定字段（关系属性和节点属性） - 支持聚合统计（COUNT、SUM、AVG等） - + 应用场景： 1. 关系属性过滤：查找满足特定条件的关系 2. 关系统计：统计满足条件的关系数量、金额总和等 3. 关系分析：分析关系的分布、趋势等 - + 示例： # 查找交易金额大于400的交易 results = client.filter_relationships( @@ -633,14 +685,14 @@ def filter_relationships( rel_conditions={"base_amt": (">", 400)}, return_fields=["tran_id", "from.node_key", "to.node_key"] ) - + # 统计is_sar为False的交易数量 results = client.filter_relationships( "TRANSFER", rel_conditions={"is_sar": ("=", False)}, aggregate="COUNT" ) - + # 查找特定日期范围的交易 results = client.filter_relationships( "TRANSFER", @@ -652,7 +704,7 @@ def filter_relationships( order_direction="DESC", limit=10 ) - + Args: rel_type: 关系类型 start_label: 起点节点标签（可选） @@ -668,35 +720,39 @@ def filter_relationships( order_by: 排序字段 order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回数量（None=不限制） - + Returns: 如果是聚合查询，返回聚合结果： [{"aggregate_type": "count", "value": 100}] - + 如果是普通查询，返回关系和节点信息： [ {"tran_id": "T001", "from_account": "A001", "to_account": "A002"}, ... ] - + 注意： - 这是关系查询，不是节点查询 - rel_conditions中的条件使用AND逻辑组合 - 如果需要OR逻辑，请使用filter_query方法 """ rel_type = self._sanitize_rel_type(rel_type) - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + # 构建节点模式 - start_pattern = f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" - end_pattern = f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" - + start_pattern = ( + f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" + ) + end_pattern = ( + f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" + ) + # 构建WHERE子句 where_parts = [] params = {} - + if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): # 处理特殊的日期范围键名（如 tran_timestamp_start, tran_timestamp_end） @@ -706,33 +762,37 @@ def filter_relationships( actual_key = self._sanitize_property_key(actual_key) else: actual_key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name = f"rel_cond_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"t.`{actual_key}` {operator.upper()} ${param_name}") + where_parts.append( + f"t.`{actual_key}` {operator.upper()} ${param_name}" + ) elif operator.upper() == "STARTS WITH": - where_parts.append(f"t.`{actual_key}` STARTS WITH ${param_name}") + where_parts.append( + f"t.`{actual_key}` STARTS WITH ${param_name}" + ) else: where_parts.append(f"t.`{actual_key}` {operator} ${param_name}") - + params[param_name] = value else: # 简单等值条件 param_name = f"rel_cond_{i}" where_parts.append(f"t.`{actual_key}` = ${param_name}") params[param_name] = condition - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # 聚合查询 if aggregate: agg_type = aggregate.upper() if agg_type not in ["COUNT", "SUM", "AVG", "MAX", "MIN"]: raise ValueError(f"Invalid aggregate type: {aggregate}") - + if agg_type == "COUNT": agg_expr = "COUNT(t)" else: @@ -740,18 +800,20 @@ def filter_relationships( raise ValueError(f"aggregate_field is required for {agg_type}") agg_field = self._sanitize_property_key(aggregate_field) agg_expr = f"{agg_type}(t.`{agg_field}`)" - + cypher = f""" MATCH {start_pattern}-[t:`{rel_type}`]->{end_pattern} {where_clause} RETURN {agg_expr} AS value """ - + result = self.run(cypher, params) if result: - return [{"aggregate_type": agg_type.lower(), "value": result[0]["value"]}] + return [ + {"aggregate_type": agg_type.lower(), "value": result[0]["value"]} + ] return [{"aggregate_type": agg_type.lower(), "value": 0}] - + # 普通查询 if return_fields: # 构建RETURN子句 @@ -773,24 +835,24 @@ def filter_relationships( # 关系属性（不带前缀，直接是属性名） prop = self._sanitize_property_key(field) return_parts.append(f"t.`{prop}` AS {prop}") - + return_clause = "RETURN " + ", ".join(return_parts) else: # 返回整个关系和节点 return_clause = "RETURN from, t AS relationship, to" - + # 可选的ORDER BY子句 order_clause = "" if order_by: order_by = self._sanitize_property_key(order_by) order_clause = f"ORDER BY t.`{order_by}` {order_direction}" - + # 可选的LIMIT子句 limit_clause = "" if limit is not None: limit_clause = "LIMIT $limit" params["limit"] = limit - + # 构建完整查询 cypher = f""" MATCH {start_pattern}-[t:`{rel_type}`]->{end_pattern} @@ -799,7 +861,7 @@ def filter_relationships( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # add gjq @@ -817,23 +879,23 @@ def aggregation_query( where: Optional[str] = None, order_by: Optional[str] = None, order_direction: str = "DESC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 聚合统计查询（支持按节点或属性分组） - + 功能： - 支持按节点分组（如：每个账户） - 支持按属性分组（如：每个branch_id） - 支持多种聚合函数（COUNT、SUM、AVG、MAX、MIN） - 支持关系聚合和节点聚合 - 支持排序和TOP-N查询 - + 应用场景： 1. 统计每个账户的交易次数/金额 2. 计算每个分支的账户数量 3. 排名查询（TOP-N） - + 示例： # 统计每个账户作为转出账户的交易次数，返回前5个 results = client.aggregation_query( @@ -847,7 +909,7 @@ def aggregation_query( order_direction="DESC", limit=5 ) - + # 计算每个账户的转出交易总金额 results = client.aggregation_query( "SUM", @@ -860,14 +922,14 @@ def aggregation_query( order_by="total", order_direction="DESC" ) - + # 统计每个branch_id下的账户数量 results = client.aggregation_query( "COUNT", group_by_property="branch_id", node_label="Account" ) - + Args: aggregate_type: 聚合类型（"COUNT", "SUM", "AVG", "MAX", "MIN"） group_by_node: 按节点分组（"start"=起点, "end"=终点, None=不按节点分组） @@ -881,7 +943,7 @@ def aggregation_query( order_by: 排序字段（"count", "total", "avg"等） order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回数量（None=不限制） - + Returns: [ {"node_key": "A001", "count": 100}, @@ -892,12 +954,12 @@ def aggregation_query( agg_type = aggregate_type.upper() if agg_type not in ["COUNT", "SUM", "AVG", "MAX", "MIN"]: raise ValueError(f"Invalid aggregate type: {aggregate_type}") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + params = {} - + # 构建聚合表达式 if agg_type == "COUNT": if rel_type: @@ -915,18 +977,18 @@ def aggregation_query( else: agg_expr = f"{agg_type}(n.`{agg_field}`) AS total" agg_alias = "total" - + # 场景1：按节点属性分组（不涉及关系） if group_by_property and not rel_type: label = self._sanitize_label(node_label) if node_label else "" label_pattern = f":`{label}`" if label else "" prop = self._sanitize_property_key(group_by_property) - + where_clause = f"WHERE {where}" if where else "" limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH (n{label_pattern}) {where_clause} @@ -934,13 +996,13 @@ def aggregation_query( ORDER BY {order_by or agg_alias} {order_direction} {limit_clause} """ - + # 场景2：按节点分组 + 关系聚合 elif group_by_node and rel_type: label = self._sanitize_label(node_label) if node_label else "" label_pattern = f":`{label}`" if label else "" rt = self._sanitize_rel_type(rel_type) - + # 构建关系模式 if direction == "out": if group_by_node == "start": @@ -954,7 +1016,7 @@ def aggregation_query( pattern = f"()<-[r:`{rt}`]-(n{label_pattern})" else: pattern = f"(n{label_pattern})-[r:`{rt}`]-()" - + # 构建返回字段 return_parts = [] if return_fields: @@ -963,14 +1025,14 @@ def aggregation_query( return_parts.append(f"n.`{field}` AS {field}") else: return_parts.append("n.node_key AS node_key") - + return_clause = ", ".join(return_parts) + f", {agg_expr}" - + where_clause = f"WHERE {where}" if where else "" limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH {pattern} {where_clause} @@ -978,10 +1040,12 @@ def aggregation_query( ORDER BY {order_by or agg_alias} {order_direction} {limit_clause} """ - + else: - raise ValueError("Must specify either group_by_property or (group_by_node + rel_type)") - + raise ValueError( + "Must specify either group_by_property or (group_by_node + rel_type)" + ) + return self.run(cypher, params) # ========================================================= @@ -1004,11 +1068,11 @@ def neighbors_n_hop( limit: Optional[int] = None, return_distinct: bool = False, exclude_start: bool = False, - return_path_length: bool = False + return_path_length: bool = False, ) -> List[JsonDict]: """ 查询节点的 N 跳邻居（支持通用修饰符和字段投影） - + 功能： - 查询节点的N跳邻居 - 支持指定返回字段（关系属性和邻居节点属性） @@ -1016,12 +1080,12 @@ def neighbors_n_hop( - 支持排序和限制结果数量 - 支持去重和排除起始节点（多跳查询推荐） - 支持返回路径长度（按距离排序时需要） - + 特别说明： - 当hops=1且需要返回每条边的详细信息时，建议使用return_fields指定字段 - 对于"转出交易明细"等场景，应使用direction="out"确保方向正确 - 多跳查询时建议设置 return_distinct=True 和 exclude_start=True - + 示例： # 查询转出交易明细（每笔交易的转入账户和金额） results = client.neighbors_n_hop( @@ -1031,7 +1095,7 @@ def neighbors_n_hop( direction="out", return_fields=["nbr.acct_id", "rel.base_amt", "rel.tran_id"] ) - + # 查询二跳邻居（去重并排除起始节点） results = client.neighbors_n_hop( "Account", "node_key", "Lee Alex", @@ -1040,7 +1104,7 @@ def neighbors_n_hop( return_distinct=True, exclude_start=True ) - + # 按距离排序的多跳查询 results = client.neighbors_n_hop( "Account", "node_key", "Collins Steven", @@ -1053,7 +1117,7 @@ def neighbors_n_hop( order_by="path_length", order_direction="ASC" ) - + Args: label: 起点节点标签 key: 起点节点属性键 @@ -1073,32 +1137,32 @@ def neighbors_n_hop( return_distinct: 是否去重（多跳查询时建议 True） exclude_start: 是否排除起始节点（多跳查询时建议 True） return_path_length: 是否返回路径长度（需要按距离排序时设为 True） - + Returns: 如果指定return_fields，返回指定字段： [{"nbr_acct_id": "A001", "rel_base_amt": 1000, ...}, ...] - + 如果不指定return_fields，返回完整信息： [{"neighbor": {...}, "minHops": 1, "samplePath": ..., "rel": {...}}, ...] - + 如果 return_path_length=True，返回包含路径长度： [{"neighbor": {...}, "path_length": 1, ...}, ...] """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + if not (1 <= hops <= 10): raise ValueError("hops must be between 1 and 10") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 特殊处理：hops=1且指定return_fields且不需要路径长度时，使用简化查询（避免路径聚合） if hops == 1 and return_fields and not return_path_length: # 构建单跳关系模式 @@ -1108,7 +1172,7 @@ def neighbors_n_hop( pattern = f"(start)<-[r{rel}]-(nbr)" else: pattern = f"(start)-[r{rel}]-(nbr)" - + # 可选的 WHERE 子句 where_parts = [] if where: @@ -1116,7 +1180,7 @@ def neighbors_n_hop( if exclude_start: where_parts.append("nbr <> start") where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # 构建RETURN子句 return_parts = [] for field in return_fields: @@ -1132,15 +1196,19 @@ def neighbors_n_hop( # 默认当作邻居节点属性 prop = self._sanitize_property_key(field) return_parts.append(f"nbr.`{prop}` AS {prop}") - - return_clause = "RETURN " + (" DISTINCT " if return_distinct else " ") + ", ".join(return_parts) - + + return_clause = ( + "RETURN " + + (" DISTINCT " if return_distinct else " ") + + ", ".join(return_parts) + ) + # 可选的 ORDER BY 子句 order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (start:`{label}` {{`{key}`: $value}}) MATCH {pattern} @@ -1158,7 +1226,7 @@ def neighbors_n_hop( pattern = f"(start)<-[r{rel}*1..{hops}]-(nbr)" else: pattern = f"(start)-[r{rel}*1..{hops}]-(nbr)" - + # 可选的 WHERE 子句 where_parts = [] if where: @@ -1166,7 +1234,7 @@ def neighbors_n_hop( if exclude_start: where_parts.append("nbr <> start") where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # 构建 RETURN 子句 if return_path_length: # 返回路径长度 @@ -1186,7 +1254,7 @@ def neighbors_n_hop( collect(p)[0] AS samplePath, relationships(collect(p)[0])[0] AS rel """ - + # 可选的 ORDER BY 子句 if order_by: if order_by == "path_length" and return_path_length: @@ -1195,10 +1263,10 @@ def neighbors_n_hop( order_clause = f"ORDER BY {order_by} {order_direction}" else: order_clause = "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (start:`{label}` {{`{key}`: $value}}) MATCH p = {pattern} @@ -1209,11 +1277,11 @@ def neighbors_n_hop( {order_clause} {limit_clause} """ - + params = {"value": value} if limit is not None: params["limit"] = limit - + return self.run(cypher, params) # ========================================================= @@ -1230,15 +1298,15 @@ def common_neighbors( order_by: Optional[str] = None, order_direction: str = "ASC", limit: Optional[int] = None, - aggregate: bool = False + aggregate: bool = False, ) -> List[JsonDict]: """ 查询两个节点的公共一跳邻居（支持通用修饰符 + 聚合排序） - + 功能： - 查询两个节点的公共邻居 - 支持按交易次数聚合排序（aggregate=True） - + Args: a: (label, key, value) 节点A b: (label, key, value) 节点B @@ -1249,31 +1317,31 @@ def common_neighbors( order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回数量（None=不限制） aggregate: 是否启用聚合模式（按交易次数统计） - + Returns: 如果 aggregate=False（默认）： [{"commonNeighbor": {...}, "relA": {...}, "relB": {...}}, ...] - + 如果 aggregate=True： [{"commonNeighbor": {...}, "count": 2}, ...] """ (la, ka, va) = a (lb, kb, vb) = b - + la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 构建路径模式 if direction == "out": pat_a = f"(A)-[rA{rel}]->(C)" @@ -1284,21 +1352,25 @@ def common_neighbors( else: pat_a = f"(A)-[rA{rel}]-(C)" pat_b = f"(B)-[rB{rel}]-(C)" - + # 可选的 WHERE 子句 where_clause = f"WHERE {where}" if where else "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + # 聚合模式：统计每个公共邻居的交易次数 if aggregate: # 如果 order_by 是 "count"，使用聚合计数 if order_by == "count": order_clause = f"ORDER BY count {order_direction}" else: - order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "ORDER BY count DESC" - + order_clause = ( + f"ORDER BY {order_by} {order_direction}" + if order_by + else "ORDER BY count DESC" + ) + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1312,7 +1384,7 @@ def common_neighbors( else: # 普通模式：返回所有关系详情 order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1323,11 +1395,11 @@ def common_neighbors( {order_clause} {limit_clause} """ - + params = {"va": va, "vb": vb} if limit is not None: params["limit"] = limit - + return self.run(cypher, params) def common_neighbors_with_rel_filter( @@ -1342,21 +1414,21 @@ def common_neighbors_with_rel_filter( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 查询两个节点的公共邻居，支持对两条关系分别设置属性过滤条件 - + 功能： - 查询两个节点的公共一跳邻居 - 支持对A→C和B→C的关系分别设置属性过滤条件 - 支持对公共邻居节点C设置过滤条件 - 支持指定返回字段 - + 应用场景： - 找出两个账户的共同交易对手，且交易金额都大于某个值 - 找出两个用户的共同好友，且关系建立时间都在某个时间段内 - + 示例： # 找出Steven Collins和Samantha Cook的共同交易邻居，且交易金额都>400 results = client.common_neighbors_with_rel_filter( @@ -1367,7 +1439,7 @@ def common_neighbors_with_rel_filter( rel_conditions={"base_amt": (">", 400)}, return_fields=["C.node_key", "C.acct_id", "rA.base_amt", "rB.base_amt"] ) - + # 找出两个账户的共同交易对手，且都是大额交易（>1000）且交易时间在2025年 results = client.common_neighbors_with_rel_filter( a=("Account", "node_key", "Collins Steven"), @@ -1380,7 +1452,7 @@ def common_neighbors_with_rel_filter( return_fields=["C.node_key", "rA.base_amt", "rA.tran_timestamp", "rB.base_amt", "rB.tran_timestamp"] ) - + Args: a: (label, key, value) 节点A b: (label, key, value) 节点B @@ -1398,35 +1470,35 @@ def common_neighbors_with_rel_filter( order_by: 排序字段（如 "C.name" 或 "rA.base_amt"） order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回数量（None=不限制） - + Returns: 如果指定return_fields，返回指定字段： [{"C_node_key": "...", "rA_base_amt": 500, "rB_base_amt": 600}, ...] - + 如果不指定return_fields，返回完整信息： [{"commonNeighbor": {...}, "relA": {...}, "relB": {...}}, ...] - + 注意： - rel_conditions中的条件会同时应用到rA和rB（AND逻辑） - 如果需要对rA和rB设置不同的条件，请使用neighbor_where参数手动指定 """ (la, ka, va) = a (lb, kb, vb) = b - + la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 构建路径模式 if direction == "out": pat_a = f"(A)-[rA{rel}]->(C)" @@ -1437,31 +1509,35 @@ def common_neighbors_with_rel_filter( else: pat_a = f"(A)-[rA{rel}]-(C)" pat_b = f"(B)-[rB{rel}]-(C)" - + # 构建WHERE子句 where_parts = [] params = {"va": va, "vb": vb} - + # 处理关系属性条件（同时应用到rA和rB） if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name_a = f"rel_cond_a_{i}" param_name_b = f"rel_cond_b_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"rA.`{key}` {operator.upper()} ${param_name_a}") - where_parts.append(f"rB.`{key}` {operator.upper()} ${param_name_b}") + where_parts.append( + f"rA.`{key}` {operator.upper()} ${param_name_a}" + ) + where_parts.append( + f"rB.`{key}` {operator.upper()} ${param_name_b}" + ) elif operator.upper() == "STARTS WITH": where_parts.append(f"rA.`{key}` STARTS WITH ${param_name_a}") where_parts.append(f"rB.`{key}` STARTS WITH ${param_name_b}") else: where_parts.append(f"rA.`{key}` {operator} ${param_name_a}") where_parts.append(f"rB.`{key}` {operator} ${param_name_b}") - + params[param_name_a] = value params[param_name_b] = value else: @@ -1472,13 +1548,13 @@ def common_neighbors_with_rel_filter( where_parts.append(f"rB.`{key}` = ${param_name_b}") params[param_name_a] = condition params[param_name_b] = condition - + # 添加公共邻居节点的过滤条件 if neighbor_where: where_parts.append(f"({neighbor_where})") - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # 构建RETURN子句 if return_fields: return_parts = [] @@ -1497,15 +1573,15 @@ def common_neighbors_with_rel_filter( return_clause = "RETURN " + ", ".join(return_parts) else: return_clause = "RETURN C AS commonNeighbor, rA AS relA, rB AS relB" - + # 可选的 ORDER BY 子句 order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1516,7 +1592,7 @@ def common_neighbors_with_rel_filter( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # ========================================================= @@ -1532,11 +1608,11 @@ def filter_query( node_where: Optional[str] = None, rel_where: Optional[str] = None, params: Optional[JsonDict] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 从起点出发，按节点/关系条件过滤 - + Args: start: (label, key, value) 起点节点 rel_type: 关系类型（可选） @@ -1546,26 +1622,26 @@ def filter_query( rel_where: 关系过滤条件（如 "r.weight > $minW"） params: 额外参数 limit: 最大返回数量（None=不限制） - + Returns: [{"start": {...}, "rel": {...}, "node": {...}}, ...] - + 警告： node_where 和 rel_where 是字符串片段，需要自行确保安全性 """ (sl, sk, sv) = start sl = self._sanitize_label(sl) sk = self._sanitize_property_key(sk) - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + tl = self._sanitize_label(node_label) if node_label else "" tlabel = f":`{tl}`" if tl else "" - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # 构建路径模式 if direction == "out": pat = f"(s)-[r{rel}]->(n{tlabel})" @@ -1573,7 +1649,7 @@ def filter_query( pat = f"(s)<-[r{rel}]-(n{tlabel})" else: pat = f"(s)-[r{rel}]-(n{tlabel})" - + # 构建 WHERE 子句 where_parts = [] if rel_where: @@ -1581,10 +1657,10 @@ def filter_query( if node_where: where_parts.append(f"({node_where})") where_clause = ("WHERE " + " AND ".join(where_parts)) if where_parts else "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (s:`{sl}` {{`{sk}`: $sv}}) MATCH {pat} @@ -1592,13 +1668,13 @@ def filter_query( RETURN s AS start, r AS rel, n AS node {limit_clause} """ - + p = {"sv": sv} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) # ========================================================= @@ -1612,11 +1688,11 @@ def subgraph_extract( rel_type: Optional[str] = None, direction: str = "both", where: Optional[str] = None, - limit_paths: int = 200 + limit_paths: int = 200, ) -> JsonDict: """ 抽取以某节点为中心的子图（支持通用修饰符） - + Args: center: (label, key, value) 中心节点 hops: 半径（跳数） @@ -1624,20 +1700,20 @@ def subgraph_extract( direction: 方向 where: WHERE 过滤条件（如 "n.balance > 1000"） limit_paths: 最大路径数 - + Returns: {"nodes": [...], "relationships": [...]} """ (cl, ck, cv) = center cl = self._sanitize_label(cl) ck = self._sanitize_property_key(ck) - + if not (1 <= hops <= 5): raise ValueError("hops must be between 1 and 5") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 构建路径模式 if direction == "out": pat = f"(c)-[r{rel}*1..{hops}]->(n)" @@ -1645,10 +1721,10 @@ def subgraph_extract( pat = f"(c)<-[r{rel}*1..{hops}]-(n)" else: pat = f"(c)-[r{rel}*1..{hops}]-(n)" - + # 可选的 WHERE 子句 where_clause = f"WHERE {where}" if where else "" - + cypher = f""" MATCH (c:`{cl}` {{`{ck}`: $cv}}) MATCH p = {pat} @@ -1660,15 +1736,15 @@ def subgraph_extract( WITH collect(DISTINCT nn) AS nodes, collect(DISTINCT rr) AS relationships RETURN nodes, relationships """ - + res = self.run(cypher, {"cv": cv, "limit_paths": limit_paths}) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), - "relationships": res[0].get("relationships", []) + "relationships": res[0].get("relationships", []), } - + return {"nodes": [], "relationships": []} def subgraph_extract_by_nodes( @@ -1680,21 +1756,21 @@ def subgraph_extract_by_nodes( include_internal: bool = True, rel_type: Optional[str] = None, direction: str = "both", - where: Optional[str] = None + where: Optional[str] = None, ) -> JsonDict: """ 基于节点列表抽取子图（包含指定节点及其相互之间的关系） - + 功能： - 提取指定节点列表中所有节点 - 提取这些节点之间的所有关系 - 可选择是否包含节点内部的关系（如 A->A） - + 应用场景： 1. 交易网络：提取账户 A、B、C 及其之间的转账记录 2. 社交网络：提取指定用户群体及其相互关系 3. 知识图谱：提取指定实体及其关联关系 - + 示例： # 提取账户 A、B、C 及其之间的转账关系 subgraph = client.subgraph_extract_by_nodes( @@ -1704,7 +1780,7 @@ def subgraph_extract_by_nodes( rel_type="TRANSFER", direction="both" ) - + # 提取用户群体的社交关系 subgraph = client.subgraph_extract_by_nodes( "User", @@ -1713,7 +1789,7 @@ def subgraph_extract_by_nodes( rel_type="FOLLOWS", include_internal=False # 不包含自环 ) - + Args: label: 节点标签 key: 节点属性键 @@ -1722,7 +1798,7 @@ def subgraph_extract_by_nodes( rel_type: 关系类型（可选，None=任意类型） direction: 方向 ("out"=单向, "in"=反向, "both"=双向) where: WHERE 过滤条件（如 "r.amount > 1000"） - + Returns: { "nodes": [节点列表], @@ -1730,7 +1806,7 @@ def subgraph_extract_by_nodes( "node_count": 节点数量, "relationship_count": 关系数量 } - + 注意： - 只返回指定节点之间的关系，不会扩展到其他节点 - 如果某个节点不存在，会在结果中忽略 @@ -1738,17 +1814,17 @@ def subgraph_extract_by_nodes( """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + if not values or len(values) == 0: raise ValueError("values list cannot be empty") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # 构建关系模式 rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 构建路径模式 if direction == "out": pattern = f"(n1)-[r{rel}]->(n2)" @@ -1756,24 +1832,24 @@ def subgraph_extract_by_nodes( pattern = f"(n1)<-[r{rel}]-(n2)" else: pattern = f"(n1)-[r{rel}]-(n2)" - + # 构建 WHERE 子句 where_parts = [] - + # 节点必须在指定列表中 where_parts.append("n1.`" + key + "` IN $values") where_parts.append("n2.`" + key + "` IN $values") - + # 是否排除自环 if not include_internal: where_parts.append("n1 <> n2") - + # 用户自定义过滤条件 if where: where_parts.append(f"({where})") - + where_clause = "WHERE " + " AND ".join(where_parts) - + # Cypher 查询 cypher = f""" MATCH (n1:`{label}`) @@ -1789,22 +1865,22 @@ def subgraph_extract_by_nodes( size(allNodes) AS node_count, size(allRels) AS relationship_count """ - + res = self.run(cypher, {"values": values}) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), "relationships": res[0].get("relationships", []), "node_count": res[0].get("node_count", 0), - "relationship_count": res[0].get("relationship_count", 0) + "relationship_count": res[0].get("relationship_count", 0), } - + return { "nodes": [], "relationships": [], "node_count": 0, - "relationship_count": 0 + "relationship_count": 0, } def subgraph_extract_by_rel_filter( @@ -1815,21 +1891,21 @@ def subgraph_extract_by_rel_filter( start_label: Optional[str] = None, end_label: Optional[str] = None, direction: str = "both", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> JsonDict: """ 基于关系属性条件抽取子图（提取满足条件的关系及相关节点） - + 功能： - 根据关系属性条件筛选关系 - 提取满足条件的关系及其起点和终点节点 - 返回子图结构（节点 + 关系） - + 应用场景： 1. 时间范围：抽取某天/某时间段的所有交易子图 2. 金额范围：抽取金额在指定范围内的交易子图 3. 类型过滤：抽取特定类型的关系子图 - + 示例： # 抽取 2025-05-01 当天所有交易构成的子图 subgraph = client.subgraph_extract_by_rel_filter( @@ -1837,13 +1913,13 @@ def subgraph_extract_by_rel_filter( {"tran_timestamp": (">=", "2025-05-01"), "tran_timestamp": ("<", "2025-05-02")} ) - + # 抽取交易金额在 300 到 500 之间的交易子图 subgraph = client.subgraph_extract_by_rel_filter( "TRANSFER", {"base_amt": (">=", 300), "base_amt": ("<=", 500)} ) - + # 抽取可疑交易子图 subgraph = client.subgraph_extract_by_rel_filter( "TRANSFER", @@ -1851,7 +1927,7 @@ def subgraph_extract_by_rel_filter( start_label="Account", end_label="Account" ) - + Args: rel_type: 关系类型 rel_conditions: 关系属性条件字典，格式：{property: (operator, value)} @@ -1860,7 +1936,7 @@ def subgraph_extract_by_rel_filter( end_label: 终点节点标签（可选） direction: 方向 ("out"=单向, "in"=反向, "both"=双向) limit: 最大关系数量（None=不限制） - + Returns: { "nodes": [节点列表], @@ -1868,20 +1944,24 @@ def subgraph_extract_by_rel_filter( "node_count": 节点数量, "relationship_count": 关系数量 } - + 注意： - 返回的节点是满足条件的关系的起点和终点节点 - 如果需要计算统计信息（如交易总数），可以从 relationship_count 获取 """ rel_type = self._sanitize_rel_type(rel_type) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # 构建节点模式 - start_pattern = f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" - end_pattern = f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" - + start_pattern = ( + f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" + ) + end_pattern = ( + f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" + ) + # 构建关系模式 if direction == "out": rel_pattern = f"{start_pattern}-[r:`{rel_type}`]->{end_pattern}" @@ -1889,11 +1969,11 @@ def subgraph_extract_by_rel_filter( rel_pattern = f"{start_pattern}<-[r:`{rel_type}`]-{end_pattern}" else: rel_pattern = f"{start_pattern}-[r:`{rel_type}`]-{end_pattern}" - + # 构建WHERE子句 where_parts = [] params = {} - + if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): # 处理特殊的日期范围键名（如 tran_timestamp_start, tran_timestamp_end） @@ -1903,32 +1983,36 @@ def subgraph_extract_by_rel_filter( actual_key = self._sanitize_property_key(actual_key) else: actual_key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name = f"rel_cond_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"r.`{actual_key}` {operator.upper()} ${param_name}") + where_parts.append( + f"r.`{actual_key}` {operator.upper()} ${param_name}" + ) elif operator.upper() == "STARTS WITH": - where_parts.append(f"r.`{actual_key}` STARTS WITH ${param_name}") + where_parts.append( + f"r.`{actual_key}` STARTS WITH ${param_name}" + ) else: where_parts.append(f"r.`{actual_key}` {operator} ${param_name}") - + params[param_name] = value else: # 简单等值条件 param_name = f"rel_cond_{i}" where_parts.append(f"r.`{actual_key}` = ${param_name}") params[param_name] = condition - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + # Cypher 查询 cypher = f""" MATCH {rel_pattern} @@ -1941,22 +2025,22 @@ def subgraph_extract_by_rel_filter( size(allNodes) AS node_count, size(allRels) AS relationship_count """ - + res = self.run(cypher, params) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), "relationships": res[0].get("relationships", []), "node_count": res[0].get("node_count", 0), - "relationship_count": res[0].get("relationship_count", 0) + "relationship_count": res[0].get("relationship_count", 0), } - + return { "nodes": [], "relationships": [], "node_count": 0, - "relationship_count": 0 + "relationship_count": 0, } # ========================================================= @@ -1968,43 +2052,43 @@ def match_path_pattern( pattern: str, where: Optional[str] = None, params: Optional[JsonDict] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 自定义路径模式查询 - + 示例: pattern = "(a:User {userId:$uid})-[:FOLLOWS]->(b:User)-[:POSTED]->(p:Post)" where = "p.createdAt >= $since" - + Args: pattern: 路径模式 where: WHERE 子句（可选） params: 参数 limit: 最大返回数量（None=不限制） - + Returns: [{"path": ...}, ...] - + 警告： pattern 和 where 是字符串片段，需要自行确保安全性 """ # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH p = {pattern} {("WHERE " + where) if where else ""} RETURN p AS path {limit_clause} """ - + p = {} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) # ========================================================= @@ -2018,15 +2102,15 @@ def aggregate_stats( where: Optional[str] = None, params: Optional[JsonDict] = None, metrics: Optional[Sequence[str]] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 聚合统计查询 - + 示例： # 统计每个国家的用户数 aggregate_stats("User", group_by="country") - + # 统计年龄>=18的用户，按城市分组 aggregate_stats( "User", @@ -2034,7 +2118,7 @@ def aggregate_stats( where="n.age >= $minAge", params={"minAge": 18} ) - + Args: label: 节点标签 group_by: 分组字段（可选） @@ -2042,21 +2126,21 @@ def aggregate_stats( params: 额外参数 metrics: 聚合指标（默认 count(*)） limit: 最大返回数量（None=不限制） - + Returns: 聚合结果列表 """ label = self._sanitize_label(label) metrics = list(metrics) if metrics else ["count(*) AS cnt"] - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + if group_by: group_by = self._sanitize_property_key(group_by) group_expr = f"n.`{group_by}` AS {group_by}" return_expr = ", ".join([group_expr] + list(metrics)) - + cypher = f""" MATCH (n:`{label}`) {("WHERE " + where) if where else ""} @@ -2067,21 +2151,22 @@ def aggregate_stats( """ else: return_expr = ", ".join(metrics) - + cypher = f""" MATCH (n:`{label}`) {("WHERE " + where) if where else ""} RETURN {return_expr} {limit_clause} """ - + p = {} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) + # ========================================================= # 两点间路径查询（完整版，不限定最短） # ========================================================= @@ -2098,11 +2183,11 @@ def paths_between( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ 查询两个节点之间的路径（支持通用修饰符 + 复杂过滤/排序/计算） - + 功能： - 支持指定最小/最大跳数 - 支持指定关系类型和方向 @@ -2110,7 +2195,7 @@ def paths_between( - ⚠️ **支持自定义返回字段**（可以计算路径总金额、最低金额等） - ⚠️ **支持自定义排序**（按金额、时间、路径长度等） - 可返回多条路径（不限定必须是最短） - + 示例： # 例1：查询路径，按转账金额最大排序 paths = client.paths_between( @@ -2121,7 +2206,7 @@ def paths_between( order_by="maxAmount", order_direction="DESC" ) - + # 例4：查询路径，要求路径上包含 is_sar 的交易 paths = client.paths_between( ("Account", "node_key", "A"), @@ -2129,7 +2214,7 @@ def paths_between( rel_type="TRANSFER", where="ANY(r IN relationships(p) WHERE r.is_sar = true)" ) - + # 例8：计算每条路径的总金额 paths = client.paths_between( ("Account", "node_key", "A"), @@ -2139,7 +2224,7 @@ def paths_between( order_by="totalAmount", order_direction="DESC" ) - + # 例9：查询路径，要求所有交易金额都 < 1000 paths = client.paths_between( ("Account", "node_key", "A"), @@ -2147,7 +2232,7 @@ def paths_between( rel_type="TRANSFER", where="ALL(r IN relationships(p) WHERE r.base_amt < 1000)" ) - + # 例10：查询路径，要求不经过 bank 节点 paths = client.paths_between( ("Account", "node_key", "A"), @@ -2155,7 +2240,7 @@ def paths_between( rel_type="TRANSFER", where="ALL(n IN nodes(p) WHERE n.bank_id <> 'bank')" ) - + Args: a: (label, key, value) 起点节点 b: (label, key, value) 终点节点 @@ -2178,7 +2263,7 @@ def paths_between( order_by: 排序字段（如 "hops", "totalAmount", "maxAmount"） order_direction: 排序方向 ("ASC"=升序, "DESC"=降序) limit: 最大返回路径数（None=不限制） - + Returns: [ { @@ -2192,7 +2277,7 @@ def paths_between( }, ... ] - + 注意： - 默认按路径长度升序返回（短路径优先） - 如果两点不连通，返回空列表 @@ -2201,26 +2286,26 @@ def paths_between( """ (la, ka, va) = a (lb, kb, vb) = b - + # 参数验证 la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if not (0 <= min_hops <= max_hops <= 10): raise ValueError("Invalid hop bounds: 0 <= min_hops <= max_hops <= 10") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + # 构建关系模式 rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # 构建路径模式 if direction == "out": pattern = f"(A)-[r{rel}*{min_hops}..{max_hops}]->(B)" @@ -2228,17 +2313,17 @@ def paths_between( pattern = f"(A)<-[r{rel}*{min_hops}..{max_hops}]-(B)" else: pattern = f"(A)-[r{rel}*{min_hops}..{max_hops}]-(B)" - + # 可选的 WHERE 子句 where_clause = f"WHERE {where}" if where else "" - + # 构建 RETURN 子句 if return_fields: # 自定义返回字段 return_parts = [] for field in return_fields: field_lower = field.lower() - + # 基础字段 if field_lower == "path": return_parts.append("p AS path") @@ -2248,27 +2333,35 @@ def paths_between( return_parts.append("pathNodes AS nodes") elif field_lower == "relationships": return_parts.append("pathRels AS relationships") - + # 计算字段：总金额 elif field_lower == "totalamount": - return_parts.append("REDUCE(s = 0, r IN pathRels | s + r.base_amt) AS totalAmount") - + return_parts.append( + "REDUCE(s = 0, r IN pathRels | s + r.base_amt) AS totalAmount" + ) + # 计算字段：最大金额 elif field_lower == "maxamount": - return_parts.append("REDUCE(m = 0, r IN pathRels | CASE WHEN r.base_amt > m THEN r.base_amt ELSE m END) AS maxAmount") - + return_parts.append( + "REDUCE(m = 0, r IN pathRels | CASE WHEN r.base_amt > m THEN r.base_amt ELSE m END) AS maxAmount" + ) + # 计算字段：最小金额 elif field_lower == "minamount": - return_parts.append("REDUCE(m = 999999, r IN pathRels | CASE WHEN r.base_amt < m THEN r.base_amt ELSE m END) AS minAmount") - + return_parts.append( + "REDUCE(m = 999999, r IN pathRels | CASE WHEN r.base_amt < m THEN r.base_amt ELSE m END) AS minAmount" + ) + # 计算字段：平均金额 elif field_lower == "avgamount": - return_parts.append("REDUCE(s = 0, r IN pathRels | s + r.base_amt) / hops AS avgAmount") - + return_parts.append( + "REDUCE(s = 0, r IN pathRels | s + r.base_amt) / hops AS avgAmount" + ) + # 其他字段：直接使用 else: return_parts.append(field) - + return_clause = "RETURN " + ", ".join(return_parts) else: # 默认返回字段 @@ -2276,16 +2369,16 @@ def paths_between( hops, pathNodes AS nodes, pathRels AS relationships""" - + # 可选的 ORDER BY 子句（默认按 hops 升序） if order_by: order_clause = f"ORDER BY {order_by} {order_direction}" else: order_clause = "ORDER BY hops ASC" - + # 可选的 LIMIT 子句 limit_clause = "LIMIT $limit" if limit is not None else "" - + # Cypher 查询 cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}), (B:`{lb}` {{`{kb}`: $vb}}) @@ -2296,13 +2389,12 @@ def paths_between( {order_clause} {limit_clause} """ - + params = {"va": va, "vb": vb} if limit is not None: params["limit"] = limit - - return self.run(cypher, params) + return self.run(cypher, params) # ========================================== @@ -2314,47 +2406,42 @@ def paths_between( Neo4jConfig( uri="bolt://localhost:7687", user="neo4j", - password="password" + password=os.getenv("NEO4J_PASSWORD", ""), ) ) as client: - # 获取 Schema schema = client.get_schema() print("=== Schema 信息 ===") print(f"节点类型: {list(schema['node_labels'].keys())}") print(f"关系类型: {list(schema['relationship_types'].keys())}") print(f"模式样例: {schema['patterns'][:3]}\n") - + # 1. 唯一键查节点 print("=== 测试1: 唯一键查节点 ===") user = client.get_node_by_unique_key("User", "userId", "u123") print(f"用户: {user}\n") - + # 2. N跳邻居 print("=== 测试2: N跳邻居 ===") neighbors = client.neighbors_n_hop( - "User", "userId", "u123", + "User", + "userId", + "u123", hops=2, rel_type="FOLLOWS", direction="out", - limit=10 + limit=10, ) print(f"找到 {len(neighbors)} 个邻居\n") - + # 3. 公共邻居 print("=== 测试3: 公共邻居 ===") common = client.common_neighbors( - ("User", "userId", "u1"), - ("User", "userId", "u2"), - rel_type="FOLLOWS" + ("User", "userId", "u1"), ("User", "userId", "u2"), rel_type="FOLLOWS" ) print(f"找到 {len(common)} 个公共邻居\n") - + # 4. 聚合统计 print("=== 测试4: 聚合统计 ===") - stats = client.aggregate_stats( - "User", - group_by="country", - limit=5 - ) - print(f"统计结果: {stats}\n") \ No newline at end of file + stats = client.aggregate_stats("User", group_by="country", limit=5) + print(f"统计结果: {stats}\n") diff --git a/aag/computing_engine/graph_query/load_data_into_neo4j.py b/aag/computing_engine/graph_query/load_data_into_neo4j.py index e506203..3dd8921 100644 --- a/aag/computing_engine/graph_query/load_data_into_neo4j.py +++ b/aag/computing_engine/graph_query/load_data_into_neo4j.py @@ -1,3 +1,4 @@ +import os import pandas as pd from neo4j import GraphDatabase import logging @@ -7,21 +8,27 @@ import yaml # 配置日志 -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) logger = logging.getLogger(__name__) + class Neo4jGraphLoader: """ 通用的Neo4j图数据加载器，支持根据schema动态加载图数据 """ - def __init__(self, - uri: str = 'bolt://219.216.65.209:7687', - username: str = 'neo4j', - password: str = '12345678', - schema: Optional[Dict[str, Any]] = None): + + def __init__( + self, + uri: str = "bolt://219.216.65.209:7687", + username: str = "neo4j", + password: str = os.getenv("NEO4J_PASSWORD", ""), + schema: Optional[Dict[str, Any]] = None, + ): """ 初始化Neo4j图数据加载器 - + Args: uri: Neo4j连接URI，格式为 bolt://host:port 或 neo4j://host:port username: 用户名 @@ -33,13 +40,15 @@ def __init__(self, self.password = password self.driver = None self.schema = schema - + self.connect_to_neo4j() - + def connect_to_neo4j(self): """连接到Neo4j""" try: - self.driver = GraphDatabase.driver(self.uri, auth=(self.username, self.password)) + self.driver = GraphDatabase.driver( + self.uri, auth=(self.username, self.password) + ) # 验证连接 with self.driver.session() as session: result = session.run("RETURN 1 AS test") @@ -52,7 +61,7 @@ def connect_to_neo4j(self): self.driver.close() self.driver = None raise Exception(f"无法连接到Neo4j数据库: {e}") - + def clear_database(self): """清空数据库（可选）""" if not self.driver: @@ -67,36 +76,41 @@ def clear_database(self): except Exception as e: logger.error(f"清空数据库失败: {e}") return False - - def _build_node_key(self, row: pd.Series, query_fields: Union[str, List[str]]) -> str: + + def _build_node_key( + self, row: pd.Series, query_fields: Union[str, List[str]] + ) -> str: """ 根据query_field构建node_key - + Args: row: DataFrame的一行数据 query_fields: 单个字段名或字段名列表 - + Returns: 构建的node_key字符串 """ if isinstance(query_fields, str): query_fields = [query_fields] - + # 将多个字段值用空格连接 values = [str(row[field]) for field in query_fields] - return ' '.join(values) - - def _get_attribute_fields(self, df: pd.DataFrame, - attribute_fields: Optional[List[str]], - exclude_fields: Optional[List[str]] = None) -> List[str]: + return " ".join(values) + + def _get_attribute_fields( + self, + df: pd.DataFrame, + attribute_fields: Optional[List[str]], + exclude_fields: Optional[List[str]] = None, + ) -> List[str]: """ 获取需要加载的属性字段列表 - + Args: df: DataFrame attribute_fields: schema中指定的属性字段列表，如果为空则使用所有列 exclude_fields: 需要排除的字段列表 - + Returns: 属性字段列表 """ @@ -107,11 +121,11 @@ def _get_attribute_fields(self, df: pd.DataFrame, # 如果没有指定，使用所有列，但排除exclude_fields exclude_fields = exclude_fields or [] return [col for col in df.columns if col not in exclude_fields] - + def create_constraints_and_indexes(self, vertex_configs: List[Dict[str, Any]]): """ 根据vertex配置创建约束和索引 - + Args: vertex_configs: vertex配置列表 """ @@ -121,10 +135,10 @@ def create_constraints_and_indexes(self, vertex_configs: List[Dict[str, Any]]): try: with self.driver.session() as session: for vertex_config in vertex_configs: - vertex_type = vertex_config.get('type', 'Node') + vertex_type = vertex_config.get("type", "Node") # 将type首字母大写作为标签 label = vertex_type.capitalize() - + # 为node_key创建唯一性约束 try: session.run(f""" @@ -135,23 +149,27 @@ def create_constraints_and_indexes(self, vertex_configs: List[Dict[str, Any]]): logger.info(f"成功创建{label}.node_key唯一性约束") except Exception as e: logger.warning(f"创建约束时出现问题（可能已存在）: {e}") - + # 创建node_key索引 try: - session.run(f"CREATE INDEX {label.lower()}_key_index IF NOT EXISTS FOR (n:{label}) ON (n.node_key)") + session.run( + f"CREATE INDEX {label.lower()}_key_index IF NOT EXISTS FOR (n:{label}) ON (n.node_key)" + ) logger.info(f"成功创建{label}.node_key索引") except Exception as e: logger.warning(f"创建索引时出现问题（可能已存在）: {e}") - + # 为id_field创建索引（如果存在） - id_field = vertex_config.get('id_field') + id_field = vertex_config.get("id_field") if id_field: try: - session.run(f"CREATE INDEX {label.lower()}_id_index IF NOT EXISTS FOR (n:{label}) ON (n.{id_field})") + session.run( + f"CREATE INDEX {label.lower()}_id_index IF NOT EXISTS FOR (n:{label}) ON (n.{id_field})" + ) logger.info(f"成功创建{label}.{id_field}索引") except Exception as e: logger.warning(f"创建索引时出现问题（可能已存在）: {e}") - + return True except Exception as e: logger.error(f"创建约束和索引时发生错误: {e}") @@ -160,7 +178,7 @@ def create_constraints_and_indexes(self, vertex_configs: List[Dict[str, Any]]): def load_vertices(self, vertex_config: Dict[str, Any]): """ 根据配置加载顶点数据 - + Args: vertex_config: 顶点配置，包含以下字段： - path: CSV文件路径 @@ -173,65 +191,69 @@ def load_vertices(self, vertex_config: Dict[str, Any]): if not self.driver: logger.error("数据库连接未建立，无法加载顶点数据") return False - + try: # 读取配置 - file_path = vertex_config.get('path') - vertex_type = vertex_config.get('type', 'Node') - query_field = vertex_config.get('query_field') - id_field = vertex_config.get('id_field') - label_field = vertex_config.get('label_field') - attribute_fields = vertex_config.get('attribute_fields', []) - + file_path = vertex_config.get("path") + vertex_type = vertex_config.get("type", "Node") + query_field = vertex_config.get("query_field") + id_field = vertex_config.get("id_field") + label_field = vertex_config.get("label_field") + attribute_fields = vertex_config.get("attribute_fields", []) + # 将type首字母大写作为Neo4j标签 neo4j_label = vertex_type.capitalize() - + # 读取数据 df = pd.read_csv(file_path) logger.info(f"读取到 {len(df)} 条{vertex_type}数据") - + # 替换NaN值为None df = df.where(pd.notna(df), None) - + # 构建node_key if isinstance(query_field, list): - df['node_key'] = df.apply(lambda row: self._build_node_key(row, query_field), axis=1) + df["node_key"] = df.apply( + lambda row: self._build_node_key(row, query_field), axis=1 + ) else: - df['node_key'] = df[query_field].astype(str) - + df["node_key"] = df[query_field].astype(str) + # 确定需要加载的属性字段 # 排除node_key（因为已经单独处理） - exclude_fields = ['node_key'] - fields_to_load = self._get_attribute_fields(df, attribute_fields, exclude_fields) - + exclude_fields = ["node_key"] + fields_to_load = self._get_attribute_fields( + df, attribute_fields, exclude_fields + ) + # 确保id_field和label_field在加载列表中 if id_field and id_field not in fields_to_load: fields_to_load.append(id_field) if label_field and label_field not in fields_to_load: fields_to_load.append(label_field) - + logger.info(f"将加载以下属性字段: {fields_to_load}") - + batch_size = 1000 total_vertices = len(df) success_count = 0 fail_count = 0 - + with self.driver.session() as session: for i in range(0, total_vertices, batch_size): - batch = df.iloc[i:i+batch_size] - + batch = df.iloc[i : i + batch_size] + # 只选择需要的列 - columns_to_use = ['node_key'] + fields_to_load + columns_to_use = ["node_key"] + fields_to_load batch_subset = batch[columns_to_use] - records = batch_subset.to_dict('records') - + records = batch_subset.to_dict("records") + # 动态构建SET子句 set_clauses = [] for field in fields_to_load: set_clauses.append(f"n.{field} = record.{field}") set_clause = ",\n ".join(set_clauses) - + # 构建Cypher查询 query = f""" UNWIND $records AS record @@ -241,28 +263,34 @@ def load_vertices(self, vertex_config: Dict[str, Any]): ON MATCH SET {set_clause} """ - + try: result = session.run(query, records=records) summary = result.consume() success_count += len(records) - logger.info(f"已处理 {min(i+batch_size, total_vertices)}/{total_vertices} 条{vertex_type}数据") + logger.info( + f"已处理 {min(i + batch_size, total_vertices)}/{total_vertices} 条{vertex_type}数据" + ) except Exception as e: logger.error(f"批量插入{vertex_type}数据失败: {e}") fail_count += len(records) - - logger.info(f"{vertex_type}数据加载完成，成功处理 {success_count} 条，失败 {fail_count} 条") + + logger.info( + f"{vertex_type}数据加载完成，成功处理 {success_count} 条，失败 {fail_count} 条" + ) return True - + except Exception as e: logger.error(f"加载{vertex_type}数据时发生错误: {e}") return False # add gjq: 修改边加载逻辑，通过id_field查找对应的node_key来匹配节点 - def load_edges(self, edge_config: Dict[str, Any], vertex_configs: List[Dict[str, Any]]): + def load_edges( + self, edge_config: Dict[str, Any], vertex_configs: List[Dict[str, Any]] + ): """ 根据配置加载边数据 - + Args: edge_config: 边配置，包含以下字段： - path: CSV文件路径 @@ -277,89 +305,105 @@ def load_edges(self, edge_config: Dict[str, Any], vertex_configs: List[Dict[str, if not self.driver: logger.error("数据库连接未建立，无法加载边数据") return False - + try: # 读取配置 - file_path = edge_config.get('path') - edge_type = edge_config.get('type', 'EDGE') - source_field = edge_config.get('source_field') - target_field = edge_config.get('target_field') - label_field = edge_config.get('label_field') - weight_field = edge_config.get('weight_field') - attribute_fields = edge_config.get('attribute_fields', []) - + file_path = edge_config.get("path") + edge_type = edge_config.get("type", "EDGE") + source_field = edge_config.get("source_field") + target_field = edge_config.get("target_field") + label_field = edge_config.get("label_field") + weight_field = edge_config.get("weight_field") + attribute_fields = edge_config.get("attribute_fields", []) + # 将type转换为大写作为Neo4j关系类型 neo4j_rel_type = edge_type.upper() - + # add gjq: 获取顶点配置信息 vertex_config = vertex_configs[0] if vertex_configs else {} - vertex_label = vertex_config.get('type', 'Node').capitalize() - vertex_id_field = vertex_config.get('id_field', 'id') - vertex_query_field = vertex_config.get('query_field') - vertex_path = vertex_config.get('path') - + vertex_label = vertex_config.get("type", "Node").capitalize() + vertex_id_field = vertex_config.get("id_field", "id") + vertex_query_field = vertex_config.get("query_field") + vertex_path = vertex_config.get("path") + # add gjq: 读取顶点数据，建立id_field到node_key的映射 logger.info(f"正在读取顶点数据以建立ID到node_key的映射...") vertex_df = pd.read_csv(vertex_path) vertex_df = vertex_df.where(pd.notna(vertex_df), None) - + # add gjq: 构建node_key列 if isinstance(vertex_query_field, list): - vertex_df['node_key'] = vertex_df.apply(lambda row: self._build_node_key(row, vertex_query_field), axis=1) + vertex_df["node_key"] = vertex_df.apply( + lambda row: self._build_node_key(row, vertex_query_field), axis=1 + ) else: - vertex_df['node_key'] = vertex_df[vertex_query_field].astype(str) - + vertex_df["node_key"] = vertex_df[vertex_query_field].astype(str) + # add gjq: 创建id到node_key的映射字典 - id_to_node_key = dict(zip(vertex_df[vertex_id_field], vertex_df['node_key'])) + id_to_node_key = dict( + zip(vertex_df[vertex_id_field], vertex_df["node_key"]) + ) logger.info(f"成功建立 {len(id_to_node_key)} 个ID到node_key的映射") - + # 读取边数据 df = pd.read_csv(file_path) logger.info(f"读取到 {len(df)} 条{edge_type}数据") - + # 替换NaN值为None df = df.where(pd.notna(df), None) - + # add gjq: 将source_field和target_field的值转换为对应的node_key - df['source_node_key'] = df[source_field].map(id_to_node_key) - df['target_node_key'] = df[target_field].map(id_to_node_key) - + df["source_node_key"] = df[source_field].map(id_to_node_key) + df["target_node_key"] = df[target_field].map(id_to_node_key) + # add gjq: 检查是否有无法映射的节点 - missing_source = df['source_node_key'].isna().sum() - missing_target = df['target_node_key'].isna().sum() + missing_source = df["source_node_key"].isna().sum() + missing_target = df["target_node_key"].isna().sum() if missing_source > 0 or missing_target > 0: - logger.warning(f"发现 {missing_source} 个源节点和 {missing_target} 个目标节点无法映射到node_key，这些边将被跳过") + logger.warning( + f"发现 {missing_source} 个源节点和 {missing_target} 个目标节点无法映射到node_key，这些边将被跳过" + ) # 过滤掉无法映射的边 - df = df.dropna(subset=['source_node_key', 'target_node_key']) + df = df.dropna(subset=["source_node_key", "target_node_key"]) logger.info(f"过滤后剩余 {len(df)} 条有效边数据") - + # 确定需要加载的属性字段 # 排除source_field、target_field、source_node_key、target_node_key - exclude_fields = [source_field, target_field, 'source_node_key', 'target_node_key'] - fields_to_load = self._get_attribute_fields(df, attribute_fields, exclude_fields) - + exclude_fields = [ + source_field, + target_field, + "source_node_key", + "target_node_key", + ] + fields_to_load = self._get_attribute_fields( + df, attribute_fields, exclude_fields + ) + # 确保label_field和weight_field在加载列表中 if label_field and label_field not in fields_to_load: fields_to_load.append(label_field) if weight_field and weight_field not in fields_to_load: fields_to_load.append(weight_field) - + logger.info(f"将加载以下边属性字段: {fields_to_load}") - + batch_size = 1000 total_edges = len(df) success_count = 0 fail_count = 0 - + with self.driver.session() as session: for i in range(0, total_edges, batch_size): - batch = df.iloc[i:i+batch_size] - + batch = df.iloc[i : i + batch_size] + # add gjq: 选择需要的列，包括source_node_key和target_node_key - columns_to_use = ['source_node_key', 'target_node_key'] + fields_to_load + columns_to_use = [ + "source_node_key", + "target_node_key", + ] + fields_to_load batch_subset = batch[columns_to_use] - records = batch_subset.to_dict('records') - + records = batch_subset.to_dict("records") + # 动态构建属性设置子句 if fields_to_load: prop_clauses = [] @@ -368,7 +412,7 @@ def load_edges(self, edge_config: Dict[str, Any], vertex_configs: List[Dict[str, properties = "{" + ", ".join(prop_clauses) + "}" else: properties = "" - + # add gjq: 构建Cypher查询 - 通过node_key匹配节点 query = f""" UNWIND $records AS record @@ -376,25 +420,31 @@ def load_edges(self, edge_config: Dict[str, Any], vertex_configs: List[Dict[str, MATCH (target:{vertex_label} {{node_key: record.target_node_key}}) CREATE (source)-[r:{neo4j_rel_type} {properties}]->(target) """ - + try: result = session.run(query, records=records) summary = result.consume() batch_success = summary.counters.relationships_created success_count += batch_success - logger.info(f"已处理 {min(i+batch_size, total_edges)}/{total_edges} 条{edge_type}数据，成功创建 {batch_success} 条关系") + logger.info( + f"已处理 {min(i + batch_size, total_edges)}/{total_edges} 条{edge_type}数据，成功创建 {batch_success} 条关系" + ) except Exception as e: logger.error(f"批量插入{edge_type}数据失败: {e}") fail_count += len(records) - - logger.info(f"{edge_type}数据加载完成，成功插入 {success_count} 条关系，失败 {fail_count} 条") + + logger.info( + f"{edge_type}数据加载完成，成功插入 {success_count} 条关系，失败 {fail_count} 条" + ) return True - + except Exception as e: logger.error(f"加载{edge_type}数据时发生错误: {e}") return False - def verify_data(self, vertex_configs: List[Dict[str, Any]], edge_configs: List[Dict[str, Any]]): + def verify_data( + self, vertex_configs: List[Dict[str, Any]], edge_configs: List[Dict[str, Any]] + ): """验证数据加载结果""" if not self.driver: logger.error("数据库连接未建立，无法验证数据") @@ -403,36 +453,42 @@ def verify_data(self, vertex_configs: List[Dict[str, Any]], edge_configs: List[D with self.driver.session() as session: # 检查每种类型的节点数量 for vertex_config in vertex_configs: - vertex_type = vertex_config.get('type', 'Node').capitalize() - result = session.run(f"MATCH (n:{vertex_type}) RETURN count(n) as count") + vertex_type = vertex_config.get("type", "Node").capitalize() + result = session.run( + f"MATCH (n:{vertex_type}) RETURN count(n) as count" + ) count = result.single()["count"] logger.info(f"{vertex_type}节点数量: {count}") - + # 检查每种类型的关系数量 for edge_config in edge_configs: - edge_type = edge_config.get('type', 'EDGE').upper() - result = session.run(f"MATCH ()-[r:{edge_type}]->() RETURN count(r) as count") + edge_type = edge_config.get("type", "EDGE").upper() + result = session.run( + f"MATCH ()-[r:{edge_type}]->() RETURN count(r) as count" + ) count = result.single()["count"] logger.info(f"{edge_type}关系数量: {count}") - + # 显示示例节点 if vertex_configs: - vertex_type = vertex_configs[0].get('type', 'Node').capitalize() - query_field = vertex_configs[0].get('query_field') + vertex_type = vertex_configs[0].get("type", "Node").capitalize() + query_field = vertex_configs[0].get("query_field") if isinstance(query_field, list): - field_names = ', '.join([f"n.{f}" for f in query_field]) + field_names = ", ".join([f"n.{f}" for f in query_field]) else: field_names = f"n.{query_field}" - - result = session.run(f"MATCH (n:{vertex_type}) RETURN n.node_key, {field_names} LIMIT 5") + + result = session.run( + f"MATCH (n:{vertex_type}) RETURN n.node_key, {field_names} LIMIT 5" + ) logger.info(f"示例{vertex_type}节点的node_key:") for record in result: logger.info(f" {dict(record)}") - + # 显示示例关系 if edge_configs and vertex_configs: - edge_type = edge_configs[0].get('type', 'EDGE').upper() - vertex_type = vertex_configs[0].get('type', 'Node').capitalize() + edge_type = edge_configs[0].get("type", "EDGE").upper() + vertex_type = vertex_configs[0].get("type", "Node").capitalize() result = session.run(f""" MATCH (source:{vertex_type})-[r:{edge_type}]->(target:{vertex_type}) RETURN source.node_key as source_key, target.node_key as target_key @@ -440,10 +496,12 @@ def verify_data(self, vertex_configs: List[Dict[str, Any]], edge_configs: List[D """) logger.info(f"示例{edge_type}关系:") for record in result: - logger.info(f" {record['source_key']} -> {record['target_key']}") - + logger.info( + f" {record['source_key']} -> {record['target_key']}" + ) + return True - + except Exception as e: logger.error(f"验证数据时发生错误: {e}") return False @@ -453,164 +511,166 @@ def close_connection(self): if self.driver: self.driver.close() logger.info("已关闭Neo4j连接") - + def load_all_data(self, clear_existing: bool = False): """ 加载所有数据的主流程 - + Args: clear_existing: 是否清空现有数据（默认False） """ if not self.schema: logger.error("未提供schema信息，无法加载数据") return False - + try: # 1. 可选：清空现有数据 if clear_existing: if not self.clear_database(): logger.warning("清空数据库失败，继续执行...") - + # 获取vertex和edge配置 - vertex_configs = self.schema.get('vertex', []) - edge_configs = self.schema.get('edge', []) - + vertex_configs = self.schema.get("vertex", []) + edge_configs = self.schema.get("edge", []) + if not vertex_configs: logger.error("schema中没有vertex配置") return False - + # 2. 创建约束和索引 if not self.create_constraints_and_indexes(vertex_configs): logger.warning("创建约束和索引失败，继续执行...") - + # 3. 加载所有顶点 for vertex_config in vertex_configs: if not self.load_vertices(vertex_config): logger.error(f"加载顶点类型 {vertex_config.get('type')} 失败") return False - + # 4. 加载所有边 for edge_config in edge_configs: if not self.load_edges(edge_config, vertex_configs): logger.error(f"加载边类型 {edge_config.get('type')} 失败") return False - + # 5. 验证数据 if not self.verify_data(vertex_configs, edge_configs): return False - + logger.info("所有数据加载完成！") return True - + except Exception as e: logger.error(f"加载数据时发生错误: {e}") return False finally: self.close_connection() -def load_from_yaml(yaml_path: str, - uri: str = 'bolt://localhost:7687', - username: str = 'neo4j', - password: str = 'password', - clear_existing: bool = False) -> bool: + +def load_from_yaml( + yaml_path: str, + uri: str = "bolt://localhost:7687", + username: str = "neo4j", + password: str = "password", + clear_existing: bool = False, +) -> bool: """ 从YAML文件加载schema并导入数据到Neo4j - + Args: yaml_path: YAML schema文件路径 uri: Neo4j连接URI username: Neo4j用户名 password: Neo4j密码 clear_existing: 是否清空现有数据 - + Returns: 是否成功 """ try: # 读取YAML文件 - with open(yaml_path, 'r', encoding='utf-8') as f: + with open(yaml_path, "r", encoding="utf-8") as f: config = yaml.safe_load(f) - + # 获取第一个数据集的schema - datasets = config.get('datasets', []) + datasets = config.get("datasets", []) if not datasets: logger.error("YAML文件中没有找到datasets") return False - + dataset = datasets[0] - schema = dataset.get('schema') - + schema = dataset.get("schema") + if not schema: logger.error("数据集中没有找到schema") return False - + logger.info(f"正在加载数据集: {dataset.get('name')}") - + # 创建加载器并加载数据 loader = Neo4jGraphLoader( - uri=uri, - username=username, - password=password, - schema=schema + uri=uri, username=username, password=password, schema=schema ) - + return loader.load_all_data(clear_existing=clear_existing) - + except Exception as e: logger.error(f"从YAML加载数据失败: {e}") return False + def main(): """主函数 - 示例用法""" # 方式1: 直接使用schema字典 schema = { - 'vertex': [ + "vertex": [ { - 'path': '/home/gaojq/AAG/aag/datasets/graphs/transaction_amlsim/1K/accounts.csv', - 'type': 'account', - 'query_field': ['last_name', 'first_name'], # 使用多个字段构建node_key - 'id_field': 'acct_id', - 'label_field': 'prior_sar_count', - 'attribute_fields': [] # 空列表表示加载所有列 + "path": "/home/gaojq/AAG/aag/datasets/graphs/transaction_amlsim/1K/accounts.csv", + "type": "account", + "query_field": ["last_name", "first_name"], # 使用多个字段构建node_key + "id_field": "acct_id", + "label_field": "prior_sar_count", + "attribute_fields": [], # 空列表表示加载所有列 } ], - 'edge': [ + "edge": [ { - 'path': '/home/gaojq/AAG/aag/datasets/graphs/transaction_amlsim/1K/transactions.csv', - 'type': 'transfer', - 'source_field': 'orig_acct', - 'target_field': 'bene_acct', - 'label_field': 'is_sar', - 'weight_field': None, - 'attribute_fields': [] # 空列表表示加载所有列 + "path": "/home/gaojq/AAG/aag/datasets/graphs/transaction_amlsim/1K/transactions.csv", + "type": "transfer", + "source_field": "orig_acct", + "target_field": "bene_acct", + "label_field": "is_sar", + "weight_field": None, + "attribute_fields": [], # 空列表表示加载所有列 } - ] + ], } - + loader = Neo4jGraphLoader( - uri='bolt://202.199.13.67:7687', - username='neo4j', - password='12345678', - schema=schema + uri="bolt://202.199.13.67:7687", + username="neo4j", + password=os.getenv("NEO4J_PASSWORD", ""), + schema=schema, ) - + success = loader.load_all_data(clear_existing=True) - + if success: logger.info("数据加载成功完成！") sys.exit(0) else: logger.error("数据加载失败！") sys.exit(1) - + # 方式2: 从YAML文件加载 # success = load_from_yaml( # yaml_path='AAG_3/AAG/aag/datasets/dataset_schemas/AMLSim1K/graph_schemas.yaml', # uri='bolt://202.199.13.67:7687', # username='neo4j', - # password='12345678', + # password=os.getenv("NEO4J_PASSWORD", ""), # 从环境变量读取密码 # clear_existing=True # ) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/aag/computing_engine/graph_query/nl_query_engine.py b/aag/computing_engine/graph_query/nl_query_engine.py index 5dfade6..b0cb772 100644 --- a/aag/computing_engine/graph_query/nl_query_engine.py +++ b/aag/computing_engine/graph_query/nl_query_engine.py @@ -28,12 +28,16 @@ # ===== 导入你的模板模块 ===== from aag.computing_engine.graph_query.graph_query import Neo4jGraphClient, Neo4jConfig -#from graph_query import Neo4jGraphClient, Neo4jConfig -from aag.computing_engine.graph_query.nl_query_engine_refactored import ParameterExtractorRouter + +# from graph_query import Neo4jGraphClient, Neo4jConfig +from aag.computing_engine.graph_query.nl_query_engine_refactored import ( + ParameterExtractorRouter, +) # add gjq -# Removed direct OpenAI client initialization -# os.environ['OPENAI_API_KEY'] = 'sk-G30rFStBigqXtuyIOkOo7Zh4QNxO8ZAjfZQ5DYPCgMXbPv8q' -# os.environ['OPENAI_BASE_URL'] = 'https://gitaigc.com/v1/' +# 已移除直接初始化 OpenAI 客户端的硬编码配置 +# 原 API Key 和 Base URL 已从环境变量读取，请勿在此处回填敏感信息 +# os.environ['OPENAI_API_KEY'] = '<从环境变量读取>' +# os.environ['OPENAI_BASE_URL'] = '<从环境变量读取>' # client = OpenAI( # api_key=os.environ.get("OPENAI_API_KEY"), # base_url=os.environ.get("OPENAI_BASE_URL") @@ -43,6 +47,7 @@ # 1. Basic Query Types (Core Query Logic) # ========================================== from aag.computing_engine.graph_query.templates import QUERY_TEMPLATES, QUERY_MODIFIERS + # Standard JSON output template STANDARD_OUTPUT_SCHEMA = { "type": "object", @@ -50,7 +55,7 @@ "properties": { "params": { "type": "object", - "description": "Query-specific required and optional parameters" + "description": "Query-specific required and optional parameters", }, "modifiers": { "type": "object", @@ -60,37 +65,44 @@ "order_direction": {"type": "string", "enum": ["ASC", "DESC"]}, "limit": {"type": "integer", "description": "Result count limit"}, "where": {"type": "string", "description": "Filter condition"}, - "aggregate": {"type": "string", "enum": ["count", "sum", "avg", "max", "min"]}, - "aggregate_field": {"type": "string", "description": "Aggregation field"} - } - } - } + "aggregate": { + "type": "string", + "enum": ["count", "sum", "avg", "max", "min"], + }, + "aggregate_field": { + "type": "string", + "description": "Aggregation field", + }, + }, + }, + }, } + # ========================================== # 2. LLM 接口抽象（你需要替换成真实实现） # ========================================== # add gjq class LLMInterface: """LLM 调用接口（使用 Reasoner 统一调用）""" - + def __init__(self, reasoner: Reasoner): """ 初始化 LLM 接口 - + Args: reasoner: Reasoner 实例，用于统一调用不同的 LLM """ self.reasoner = reasoner - + def call(self, prompt: str, **kwargs) -> str: """ 调用 LLM，返回文本响应 - + 使用 Reasoner 的 generate_response 方法 """ response = self.reasoner.generate_response(prompt) - if hasattr(response, 'text'): + if hasattr(response, "text"): return response.text return str(response) @@ -101,16 +113,16 @@ def call(self, prompt: str, **kwargs) -> str: # add gjq class QueryTypeClassifier: """使用 LLM 判断查询类型""" - + def __init__(self, reasoner: Reasoner): """ 初始化查询类型分类器 - + Args: reasoner: Reasoner 实例，用于调用 LLM """ self.reasoner = reasoner - + def classify(self, question: str) -> str: """ 返回查询类型（如 "node_lookup"） @@ -122,99 +134,133 @@ def classify(self, question: str) -> str: # ⚠️ CRITICAL FIX: 修正 LLM 可能返回的错误类型名 # LLM 可能返回 "subgraph_extract_by_nodes" 而不是正确的 "subgraph_by_nodes" if query_type == "subgraph_extract_by_nodes": - logging.warning(f"LLM 返回了错误的类型名 '{query_type}'，自动修正为 'subgraph_by_nodes'") + logging.warning( + f"LLM 返回了错误的类型名 '{query_type}'，自动修正为 'subgraph_by_nodes'" + ) query_type = "subgraph_by_nodes" - + if query_type not in QUERY_TEMPLATES: # 降级到规则匹配 logging.warning(f"未知查询类型: {query_type}，使用规则匹配降级") query_type = self._fallback_classify(question) - + return query_type - + def _fallback_classify(self, question: str) -> str: """规则匹配降级方案""" q = question.lower() - + # 优先级从高到低 # 0. 检测是否有具体的起点节点（人名、ID等） - name_pattern = r'\b([A-Z][a-z]+\s+[A-Z][a-z]+)\b' + name_pattern = r"\b([A-Z][a-z]+\s+[A-Z][a-z]+)\b" has_specific_node = bool(re.search(name_pattern, question)) - + # 1. 检测聚合查询（最高优先级） # 1.1 分组聚合：包含"每个"关键词 if any(kw in q for kw in ["统计每个", "计算每个", "每个", "各个", "分别"]): # ⚠️ CRITICAL: 即使有"每个"，如果有具体节点，也应该是 neighbor_query if has_specific_node: - logging.info(f"检测到具体节点 + '每个'关键词，使用 neighbor_query + aggregate") + logging.info( + f"检测到具体节点 + '每个'关键词，使用 neighbor_query + aggregate" + ) return "neighbor_query" return "aggregation_query" - + # 1.2 全局聚合：包含聚合关键词（统计、计算、数量、总额等） - if any(kw in q for kw in ["统计", "计算", "数量", "次数", "总数", "总金额", "金额合计", "平均", "最大", "最小", "top", "前", "排序", "排名", "占比"]): + if any( + kw in q + for kw in [ + "统计", + "计算", + "数量", + "次数", + "总数", + "总金额", + "金额合计", + "平均", + "最大", + "最小", + "top", + "前", + "排序", + "排名", + "占比", + ] + ): # ⚠️ CRITICAL: 如果有具体的起点节点，应该是 neighbor_query + aggregate if has_specific_node: - logging.info(f"检测到具体节点 + 聚合关键词，使用 neighbor_query + aggregate") + logging.info( + f"检测到具体节点 + 聚合关键词，使用 neighbor_query + aggregate" + ) return "neighbor_query" - + # 排除"找出"、"列出"、"返回"等返回列表的关键词 - if not any(kw in q for kw in ["找出", "列出", "返回...的交易", "查询...的交易"]): + if not any( + kw in q for kw in ["找出", "列出", "返回...的交易", "查询...的交易"] + ): return "aggregation_query" - + # 2. 检测关系过滤查询（关键词：交易、转账、关系属性） if any(kw in q for kw in ["交易", "转账", "金额", "is_sar", "tran_"]): # ⚠️ 关键区分：如果包含人名或具体节点标识，是邻居查询 - name_pattern = r'\b([A-Z][a-z]+\s+[A-Z][a-z]+)\b' + name_pattern = r"\b([A-Z][a-z]+\s+[A-Z][a-z]+)\b" if re.search(name_pattern, question): # 包含人名，很可能是邻居查询 if any(kw in q for kw in ["对手", "邻居", "直接", "发生过"]): return "neighbor_query" - + # 如果包含关系属性条件（大于、小于等），且要求返回列表 if any(kw in q for kw in ["大于", "小于", "等于", ">", "<", "="]): if any(kw in q for kw in ["找出", "列出", "返回", "查询"]): return "relationship_filter" - + # 如果包含"所有"且没有明确起点，是关系过滤 if "所有" in q and not any(kw in q for kw in ["的", "作为"]): return "relationship_filter" - + # 检测公共邻居查询（支持简单模式和过滤模式） if "公共" in q or "共同" in q or ("既" in q and "又" in q): # 检查是否是节点属性条件（如 "为true"、"为false"） # 这种情况目前不支持，需要自定义Cypher if any(kw in q for kw in ["为true", "为false"]): - logging.warning(f"检测到复杂的节点属性条件公共邻居查询，当前模板不支持，建议使用自定义Cypher") + logging.warning( + f"检测到复杂的节点属性条件公共邻居查询，当前模板不支持，建议使用自定义Cypher" + ) return "node_lookup" # 降级处理 - + # 统一返回 common_neighbor，由参数提取器判断是简单模式还是过滤模式 return "common_neighbor" - + if any(kw in q for kw in ["路径", "怎么到", "到达"]): return "path_query" - + # ⚠️ CRITICAL FIX: 优先检测多节点子图查询（在检测单节点邻居查询之前） # 检测多节点子图查询的关键特征： # 1. 包含"子图"关键词 # 2. 包含节点列表标识（如 "A、B、C" 或 "[A,B,C]" 或 "之间"） # 3. 包含"包含"、"及其"、"相互"、"连边"等关键词 - if "子图" in q or any(kw in q for kw in ["包含节点", "节点列表", "相互之间", "连边"]): + if "子图" in q or any( + kw in q for kw in ["包含节点", "节点列表", "相互之间", "连边"] + ): # 检测是否包含多个节点标识 # 方式1：逗号分隔（如 "A、B、C" 或 "A,B,C"） # 方式2：列表格式（如 "[A,B,C]"） # 方式3：关键词（如 "之间"、"相互"、"互相"） - if any(kw in q for kw in ["、", ",", "[", "之间", "相互", "互相", "及其", "连边"]): + if any( + kw in q + for kw in ["、", ",", "[", "之间", "相互", "互相", "及其", "连边"] + ): logging.info("检测到多节点子图查询关键词，返回 subgraph_by_nodes") return "subgraph_by_nodes" # 如果只有"子图"但没有多节点标识，是单中心子图 if "子图" in q: return "subgraph" - + if any(kw in q for kw in ["邻居", "朋友", "关注", "关系"]): return "neighbor_query" if any(kw in q for kw in ["找", "查", "获取"]): return "node_lookup" - + # 默认使用邻居查询 return "neighbor_query" @@ -224,23 +270,23 @@ def _fallback_classify(self, question: str) -> str: # ========================================== class SchemaAnalyzer: """图 Schema 分析工具""" - + def __init__(self, client: Neo4jGraphClient): self.client = client self._schema_cache = None - + def get_schema(self) -> Dict: """获取并缓存 Schema 信息""" if self._schema_cache is None: self._schema_cache = self.client.get_schema() return self._schema_cache - + def format_for_llm(self) -> str: """格式化 Schema 信息给 LLM（增强版）""" schema = self.get_schema() - + formatted = "## 图数据库 Schema 信息\n\n" - + # 节点类型及属性（包含示例值） formatted += "### 节点类型及属性\n" for label, info in schema["node_labels"].items(): @@ -251,7 +297,7 @@ def format_for_llm(self) -> str: sample = samples.get(prop, "") sample_str = f" (示例: {sample})" if sample else "" formatted += f" - `{prop}`{sample_str}\n" - + # 关系类型及属性（包含示例值） formatted += "\n### 关系类型及属性\n" for rel_type, info in schema["relationship_types"].items(): @@ -265,33 +311,37 @@ def format_for_llm(self) -> str: formatted += f" - `{prop}`{sample_str}\n" else: formatted += f"- **{rel_type}**: 无属性\n" - + # 关系模式 formatted += "\n### 常见关系模式\n" for pattern in schema["patterns"][:10]: formatted += f"- {pattern}\n" - + return formatted # add gjq class CypherValidator: """Use LLM to validate generated Cypher statements""" - - def __init__(self, reasoner: Reasoner, schema_analyzer: Optional['SchemaAnalyzer'] = None): + + def __init__( + self, reasoner: Reasoner, schema_analyzer: Optional["SchemaAnalyzer"] = None + ): self.reasoner = reasoner self.schema_analyzer = schema_analyzer - - def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> Dict: + + def validate( + self, cypher: str, question: str, query_type: str, params: Dict + ) -> Dict: """ Validate if Cypher statement conforms to syntax rules and user question requirements - + Args: cypher: Generated Cypher statement question: User's original question query_type: Query type params: Extracted parameters - + Returns: Validation result dictionary { "is_valid": bool, @@ -304,7 +354,7 @@ def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> schema_info = "" if self.schema_analyzer: schema_info = self.schema_analyzer.format_for_llm() - + # Get query template information template_info = "" template_cypher_example = "" @@ -313,14 +363,14 @@ def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> template_info = f""" ## Query Template Information - **Template Name**: {query_type} -- **Description**: {template.get('description', '')} -- **Method Used**: {template.get('method', '')} -- **Required Parameters**: {template.get('required_params', [])} -- **Optional Parameters**: {template.get('optional_params', [])} +- **Description**: {template.get("description", "")} +- **Method Used**: {template.get("method", "")} +- **Required Parameters**: {template.get("required_params", [])} +- **Optional Parameters**: {template.get("optional_params", [])} ⚠️ **Important**: Please refer to this template information to validate parameters and generate Cypher statements. """ - + # Provide Cypher examples based on query type if query_type == "subgraph": template_cypher_example = """ @@ -354,7 +404,7 @@ def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> ORDER BY c.acct_id ``` """ - + # add gjq # Use Reasoner's nl_query_validate_cypher method, directly call model_deployment.py method try: @@ -365,7 +415,7 @@ def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> params=params, schema_info=schema_info, template_info=template_info, - template_cypher_example=template_cypher_example + template_cypher_example=template_cypher_example, ) return result except Exception as e: @@ -375,63 +425,77 @@ def validate(self, cypher: str, question: str, query_type: str, params: Dict) -> "is_valid": True, "issues": [], "suggestions": [], - "corrected_cypher": None + "corrected_cypher": None, } - # ========================================== # 6. 查询执行器 # ========================================== class QueryExecutor: """执行查询模板（支持聚合）""" - - def __init__(self, client: Neo4jGraphClient, schema_analyzer: Optional['SchemaAnalyzer'] = None): + + def __init__( + self, + client: Neo4jGraphClient, + schema_analyzer: Optional["SchemaAnalyzer"] = None, + ): self.client = client self.schema_analyzer = schema_analyzer - - def _apply_aggregation(self, results: List[JsonDict], aggregate_type: str, aggregate_field: Optional[str] = None) -> List[JsonDict]: + + def _apply_aggregation( + self, + results: List[JsonDict], + aggregate_type: str, + aggregate_field: Optional[str] = None, + ) -> List[JsonDict]: """ 对查询结果应用聚合 - + Args: results: 原始查询结果 aggregate_type: 聚合类型 (count, sum, avg, max, min) aggregate_field: 聚合字段（可选，count 不需要） - + Returns: 聚合后的结果 """ # 统一转换为小写，避免大小写不匹配问题 aggregate_type = aggregate_type.lower() - + if not results: return [{"aggregate_type": aggregate_type, "value": 0}] - + if aggregate_type == "count": return [{"aggregate_type": "count", "value": len(results)}] - + # 其他聚合类型需要指定字段 if not aggregate_field: return [{"error": f"聚合类型 {aggregate_type} 需要指定 aggregate_field"}] - + # 提取字段值（支持嵌套字段，如 rel.base_amt） values = [] for item in results: try: # 支持嵌套访问，如 "rel.base_amt" value = item - for key in aggregate_field.split('.'): + for key in aggregate_field.split("."): value = value.get(key, {}) - + if isinstance(value, (int, float)): values.append(value) except (AttributeError, TypeError): continue - + if not values: - return [{"aggregate_type": aggregate_type, "value": None, "note": "没有找到有效的数值"}] - + return [ + { + "aggregate_type": aggregate_type, + "value": None, + "note": "没有找到有效的数值", + } + ] + # 执行聚合 if aggregate_type == "sum": result_value = sum(values) @@ -443,19 +507,21 @@ def _apply_aggregation(self, results: List[JsonDict], aggregate_type: str, aggre result_value = min(values) else: return [{"error": f"未知的聚合类型: {aggregate_type}"}] - - return [{ - "aggregate_type": aggregate_type, - "field": aggregate_field, - "value": result_value, - "count": len(values) - }] - + + return [ + { + "aggregate_type": aggregate_type, + "field": aggregate_field, + "value": result_value, + "count": len(values), + } + ] + # add gjq def execute(self, query_type: str, params: Dict) -> Dict: """根据类型和参数执行查询（支持聚合）""" template = QUERY_TEMPLATES[query_type] - + # ⚠️ CRITICAL: 对于 aggregation_query，不要 pop aggregate_field # 因为后续转换为 relationship_filter 时还需要使用 if query_type == "aggregation_query": @@ -466,7 +532,7 @@ def execute(self, query_type: str, params: Dict) -> Dict: # 其他查询类型，pop 出来 aggregate_type = params.pop("aggregate", None) aggregate_field = params.pop("aggregate_field", None) - + try: # 根据不同方法调整参数格式 if query_type == "node_lookup": @@ -476,7 +542,7 @@ def execute(self, query_type: str, params: Dict) -> Dict: # LLM 可能返回 {"field": "direction"} 格式，需要拆分 order_by_param = params.get("order_by") order_direction_param = params.get("order_direction", "ASC") - + if isinstance(order_by_param, dict): # 如果 order_by 是字典格式 {"field": "direction"} # 提取字段名和方向 @@ -487,7 +553,7 @@ def execute(self, query_type: str, params: Dict) -> Dict: order_direction_param = direction else: order_by_param = None - + # 多节点筛选模式 results = self.client.filter_nodes_by_properties( params["label"], @@ -495,7 +561,7 @@ def execute(self, query_type: str, params: Dict) -> Dict: return_fields=params.get("return_fields"), order_by=order_by_param, order_direction=order_direction_param, - limit=params.get("limit") + limit=params.get("limit"), ) else: # 单节点查找模式（支持return_fields） @@ -503,10 +569,10 @@ def execute(self, query_type: str, params: Dict) -> Dict: params["label"], params["key"], params["value"], - return_fields=params.get("return_fields") + return_fields=params.get("return_fields"), ) results = [result] if result else [] - + elif query_type == "relationship_filter": # 关系过滤查询 results = self.client.filter_relationships( @@ -519,18 +585,22 @@ def execute(self, query_type: str, params: Dict) -> Dict: aggregate_field=params.get("aggregate_field"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), - limit=params.get("limit") + limit=params.get("limit"), ) - + elif query_type == "aggregation_query": # 聚合统计查询 # 检查是否有分组参数 - has_grouping = params.get("group_by_node") or params.get("group_by_property") - + has_grouping = params.get("group_by_node") or params.get( + "group_by_property" + ) + # 如果没有分组参数，这应该是全局聚合查询 if not has_grouping: - logging.warning("检测到 aggregation_query 缺少分组参数，转换为全局聚合查询") - + logging.warning( + "检测到 aggregation_query 缺少分组参数，转换为全局聚合查询" + ) + # 场景1：节点属性聚合（有 conditions） if params.get("conditions"): logging.info("转换为 node_lookup + aggregate") @@ -541,81 +611,121 @@ def execute(self, query_type: str, params: Dict) -> Dict: return_fields=params.get("return_fields"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), - limit=params.get("limit") + limit=params.get("limit"), ) - + # 应用聚合 agg_type = params.get("aggregate_type", "COUNT").upper() if agg_type == "COUNT": - results = [{"aggregate_type": "count", "value": len(results)}] + results = [ + {"aggregate_type": "count", "value": len(results)} + ] else: # 其他聚合类型需要 aggregate_field agg_field = params.get("aggregate_field") if agg_field: - results = self._apply_aggregation(results, agg_type, agg_field) + results = self._apply_aggregation( + results, agg_type, agg_field + ) else: - results = [{"error": f"聚合类型 {agg_type} 需要指定 aggregate_field"}] - + results = [ + { + "error": f"聚合类型 {agg_type} 需要指定 aggregate_field" + } + ] + # 场景2：关系属性聚合（有 rel_type） elif params.get("rel_type"): logging.info("转换为 relationship_filter + aggregate") - + # ⚠️ CRITICAL: 获取聚合类型和字段 # 注意：aggregate_field 可能在 params 中，也可能已经被 pop 出去了 agg_type = params.get("aggregate_type") agg_field = params.get("aggregate_field") - + # 如果 aggregate_field 不在 params 中，尝试从外层获取 if not agg_field and aggregate_field: agg_field = aggregate_field logging.info(f"从外层获取 aggregate_field: {agg_field}") - - logging.info(f"aggregate_type: {agg_type}, aggregate_field: {agg_field}") - + + logging.info( + f"aggregate_type: {agg_type}, aggregate_field: {agg_field}" + ) + # ⚠️ CRITICAL: 确保 aggregate_field 正确传递 # 对于 SUM/AVG/MAX/MIN，aggregate_field 是必需的 - if agg_type and agg_type.upper() in ["SUM", "AVG", "MAX", "MIN"]: + if agg_type and agg_type.upper() in [ + "SUM", + "AVG", + "MAX", + "MIN", + ]: if not agg_field: # 尝试从 rel_type 的 Schema 中推断默认字段 if self.schema_analyzer: schema = self.schema_analyzer.get_schema() - rel_info = schema.get("relationship_types", {}).get(params["rel_type"], {}) + rel_info = schema.get("relationship_types", {}).get( + params["rel_type"], {} + ) rel_props = rel_info.get("properties", []) - + # 优先使用常见的金额字段 - for common_field in ["base_amt", "amount", "value", "total"]: + for common_field in [ + "base_amt", + "amount", + "value", + "total", + ]: if common_field in rel_props: agg_field = common_field - logging.warning(f"aggregate_field 未指定，自动推断为: {agg_field}") + logging.warning( + f"aggregate_field 未指定，自动推断为: {agg_field}" + ) break - + if not agg_field: # 如果还是没有，使用第一个数值型属性 for prop in rel_props: - if prop not in ["tran_id", "alert_id", "is_sar", "tx_type", "tran_timestamp"]: + if prop not in [ + "tran_id", + "alert_id", + "is_sar", + "tx_type", + "tran_timestamp", + ]: agg_field = prop - logging.warning(f"aggregate_field 未指定，使用第一个可能的数值字段: {agg_field}") + logging.warning( + f"aggregate_field 未指定，使用第一个可能的数值字段: {agg_field}" + ) break else: - logging.error("schema_analyzer 未初始化，无法自动推断 aggregate_field") - + logging.error( + "schema_analyzer 未初始化，无法自动推断 aggregate_field" + ) + # 转换为 relationship_filter 查询 results = self.client.filter_relationships( params["rel_type"], - start_label=params.get("start_label") or params.get("node_label"), - end_label=params.get("end_label") or params.get("node_label"), + start_label=params.get("start_label") + or params.get("node_label"), + end_label=params.get("end_label") + or params.get("node_label"), rel_conditions=params.get("rel_conditions"), aggregate=agg_type, aggregate_field=agg_field, return_fields=params.get("return_fields"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), - limit=params.get("limit") + limit=params.get("limit"), ) - + else: - results = [{"error": "aggregation_query 缺少必要参数：需要 group_by_node/group_by_property 或 conditions 或 rel_type"}] - + results = [ + { + "error": "aggregation_query 缺少必要参数：需要 group_by_node/group_by_property 或 conditions 或 rel_type" + } + ] + else: # 正常的分组聚合查询 results = self.client.aggregation_query( @@ -630,9 +740,9 @@ def execute(self, query_type: str, params: Dict) -> Dict: where=params.get("where"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "DESC"), - limit=params.get("limit") + limit=params.get("limit"), ) - + elif query_type == "neighbor_query": results = self.client.neighbors_n_hop( params["label"], @@ -648,13 +758,15 @@ def execute(self, query_type: str, params: Dict) -> Dict: limit=params.get("limit"), return_distinct=params.get("return_distinct", False), exclude_start=params.get("exclude_start", False), - return_path_length=params.get("return_path_length", False) + return_path_length=params.get("return_path_length", False), ) - + # 如果需要聚合，处理结果 if aggregate_type: - results = self._apply_aggregation(results, aggregate_type, aggregate_field) - + results = self._apply_aggregation( + results, aggregate_type, aggregate_field + ) + elif query_type == "common_neighbor": # 判断是简单模式还是过滤模式 if "rel_conditions" in params and params["rel_conditions"]: @@ -669,14 +781,16 @@ def execute(self, query_type: str, params: Dict) -> Dict: return_fields=params.get("return_fields"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), - limit=params.get("limit") + limit=params.get("limit"), ) else: # 简单模式：使用 common_neighbors # ⚠️ CRITICAL: 检测是否需要聚合排序 # 如果有 aggregate 参数且 order_by 是 "count"，启用聚合模式 - enable_aggregate = aggregate_type and aggregate_type.lower() == "count" - + enable_aggregate = ( + aggregate_type and aggregate_type.lower() == "count" + ) + results = self.client.common_neighbors( (params["label"], params["key"], params["v1"]), (params["label"], params["key"], params["v2"]), @@ -686,15 +800,17 @@ def execute(self, query_type: str, params: Dict) -> Dict: order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), limit=params.get("limit"), - aggregate=enable_aggregate + aggregate=enable_aggregate, ) - + # ⚠️ CRITICAL: 公共邻居的聚合处理 # 如果需要聚合，但不是 count 类型（已经在上面处理了） if aggregate_type and aggregate_type.lower() != "count": # 其他聚合类型（SUM、AVG等） - results = self._apply_aggregation(results, aggregate_type, aggregate_field) - + results = self._apply_aggregation( + results, aggregate_type, aggregate_field + ) + elif query_type == "path_query": results = self.client.paths_between( (params["label"], params["key"], params["v1"]), @@ -707,13 +823,15 @@ def execute(self, query_type: str, params: Dict) -> Dict: return_fields=params.get("return_fields"), order_by=params.get("order_by"), order_direction=params.get("order_direction", "ASC"), - limit=params.get("limit") + limit=params.get("limit"), ) - + # 如果需要聚合，处理结果 if aggregate_type: - results = self._apply_aggregation(results, aggregate_type, aggregate_field) - + results = self._apply_aggregation( + results, aggregate_type, aggregate_field + ) + elif query_type == "global_stats": results = self.client.aggregate_stats( params["label"], @@ -721,9 +839,9 @@ def execute(self, query_type: str, params: Dict) -> Dict: where=params.get("where"), params=params.get("params"), metrics=params.get("metrics"), - limit=params.get("limit") + limit=params.get("limit"), ) - + elif query_type == "subgraph": # 判断是单中心节点模式还是关系属性过滤模式 if "rel_conditions" in params and params["rel_conditions"]: @@ -734,7 +852,7 @@ def execute(self, query_type: str, params: Dict) -> Dict: start_label=params.get("start_label"), end_label=params.get("end_label"), direction=params.get("direction", "both"), - limit=params.get("limit") + limit=params.get("limit"), ) else: # 单中心节点模式：使用 subgraph_extract @@ -744,21 +862,25 @@ def execute(self, query_type: str, params: Dict) -> Dict: rel_type=params.get("rel_type"), direction=params.get("direction", "both"), where=params.get("where"), - limit_paths=params.get("limit_paths", 200) + limit_paths=params.get("limit_paths", 200), ) - + # 如果需要聚合（如计算交易总数），应用聚合 if aggregate_type: if aggregate_type.lower() == "count": # 计算关系总数（交易总数） - rel_count = result.get("relationship_count", len(result.get("relationships", []))) + rel_count = result.get( + "relationship_count", len(result.get("relationships", [])) + ) results = [{"aggregate_type": "count", "value": rel_count}] else: # 其他聚合类型（如 SUM、AVG 等） - results = self._apply_aggregation([result], aggregate_type, aggregate_field) + results = self._apply_aggregation( + [result], aggregate_type, aggregate_field + ) else: results = [result] - + elif query_type == "subgraph_by_nodes": result = self.client.subgraph_extract_by_nodes( params["label"], @@ -767,21 +889,25 @@ def execute(self, query_type: str, params: Dict) -> Dict: include_internal=params.get("include_internal", True), rel_type=params.get("rel_type"), direction=params.get("direction", "both"), - where=params.get("where") + where=params.get("where"), ) - + # ⚠️ CRITICAL FIX: 如果需要聚合（如计算节点总数），应用聚合 if aggregate_type: if aggregate_type.lower() == "count": # 计算节点总数 - node_count = result.get("node_count", len(result.get("nodes", []))) + node_count = result.get( + "node_count", len(result.get("nodes", [])) + ) results = [{"aggregate_type": "count", "value": node_count}] else: # 其他聚合类型（如 SUM、AVG 等） - results = self._apply_aggregation([result], aggregate_type, aggregate_field) + results = self._apply_aggregation( + [result], aggregate_type, aggregate_field + ) else: results = [result] - + elif query_type == "filter_query": results = self.client.filter_query( (params["label"], params["key"], params["value"]), @@ -791,30 +917,32 @@ def execute(self, query_type: str, params: Dict) -> Dict: node_where=params.get("node_where"), rel_where=params.get("rel_where"), params=params.get("params"), - limit=params.get("limit") + limit=params.get("limit"), ) - + # 如果需要聚合，处理结果 if aggregate_type: - results = self._apply_aggregation(results, aggregate_type, aggregate_field) - + results = self._apply_aggregation( + results, aggregate_type, aggregate_field + ) + else: return {"success": False, "error": f"未知查询类型: {query_type}"} - + return { "success": True, "query_type": query_type, "params": params, "results": results, - "count": len(results) + "count": len(results), } - + except Exception as e: return { "success": False, "error": str(e), "query_type": query_type, - "params": params + "params": params, } @@ -824,23 +952,27 @@ def execute(self, query_type: str, params: Dict) -> Dict: # add gjq class NaturalLanguageQueryEngine: """自然语言查询引擎（双 LLM + Schema 工具）""" - - def __init__(self, db_client: Neo4jGraphClient, llm: LLMInterface, enable_validation: bool = True): + + def __init__( + self, + db_client: Neo4jGraphClient, + llm: LLMInterface, + enable_validation: bool = True, + ): self.client = db_client self.llm = llm # add gjq: log_manager 未定义，移除此行 # self.log_manager = log_manager self.enable_validation = enable_validation - + # 初始化各组件 self.schema_analyzer = SchemaAnalyzer(db_client) # add gjq: QueryTypeClassifier 需要 Reasoner 对象，从 LLMInterface 中获取 self.type_classifier = QueryTypeClassifier(llm.reasoner) - - + self.param_extractor = ParameterExtractorRouter(llm, self.schema_analyzer) logging.info("✅ 使用重构后的参数提取器路由器") - + # 初始化 Cypher 验证器（传入 schema_analyzer） # add gjq: CypherValidator 需要 Reasoner 对象，从 LLMInterface 中获取 if enable_validation: @@ -849,22 +981,22 @@ def __init__(self, db_client: Neo4jGraphClient, llm: LLMInterface, enable_valida else: self.cypher_validator = None logging.info("⚠️ Cypher 验证器已禁用") - + self.executor = QueryExecutor(db_client, self.schema_analyzer) - + def initialize(self): """初始化：加载 Schema""" print("\n🔍 正在分析图数据库结构...") schema = self.schema_analyzer.get_schema() - + print(f"✅ Schema 加载完成") print(f" 节点类型: {list(schema['node_labels'].keys())}") print(f" 关系类型: {list(schema['relationship_types'].keys())}\n") - + def ask(self, question: str) -> Dict: """ 主入口：处理自然语言问题 - + 流程： 1. LLM1 判断查询类型 2. 获取 Schema 信息 @@ -873,24 +1005,24 @@ def ask(self, question: str) -> Dict: """ print(f"\n💬 用户问题: {question}") print("=" * 80) - + # Step 1: LLM1 判断查询类型 print("🤖 LLM1 正在判断查询类型...") query_type = self.type_classifier.classify(question) print(f" 查询类型: {query_type}") print(f" 说明: {QUERY_TEMPLATES[query_type]['description']}") - + # Step 2: 获取 Schema（已缓存，不会重复查询） print("\n📊 获取图 Schema 信息...") schema_info = self.schema_analyzer.format_for_llm() print(" Schema 已加载") - + # 调试：打印 Schema 信息（可选） if os.environ.get("DEBUG_SCHEMA"): - print("\n" + "="*80) + print("\n" + "=" * 80) print(schema_info) - print("="*80) - + print("=" * 80) + # Step 3: LLM2 提取参数 print("\n🤖 LLM2 正在提取参数...") try: @@ -899,46 +1031,52 @@ def ask(self, question: str) -> Dict: except Exception as e: error_msg = str(e) print(f" ❌ 参数提取失败: {error_msg}") - + # ⚠️ CRITICAL FIX: 检测查询类型判断错误，自动重试 if "查询类型判断错误" in error_msg and "neighbor_query" in error_msg: print("\n🔄 检测到查询类型判断错误，自动切换到 neighbor_query 重试...") query_type = "neighbor_query" print(f" 新查询类型: {query_type}") print(f" 说明: {QUERY_TEMPLATES[query_type]['description']}") - + try: params = self.param_extractor.extract(question, query_type) - print(f" 提取参数: {json.dumps(params, ensure_ascii=False, indent=2)}") + print( + f" 提取参数: {json.dumps(params, ensure_ascii=False, indent=2)}" + ) except Exception as e2: print(f" ❌ 重试后仍然失败: {e2}") return {"success": False, "error": f"参数提取失败: {e2}"} else: return {"success": False, "error": f"参数提取失败: {error_msg}"} - + # Step 3.5: 验证参数（可选） corrected_cypher = None # ⚠️ 对于 path_query、subgraph 和 common_neighbor 查询类型，跳过 LLM3 验证 # 原因：这些查询类型使用固定的模板，不需要 LLM3 动态修正 # 解决方案：直接增强模板本身的能力，支持更多参数 skip_validation_types = ["path_query", "subgraph"] - - if self.enable_validation and self.cypher_validator and query_type not in skip_validation_types: + + if ( + self.enable_validation + and self.cypher_validator + and query_type not in skip_validation_types + ): print("\n🔍 LLM3 正在验证参数...") # 构建一个模拟的 Cypher 语句用于验证 # 注意：这里我们验证的是参数的完整性和正确性，而不是实际的 Cypher 语句 validation_result = self._validate_params(question, query_type, params) - + if not validation_result["is_valid"]: print(f" ⚠️ 发现问题:") for issue in validation_result["issues"]: print(f" - {issue}") - + if validation_result["suggestions"]: print(f" 💡 修改建议:") for suggestion in validation_result["suggestions"]: print(f" - {suggestion}") - + # 如果验证器提供了修正后的 Cypher 语句，使用它 if validation_result.get("corrected_cypher"): print(f"\n 🔧 使用修正后的 Cypher 语句:") @@ -950,10 +1088,10 @@ def ask(self, question: str) -> Dict: print(f" ✅ 参数验证通过") elif query_type in skip_validation_types: print(f"\n⏭️ 跳过 LLM3 验证（查询类型: {query_type}）") - + # Step 4: 执行查询 print("\n⚙️ 执行查询...") - + # 如果有修正后的 Cypher 语句，直接执行它 if corrected_cypher: print(f" 使用修正后的 Cypher 语句执行...") @@ -966,7 +1104,7 @@ def ask(self, question: str) -> Dict: "params": params, "results": results, "count": len(results), - "corrected": True + "corrected": True, } except Exception as e: result = { @@ -974,115 +1112,121 @@ def ask(self, question: str) -> Dict: "error": str(e), "query_type": query_type, "params": params, - "corrected_cypher": corrected_cypher + "corrected_cypher": corrected_cypher, } else: # 使用原始的模板方法执行 result = self.executor.execute(query_type, params) - + if result["success"]: print(f"✅ 查询成功！返回 {result['count']} 条结果\n") - + # 显示结果预览 for i, item in enumerate(result["results"][:3], 1): print(f"{i}. {item}") - + if result["count"] > 3: print(f"... 还有 {result['count'] - 3} 条结果") else: print(f"❌ 查询失败: {result['error']}") - + return result - + def _validate_params(self, question: str, query_type: str, params: Dict) -> Dict: """ 验证提取的参数是否完整和正确 - + Args: question: 用户原始问题 query_type: 查询类型 params: 提取的参数 - + Returns: 验证结果字典 """ # 构建一个描述性的"Cypher"用于验证 cypher_description = self._build_cypher_description(query_type, params) - + # 调用验证器 - return self.cypher_validator.validate(cypher_description, question, query_type, params) - + return self.cypher_validator.validate( + cypher_description, question, query_type, params + ) + def _build_cypher_description(self, query_type: str, params: Dict) -> str: """ 根据查询类型和参数构建 Cypher 描述 - + 这不是实际的 Cypher 语句，而是一个描述性的文本， 用于让 LLM 理解我们要执行的查询 """ description = f"查询类型: {query_type}\n\n" description += "查询参数:\n" - + for key, value in params.items(): description += f" - {key}: {value}\n" - + # 根据查询类型添加特定的描述 if query_type == "common_neighbor": description += "\n预期行为:\n" - description += f" 1. 找出节点 {params.get('v1')} 和 {params.get('v2')} 的公共邻居\n" - + description += ( + f" 1. 找出节点 {params.get('v1')} 和 {params.get('v2')} 的公共邻居\n" + ) + if "rel_conditions" in params and params["rel_conditions"]: - description += f" 2. 对两条关系都应用过滤条件: {params['rel_conditions']}\n" + description += ( + f" 2. 对两条关系都应用过滤条件: {params['rel_conditions']}\n" + ) description += f" 3. 只返回满足条件的公共邻居\n" else: description += f" 2. 返回所有公共邻居（无过滤条件）\n" - + elif query_type == "relationship_filter": description += "\n预期行为:\n" description += f" 1. 筛选关系类型: {params.get('rel_type')}\n" - + if "rel_conditions" in params and params["rel_conditions"]: description += f" 2. 应用关系属性过滤: {params['rel_conditions']}\n" - + if "return_fields" in params: description += f" 3. 返回字段: {params['return_fields']}\n" - + elif query_type == "aggregation_query": description += "\n预期行为:\n" description += f" 1. 聚合类型: {params.get('aggregate_type')}\n" description += f" 2. 分组依据: {params.get('group_by_node')}\n" - + if "aggregate_field" in params: description += f" 3. 聚合字段: {params['aggregate_field']}\n" - + elif query_type == "subgraph": description += "\n预期行为:\n" description += f" 1. 抽取以节点 {params.get('value')} 为中心的子图\n" description += f" 2. 跳数: {params.get('hops', 2)}\n" description += f" 3. 节点标签: {params.get('label')}\n" description += f" 4. 节点属性键: {params.get('key')}\n" - + if "rel_type" in params: description += f" 5. 关系类型: {params.get('rel_type')}\n" - + if "direction" in params: description += f" 6. 关系方向: {params.get('direction')}\n" - + elif query_type == "neighbor_query": description += "\n预期行为:\n" description += f" 1. 查询节点 {params.get('value')} 的邻居\n" description += f" 2. 跳数: {params.get('hops', 1)}\n" description += f" 3. 节点标签: {params.get('label')}\n" description += f" 4. 节点属性键: {params.get('key')}\n" - + if "rel_type" in params: description += f" 5. 关系类型: {params.get('rel_type')}\n" - + if "direction" in params: description += f" 6. 关系方向: {params.get('direction')}\n" - + if "order_by" in params: description += f" 7. 排序字段: {params.get('order_by')}\n" - + return description @@ -1095,65 +1239,71 @@ def main(): print("=" * 80) print("🚀 Neo4j 自然语言查询引擎（双 LLM 架构）") print("=" * 80) - + # 数据库配置 - uri = input("\nNeo4j URI (默认 bolt://localhost:7687): ").strip() or "bolt://localhost:7687" + uri = ( + input("\nNeo4j URI (默认 bolt://localhost:7687): ").strip() + or "bolt://localhost:7687" + ) user = input("用户名 (默认 neo4j): ").strip() or "neo4j" password = input("密码: ").strip() - + try: # 初始化数据库客户端 config = Neo4jConfig(uri=uri, user=user, password=password) db_client = Neo4jGraphClient(config) - + # add gjq # 初始化 Reasoner（使用配置文件或默认配置） # 这里需要根据实际情况配置 ReasonerConfig # 示例：使用 OpenAI from aag.config.engine_config import LLMConfig - + llm_config = LLMConfig( - provider='openai', + provider="openai", ollama={}, openai={ - 'base_url': os.environ.get('OPENAI_BASE_URL', 'https://api.openai.com/v1'), - 'api_key': os.environ.get('OPENAI_API_KEY'), - 'model': 'gpt-4o-mini' - } + "base_url": os.environ.get( + "OPENAI_BASE_URL", "https://api.openai.com/v1" + ), + "api_key": os.environ.get("OPENAI_API_KEY"), + "model": "gpt-4o-mini", + }, ) reasoner_config = ReasonerConfig(llm=llm_config) reasoner = Reasoner(reasoner_config) - + # 创建查询引擎 engine = NaturalLanguageQueryEngine(db_client, reasoner) engine.initialize() - + print("\n✨ 已连接！输入 'exit' 退出，'help' 查看示例") print("=" * 80) - + while True: question = input("\n❓ 请输入问题: ").strip() - + if not question: continue - - if question.lower() in ['exit', 'quit', '退出']: + + if question.lower() in ["exit", "quit", "退出"]: print("\n👋 再见！") break - - if question.lower() == 'help': + + if question.lower() == "help": print("\n📚 示例问题：") for qtype, info in QUERY_TEMPLATES.items(): print(f" - {info['example']}") continue - + engine.ask(question) - + db_client.close() - + except Exception as e: print(f"\n❌ 错误: {e}") import traceback + traceback.print_exc() diff --git a/aag/computing_engine/networkx_server/graph_computation_processor.py b/aag/computing_engine/networkx_server/graph_computation_processor.py index a8bf39e..363044a 100644 --- a/aag/computing_engine/networkx_server/graph_computation_processor.py +++ b/aag/computing_engine/networkx_server/graph_computation_processor.py @@ -1,6 +1,6 @@ import networkx as nx import logging -import community +import community from typing import List, Tuple, Dict, Any, Optional, Union from aag.expert_search_engine.database.datatype import * from aag.expert_search_engine.database.nebulagraph import NebulaGraphClient @@ -13,7 +13,7 @@ class GraphComputationProcessor(GraphProcessor): """ 图处理器类：处理从图数据库提取的图数据，转换为networkx格式并运行图算法 """ - + def __init__(self): self.vertices = None self.edges = None @@ -21,17 +21,28 @@ def __init__(self): self.is_directed = False self.is_multiedge = False self.query_vertices = None - - def create_graph_from_edges(self, vertices: List[VertexData], edges: List[EdgeData], directed: bool = True, multiedge: bool = False): + # Louvain 社区检测结果缓存：避免同一图上重复计算 + # 键为 (节点数, 边数, resolution, random_state) 元组，值为结果字典 + self._louvain_cache: Dict[ + Tuple[int, int, float, Optional[int]], Dict[str, Any] + ] = {} + + def create_graph_from_edges( + self, + vertices: List[VertexData], + edges: List[EdgeData], + directed: bool = True, + multiedge: bool = False, + ): """ 从边列表创建networkx图 - + Args: vertices: 顶点列表 edges: 边列表 directed: 是否为有向图 multiedge: 是否为多重边图 - + Returns: networkx图对象 """ @@ -43,7 +54,7 @@ def create_graph_from_edges(self, vertices: List[VertexData], edges: List[EdgeDa self.graph = nx.MultiDiGraph() if self.is_multiedge else nx.DiGraph() else: self.graph = nx.MultiGraph() if self.is_multiedge else nx.Graph() - + # 添加顶点及属性 for v in vertices: self.graph.add_node(v.vid, **v.properties) @@ -52,7 +63,7 @@ def create_graph_from_edges(self, vertices: List[VertexData], edges: List[EdgeDa for idx, e in enumerate(edges): edge_attrs = e.properties.copy() if e.rank is not None: - edge_attrs['rank'] = e.rank + edge_attrs["rank"] = e.rank if self.is_multiedge: key = e.rank if e.rank is not None else 0 @@ -60,52 +71,53 @@ def create_graph_from_edges(self, vertices: List[VertexData], edges: List[EdgeDa else: self.graph.add_edge(e.src, e.dst, **edge_attrs) - logger.debug(f"✅ 图构建完成: {self.graph.number_of_nodes()} 个节点，{self.graph.number_of_edges()} 条边") + logger.debug( + f"✅ 图构建完成: {self.graph.number_of_nodes()} 个节点，{self.graph.number_of_edges()} 条边" + ) return self.graph - + except Exception as e: logger.error(f"创建图时发生错误: {e}") raise - - def run_pagerank(self, alpha: float = 0.85, max_iter: int = 100, tol: float = 1e-6) -> Dict[Any, float]: + + def run_pagerank( + self, alpha: float = 0.85, max_iter: int = 100, tol: float = 1e-6 + ) -> Dict[Any, float]: """ 运行PageRank算法 - + Args: alpha: 阻尼系数，默认为0.85 max_iter: 最大迭代次数，默认为100 tol: 收敛容差，默认为1e-6 - + Returns: PageRank分数字典，键为节点，值为分数 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: pagerank_scores = nx.pagerank( - self.graph, - alpha=alpha, - max_iter=max_iter, - tol=tol + self.graph, alpha=alpha, max_iter=max_iter, tol=tol ) logger.info(f"PageRank算法完成，计算了 {len(pagerank_scores)} 个节点的分数") return pagerank_scores - + except Exception as e: logger.error(f"运行PageRank算法时发生错误: {e}") raise - + def run_connected_components(self) -> List[set]: """ 运行连通分量算法 - + Returns: 连通分量列表，每个连通分量是一个节点集合 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: if self.is_directed: # 对于有向图，使用强连通分量 @@ -115,170 +127,195 @@ def run_connected_components(self) -> List[set]: # 对于无向图，使用连通分量 components = list(nx.connected_components(self.graph)) logger.info(f"找到 {len(components)} 个连通分量") - + return components - + except Exception as e: logger.error(f"运行连通分量算法时发生错误: {e}") raise - + def run_shortest_path(self, source: Any, target: Any) -> Optional[List[Any]]: """ 计算两个节点之间的最短路径 - + Args: source: 源节点 target: 目标节点 - + Returns: 最短路径节点列表，如果不存在路径则返回None """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: - if self.is_directed: - path = nx.shortest_path(self.graph, source, target) - else: - path = nx.shortest_path(self.graph, source, target) - + # TODO: networkx.shortest_path 对 Directed/Undirected 图自动适配； + # 若未来需要区分算法（如 Dijkstra/Bellman-Ford），在此补充有向图专属逻辑。 + path = nx.shortest_path(self.graph, source, target) + logger.info(f"从 {source} 到 {target} 的最短路径长度为 {len(path) - 1}") return path - + except nx.NetworkXNoPath: logger.warning(f"从 {source} 到 {target} 不存在路径") return None except Exception as e: logger.error(f"计算最短路径时发生错误: {e}") raise - + def run_betweenness_centrality(self) -> Dict[Any, float]: """ 计算介数中心性 - + Returns: 介数中心性分数字典 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: betweenness_scores = nx.betweenness_centrality(self.graph) - logger.info(f"介数中心性计算完成，计算了 {len(betweenness_scores)} 个节点的分数") + logger.info( + f"介数中心性计算完成，计算了 {len(betweenness_scores)} 个节点的分数" + ) return betweenness_scores - + except Exception as e: logger.error(f"计算介数中心性时发生错误: {e}") raise - + def run_closeness_centrality(self) -> Dict[Any, float]: """ 计算接近中心性 - + Returns: 接近中心性分数字典 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: closeness_scores = nx.closeness_centrality(self.graph) - logger.info(f"接近中心性计算完成，计算了 {len(closeness_scores)} 个节点的分数") + logger.info( + f"接近中心性计算完成，计算了 {len(closeness_scores)} 个节点的分数" + ) return closeness_scores - + except Exception as e: logger.error(f"计算接近中心性时发生错误: {e}") raise - + def run_degree_centrality(self) -> Dict[Any, float]: """ 计算度中心性 - + Returns: 度中心性分数字典 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + try: degree_scores = nx.degree_centrality(self.graph) logger.info(f"度中心性计算完成，计算了 {len(degree_scores)} 个节点的分数") return degree_scores - + except Exception as e: logger.error(f"计算度中心性时发生错误: {e}") raise - def run_louvain_community_detection(self, resolution: float = 1.0, random_state: Optional[int] = None) -> Dict[str, Any]: + def run_louvain_community_detection( + self, resolution: float = 1.0, random_state: Optional[int] = None + ) -> Dict[str, Any]: """ 使用Louvain算法进行社区检测 - + Args: resolution: 分辨率参数，控制社区大小。值越大，社区越小 random_state: 随机种子，用于结果的可重现性 - + Returns: 包含社区检测结果的字典： - 'communities': 社区列表，每个社区是一个节点集合 - 'modularity': 模块度分数 - 'node_communities': 节点到社区的映射字典 - - 'community_count': 社区数量 + - 'community_count': 社区数量 """ if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + + # 缓存检查：相同图（节点数+边数）且参数相同时复用上次结果 + cache_key = ( + self.graph.number_of_nodes(), + self.graph.number_of_edges(), + resolution, + random_state, + ) + if cache_key in self._louvain_cache: + logger.debug( + f"Louvain社区检测命中缓存 (nodes={cache_key[0]}, edges={cache_key[1]})" + ) + return self._louvain_cache[cache_key] + try: # 将图转换为无向图（Louvain算法通常用于无向图） if self.is_directed: undirected_graph = self.graph.to_undirected() else: undirected_graph = self.graph - + # 运行Louvain算法 - partition = community.best_partition(undirected_graph, resolution=resolution, random_state=random_state) - + partition = community.best_partition( + undirected_graph, resolution=resolution, random_state=random_state + ) + # 计算模块度 modularity = self._calculate_modularity(undirected_graph, partition) - + # 构建社区列表 communities = {} for node, community_id in partition.items(): if community_id not in communities: communities[community_id] = set() communities[community_id].add(node) - + # 转换为列表格式 community_list = list(communities.values()) - + # 构建节点到社区的映射 node_communities = {node: comm_id for node, comm_id in partition.items()} - + result = { - 'communities': community_list, - 'modularity': modularity, - 'node_communities': node_communities, - 'community_count': len(community_list), - 'resolution': resolution + "communities": community_list, + "modularity": modularity, + "node_communities": node_communities, + "community_count": len(community_list), + "resolution": resolution, } - - logger.info(f"Louvain社区检测完成，发现 {len(community_list)} 个社区，模块度: {modularity:.4f}") - print(f"Louvain社区检测完成，发现 {len(community_list)} 个社区，模块度: {modularity:.4f}") + + logger.info( + f"Louvain社区检测完成，发现 {len(community_list)} 个社区，模块度: {modularity:.4f}" + ) + # 存入缓存，后续相同图+参数可直接复用 + self._louvain_cache[cache_key] = result return result - + except ImportError: logger.error("未安装community模块，请运行: pip install python-louvain") raise except Exception as e: logger.error(f"Louvain社区检测时发生错误: {e}") raise - - def get_community_by_specific_id(self, query_vertex_id: Optional[int] = None) -> Dict[str, Any]: + + def get_community_by_specific_id( + self, query_vertex_id: Optional[int] = None + ) -> Dict[str, Any]: """ 获取指定节点ID的社区结果 - + Args: query_vertex_id: 查询的节点ID，如果为None则使用self.query_vertices - + Returns: 包含指定节点社区信息的字典： - 'vertex_id': 查询的节点ID @@ -290,73 +327,81 @@ def get_community_by_specific_id(self, query_vertex_id: Optional[int] = None) -> """ if query_vertex_id is None: query_vertex_id = self.query_vertices - + if query_vertex_id is None: - raise ValueError("未指定查询节点ID，请设置query_vertex_id参数或self.query_vertices") - + raise ValueError( + "未指定查询节点ID，请设置query_vertex_id参数或self.query_vertices" + ) + if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + # 运行Louvain算法获取社区检测结果 louvain_result = self.run_louvain_community_detection() - + # 获取节点到社区的映射 - node_communities = louvain_result['node_communities'] - + node_communities = louvain_result["node_communities"] + # 检查查询节点是否存在于图中 if query_vertex_id not in node_communities: return { - 'error': f'节点 {query_vertex_id} 不存在于图中', - 'vertex_id': query_vertex_id, - 'available_nodes': list(node_communities.keys()) + "error": f"节点 {query_vertex_id} 不存在于图中", + "vertex_id": query_vertex_id, + "available_nodes": list(node_communities.keys()), } - + # 获取节点所属的社区ID community_id = node_communities[query_vertex_id] - + # 找到该社区的所有成员 community_members = set() for node, comm_id in node_communities.items(): if comm_id == community_id: community_members.add(node) - + # 获取节点的邻居 if query_vertex_id in self.graph: undirected_graph = self.graph.to_undirected() all_neighbors = set(undirected_graph.neighbors(query_vertex_id)) - + # 分类邻居：社区内和社区外 neighbors_in_community = all_neighbors.intersection(community_members) neighbors_outside_community = all_neighbors - community_members else: neighbors_in_community = set() neighbors_outside_community = set() - + result = { - 'vertex_id': query_vertex_id, - 'community_id': community_id, - 'community_members': sorted(list(community_members)), - 'community_size': len(community_members), - 'neighbors_in_community': sorted(list(neighbors_in_community)), - 'neighbors_outside_community': sorted(list(neighbors_outside_community)), - 'total_neighbors': len(neighbors_in_community) + len(neighbors_outside_community), - 'community_cohesion': len(neighbors_in_community) / max(1, len(all_neighbors)) if 'all_neighbors' in locals() else 0.0 + "vertex_id": query_vertex_id, + "community_id": community_id, + "community_members": sorted(list(community_members)), + "community_size": len(community_members), + "neighbors_in_community": sorted(list(neighbors_in_community)), + "neighbors_outside_community": sorted(list(neighbors_outside_community)), + "total_neighbors": len(neighbors_in_community) + + len(neighbors_outside_community), + "community_cohesion": len(neighbors_in_community) + / max(1, len(all_neighbors)) + if "all_neighbors" in locals() + else 0.0, } - - logger.info(f"节点 {query_vertex_id} 的社区信息: 社区ID={community_id}, 社区大小={len(community_members)}") - print(f"节点 {query_vertex_id} 的社区信息: 社区ID={community_id}, 社区大小={len(community_members)}") - + + logger.info( + f"节点 {query_vertex_id} 的社区信息: 社区ID={community_id}, 社区大小={len(community_members)}" + ) + return result - - def _calculate_modularity(self, graph: nx.Graph, partition: Dict[Any, int]) -> float: + def _calculate_modularity( + self, graph: nx.Graph, partition: Dict[Any, int] + ) -> float: """ 计算模块度 - + Args: graph: 无向图 partition: 节点到社区的映射 - + Returns: 模块度分数 """ @@ -365,222 +410,235 @@ def _calculate_modularity(self, graph: nx.Graph, partition: Dict[Any, int]) -> f except ImportError: # 如果community模块不可用，使用自定义实现 return self._custom_modularity(graph, partition) - + def _custom_modularity(self, graph: nx.Graph, partition: Dict[Any, int]) -> float: """ 自定义模块度计算实现 - + Args: graph: 无向图 partition: 节点到社区的映射 - + Returns: 模块度分数 """ m = graph.number_of_edges() if m == 0: return 0.0 - + # 计算总度数 total_degree = sum(dict(graph.degree()).values()) - + # 按社区分组 communities = {} for node, comm_id in partition.items(): if comm_id not in communities: communities[comm_id] = set() communities[comm_id].add(node) - + modularity = 0.0 - + for community in communities.values(): # 计算社区内边数 internal_edges = 0 community_degree = 0 - + for node in community: community_degree += graph.degree(node) for neighbor in graph.neighbors(node): if neighbor in community: internal_edges += 1 - + # 避免重复计算边 internal_edges //= 2 - + # 计算模块度贡献 - expected_edges = (community_degree ** 2) / (2 * m) + expected_edges = (community_degree**2) / (2 * m) modularity += (internal_edges - expected_edges) / m - + return modularity - - def get_community_statistics(self, louvain_result: Dict[str, Any]) -> Dict[str, Any]: + + def get_community_statistics( + self, louvain_result: Dict[str, Any] + ) -> Dict[str, Any]: """ 获取社区检测的统计信息 - + Args: louvain_result: Louvain算法的结果 - + Returns: 统计信息字典 """ - communities = louvain_result['communities'] - + communities = louvain_result["communities"] + if not communities: return {"error": "没有找到社区"} - + # 计算社区大小 community_sizes = [len(comm) for comm in communities] - + # 计算统计信息 stats = { - 'total_communities': len(communities), - 'largest_community_size': max(community_sizes), - 'smallest_community_size': min(community_sizes), - 'average_community_size': sum(community_sizes) / len(community_sizes), - 'modularity': louvain_result['modularity'], - 'size_distribution': { - 'small': len([s for s in community_sizes if s <= 5]), - 'medium': len([s for s in community_sizes if 5 < s <= 20]), - 'large': len([s for s in community_sizes if s > 20]) - } + "total_communities": len(communities), + "largest_community_size": max(community_sizes), + "smallest_community_size": min(community_sizes), + "average_community_size": sum(community_sizes) / len(community_sizes), + "modularity": louvain_result["modularity"], + "size_distribution": { + "small": len([s for s in community_sizes if s <= 5]), + "medium": len([s for s in community_sizes if 5 < s <= 20]), + "large": len([s for s in community_sizes if s > 20]), + }, } - + return stats - + def run_algorithm(self, algorithm: str, **kwargs) -> Union[Dict, List, Any]: """ 运行指定的图算法 - + Args: algorithm: 算法名称，支持 'pagerank', 'cc', 'shortest_path', 'betweenness', 'closeness', 'degree' **kwargs: 算法特定参数 - + Returns: 算法结果 """ algorithm_map = { - 'pagerank': self.run_pagerank, - 'cc': self.run_connected_components, - 'shortest_path': self.run_shortest_path, - 'betweenness': self.run_betweenness_centrality, - 'closeness': self.run_closeness_centrality, - 'degree': self.run_degree_centrality, - 'louvain': self.run_louvain_community_detection + "pagerank": self.run_pagerank, + "cc": self.run_connected_components, + "shortest_path": self.run_shortest_path, + "betweenness": self.run_betweenness_centrality, + "closeness": self.run_closeness_centrality, + "degree": self.run_degree_centrality, + "louvain": self.run_louvain_community_detection, } - + if algorithm not in algorithm_map: - raise ValueError(f"不支持的算法: {algorithm}。支持的算法: {list(algorithm_map.keys())}") - + raise ValueError( + f"不支持的算法: {algorithm}。支持的算法: {list(algorithm_map.keys())}" + ) + try: return algorithm_map[algorithm](**kwargs) except Exception as e: logger.error(f"运行算法 {algorithm} 时发生错误: {e}") - raise - + raise + def get_graph_info(self) -> Dict[str, Any]: """ 获取图的基本信息 - + Returns: 包含图信息的字典 """ if self.graph is None: return {"error": "图未初始化"} - + info = { "节点数": self.graph.number_of_nodes(), "边数": self.graph.number_of_edges(), "图类型": "有向图" if self.is_directed else "无向图", "是否连通": nx.is_connected(self.graph) if not self.is_directed else None, - "是否强连通": nx.is_strongly_connected(self.graph) if self.is_directed else None + "是否强连通": nx.is_strongly_connected(self.graph) + if self.is_directed + else None, } - + return info - - def execute_plan(self, edges: List[Tuple], plan: List[Dict[str, Any]]) -> Dict[str, Any]: + + def execute_plan( + self, edges: List[Tuple], plan: List[Dict[str, Any]] + ) -> Dict[str, Any]: """ 根据执行计划按顺序执行算法 - + Args: plan: 执行计划列表，每个元素是一个字典，包含： - 'algorithm': 算法名称 - 'params': 算法参数字典（可选） - 'name': 结果名称（可选，默认为算法名称） - + Returns: 包含所有算法结果的字典 """ self.create_graph_from_edges(edges) if self.graph is None: raise ValueError("图未初始化，请先调用create_graph_from_edges方法") - + results = {} - + try: for i, step in enumerate(plan): - algorithm = step.get('algorithm') - params = step.get('params', {}) - result_name = step.get('name', algorithm) - + algorithm = step.get("algorithm") + params = step.get("params", {}) + result_name = step.get("name", algorithm) + if not algorithm: - logger.error(f"步骤 {i+1} 缺少算法名称") - raise ValueError(f"步骤 {i+1} 缺少算法名称") - - logger.info(f"执行步骤 {i+1}: {algorithm}") - + logger.error(f"步骤 {i + 1} 缺少算法名称") + raise ValueError(f"步骤 {i + 1} 缺少算法名称") + + logger.info(f"执行步骤 {i + 1}: {algorithm}") + # 执行算法 result = self.run_algorithm(algorithm, **params) results[result_name] = result - - logger.info(f"步骤 {i+1} 完成: {algorithm}") - + + logger.info(f"步骤 {i + 1} 完成: {algorithm}") + logger.info(f"执行计划完成，共执行 {len(results)} 个算法") return results - + except Exception as e: logger.error(f"执行计划时发生错误: {e}") raise - - def create_execution_plan(self, algorithms: List[str], - algorithm_params: Optional[Dict[str, Dict]] = None, - result_names: Optional[List[str]] = None) -> List[Dict[str, Any]]: + + def create_execution_plan( + self, + algorithms: List[str], + algorithm_params: Optional[Dict[str, Dict]] = None, + result_names: Optional[List[str]] = None, + ) -> List[Dict[str, Any]]: """ 创建执行计划 - + Args: algorithms: 算法名称列表 algorithm_params: 算法参数字典，键为算法名称，值为参数字典 result_names: 结果名称列表，如果为None则使用算法名称 - + Returns: 执行计划列表 """ plan = [] algorithm_params = algorithm_params or {} - + for i, algorithm in enumerate(algorithms): step = { - 'algorithm': algorithm, - 'params': algorithm_params.get(algorithm, {}), - 'name': result_names[i] if result_names and i < len(result_names) else algorithm + "algorithm": algorithm, + "params": algorithm_params.get(algorithm, {}), + "name": result_names[i] + if result_names and i < len(result_names) + else algorithm, } plan.append(step) - - return plan - + + return plan + def get_vertex_map(self, vertices: List[VertexData]): """ 将 self.vertices 转换为以 vid 为键、VertexData 为值的字典。 """ - self.vertices = {vertex.vid: vertex for vertex in vertices} + self.vertices = {vertex.vid: vertex for vertex in vertices} -if __name__ == '__main__': +if __name__ == "__main__": testclient = NebulaGraphClient(space_name="AMLSim1K") vertices = testclient.get_all_vertices() edges = testclient.get_all_edges() print(f"vertices number: {len(vertices)}, edge number: {len(edges)}") - + graphprocessor = GraphComputationProcessor() graphprocessor.create_graph_from_edges(vertices, edges, True, True) info = graphprocessor.get_graph_info() @@ -603,11 +661,92 @@ def get_vertex_map(self, vertices: List[VertexData]): # for k,v in result.items(): # print(f"{k}: {v}") - com_member = [6, 15, 23, 28, 68, 86, 93, 118, 124, 136, 139, 155, 171, 221, 243, 244, 310, 360, 371, 385, 386, 402, 409, 437, 462, 467, 508, 545, 570, 585, 622, 623, 630, 631, 635, 654, 675, 679, 744, 746, 769, 797, 799, 802, 815, 839, 847, 854, 878, 884, 895, 934, 956, 959, 965, 982, 983, 1040, 1045, 1048, 1052, 1079, 1104, 1130, 1138, 1164, 1189, 1202, 1264, 1273, 1280, 1282, 1293, 1297, 1307, 1329, 1351, 1379, 1407, 1418] + com_member = [ + 6, + 15, + 23, + 28, + 68, + 86, + 93, + 118, + 124, + 136, + 139, + 155, + 171, + 221, + 243, + 244, + 310, + 360, + 371, + 385, + 386, + 402, + 409, + 437, + 462, + 467, + 508, + 545, + 570, + 585, + 622, + 623, + 630, + 631, + 635, + 654, + 675, + 679, + 744, + 746, + 769, + 797, + 799, + 802, + 815, + 839, + 847, + 854, + 878, + 884, + 895, + 934, + 956, + 959, + 965, + 982, + 983, + 1040, + 1045, + 1048, + 1052, + 1079, + 1104, + 1130, + 1138, + 1164, + 1189, + 1202, + 1264, + 1273, + 1280, + 1282, + 1293, + 1297, + 1307, + 1329, + 1351, + 1379, + 1407, + 1418, + ] sar_count = 0 print(len(com_member)) - for v in vertices: + for v in vertices: if v.vid in com_member: - if v.properties['prior_sar_count']: - sar_count+=1 - print(sar_count) \ No newline at end of file + if v.properties["prior_sar_count"]: + sar_count += 1 + print(sar_count) diff --git a/aag/data_pipeline/data_transformer/dataset_manager.py b/aag/data_pipeline/data_transformer/dataset_manager.py index d32b937..3f7cbd6 100644 --- a/aag/data_pipeline/data_transformer/dataset_manager.py +++ b/aag/data_pipeline/data_transformer/dataset_manager.py @@ -2,14 +2,20 @@ from typing import Any, Tuple, Union from pathlib import Path import yaml + # add gjq: 导入日志和dataclass工具 import logging from dataclasses import asdict from aag.data_pipeline.data_transformer.graph_loader import GraphDataLoader from aag.data_pipeline.data_transformer.table_loader import TableDataLoader from aag.data_pipeline.data_transformer.text_loader import TextDataLoader -from aag.config.data_upload_config import DataUploadConfig, DatasetConfig, load_data_upload_config -from aag.utils.path_utils import DATASETS_INDEX_PATH +from aag.config.data_upload_config import ( + DataUploadConfig, + DatasetConfig, + load_data_upload_config, +) +from aag.utils.path_utils import DATASETS_INDEX_PATH + # add gjq: 导入Neo4j图数据加载器 from aag.computing_engine.graph_query.load_data_into_neo4j import Neo4jGraphLoader @@ -24,17 +30,17 @@ class DatasetManager: def __init__(self, datasets_index_path: Optional[Path] = None): """ Initialize dataset manager - + Args: - datasets_index_path: Path to datasets.yaml index file. + datasets_index_path: Path to datasets.yaml index file. If None, uses default path from path_utils """ if datasets_index_path is None: datasets_index_path = DATASETS_INDEX_PATH - + self.datasets_index_path = Path(datasets_index_path) self.datasets_schema_path = self.datasets_index_path.parent - + self.graph_loader = GraphDataLoader() self.table_loader = TableDataLoader() self.text_loader = TextDataLoader() @@ -44,14 +50,13 @@ def __init__(self, datasets_index_path: Optional[Path] = None): "table": self.table_loader, "text": self.text_loader, } - + # Cache dataset index for quick lookup self.datasets_index: Dict[str, Dict] = {} - + # Load all datasets from the new architecture self._load_all_datasets() - def _load_all_datasets(self): """ Load all datasets from the new architecture: @@ -59,49 +64,49 @@ def _load_all_datasets(self): 2. For each dataset, load corresponding schema files from dataset folder """ if not self.datasets_index_path.exists(): - print(f"[WARNING] Datasets index file not found: {self.datasets_index_path}") + logger.warning(f"Datasets index file not found: {self.datasets_index_path}") return - + try: - with open(self.datasets_index_path, 'r', encoding='utf-8') as f: + with open(self.datasets_index_path, "r", encoding="utf-8") as f: index_data = yaml.safe_load(f) or {} - - datasets_list = index_data.get('datasets', []) - + + datasets_list = index_data.get("datasets", []) + for dataset_info in datasets_list: - dataset_name = dataset_info.get('name') - dataset_type = dataset_info.get('type', '').lower() - data_path = dataset_info.get('data_path', dataset_name) - + dataset_name = dataset_info.get("name") + dataset_type = dataset_info.get("type", "").lower() + data_path = dataset_info.get("data_path", dataset_name) + if not dataset_name or dataset_type not in self.VALID_DATA_TYPES: continue self.datasets_index[dataset_name] = dataset_info dataset_folder = self.datasets_schema_path / data_path - + # Load schema based on dataset type if dataset_type == "graph": schema_file = dataset_folder / "graph_schemas.yaml" if schema_file.exists(): self._load_graph_schemas(schema_file, dataset_name) - + elif dataset_type == "text": schema_file = dataset_folder / "text_schemas.yaml" if schema_file.exists(): self._load_text_schemas(schema_file, dataset_name) - + elif dataset_type == "table": schema_file = dataset_folder / "table_schemas.yaml" if schema_file.exists(): self._load_table_schemas(schema_file, dataset_name) - + except Exception as e: raise RuntimeError(f"[ERROR] Failed to load datasets from index: {e}") - + def _load_graph_schemas(self, schema_file: Path, dataset_name: str): """ Load graph schemas from a dataset's graph_schemas.yaml - + For graph datasets, graph_schemas.yaml should contain only one entry. Store as dataset_name -> List[DatasetConfig] (list with single config) """ @@ -110,11 +115,13 @@ def _load_graph_schemas(self, schema_file: Path, dataset_name: str): if not schemas.datasets: print(f"[WARNING] No datasets found in {schema_file}") return - + # Graph schemas should have only one entry if len(schemas.datasets) > 1: - print(f"[WARNING] Graph schema file {schema_file} contains {len(schemas.datasets)} entries, expected 1. Using first entry.") - + print( + f"[WARNING] Graph schema file {schema_file} contains {len(schemas.datasets)} entries, expected 1. Using first entry." + ) + ds_config = schemas.datasets[0] # Update path to be absolute if it's relative self._resolve_paths(ds_config, schema_file.parent) @@ -122,11 +129,11 @@ def _load_graph_schemas(self, schema_file: Path, dataset_name: str): self.graph_loader.dataset_schemas[dataset_name] = [ds_config] except Exception as e: print(f"[WARNING] Failed to load graph schemas from {schema_file}: {e}") - + def _load_text_schemas(self, schema_file: Path, dataset_name: str): """ Load text schemas from a dataset's text_schemas.yaml - + For text datasets, text_schemas.yaml may contain multiple file entries. Store as dataset_name -> List[DatasetConfig] (list of file configs) """ @@ -135,23 +142,23 @@ def _load_text_schemas(self, schema_file: Path, dataset_name: str): if not schemas.datasets: print(f"[WARNING] No datasets found in {schema_file}") return - + # Collect all file-level configs for this dataset file_configs = [] for ds_config in schemas.datasets: # Update path to be absolute if it's relative self._resolve_paths(ds_config, schema_file.parent) file_configs.append(ds_config) - + # Store with dataset-level name as key, value is list of file configs self.text_loader.dataset_schemas[dataset_name] = file_configs except Exception as e: print(f"[WARNING] Failed to load text schemas from {schema_file}: {e}") - + def _load_table_schemas(self, schema_file: Path, dataset_name: str): """ Load table schemas from a dataset's table_schemas.yaml - + Store as dataset_name -> List[DatasetConfig] (list with single config) """ try: @@ -159,52 +166,52 @@ def _load_table_schemas(self, schema_file: Path, dataset_name: str): if not schemas.datasets: print(f"[WARNING] No datasets found in {schema_file}") return - + # Collect all table configs for this dataset table_configs = [] for ds_config in schemas.datasets: # Update path to be absolute if it's relative self._resolve_paths(ds_config, schema_file.parent) table_configs.append(ds_config) - + # Store with dataset-level name as key, value is list of configs self.table_loader.dataset_schemas[dataset_name] = table_configs except Exception as e: print(f"[WARNING] Failed to load table schemas from {schema_file}: {e}") - + def _resolve_paths(self, dataset_config: DatasetConfig, schema_dir: Path): """ Resolve relative paths in dataset config to absolute paths - + Args: dataset_config: Dataset configuration object schema_dir: Directory where schema file is located """ if dataset_config.type == "text": # For text, schema.path is relative to schema file - if hasattr(dataset_config.schema, 'path'): + if hasattr(dataset_config.schema, "path"): rel_path = dataset_config.schema.path if not Path(rel_path).is_absolute(): abs_path = (schema_dir / rel_path).resolve() dataset_config.schema.path = str(abs_path) - + elif dataset_config.type == "graph": # For graph, vertex and edge paths are relative to schema file - if hasattr(dataset_config.schema, 'vertex'): + if hasattr(dataset_config.schema, "vertex"): for v in dataset_config.schema.vertex: - if hasattr(v, 'path') and not Path(v.path).is_absolute(): + if hasattr(v, "path") and not Path(v.path).is_absolute(): abs_path = (schema_dir / v.path).resolve() v.path = str(abs_path) - - if hasattr(dataset_config.schema, 'edge'): + + if hasattr(dataset_config.schema, "edge"): for e in dataset_config.schema.edge: - if hasattr(e, 'path') and not Path(e.path).is_absolute(): + if hasattr(e, "path") and not Path(e.path).is_absolute(): abs_path = (schema_dir / e.path).resolve() e.path = str(abs_path) - + elif dataset_config.type == "table": # For table, schema.path is relative to schema file - if hasattr(dataset_config.schema, 'path'): + if hasattr(dataset_config.schema, "path"): rel_path = dataset_config.schema.path if not Path(rel_path).is_absolute(): abs_path = (schema_dir / rel_path).resolve() @@ -213,10 +220,10 @@ def _resolve_paths(self, dataset_config: DatasetConfig, schema_dir: Path): def load_dataset(self, data_upload_config: DataUploadConfig) -> None: """ Load datasets from DataUploadConfig and dispatch to corresponding loaders - + Args: data_upload_config: User uploaded data configuration object - + Raises: ValueError: Unsupported dataset type """ @@ -230,19 +237,18 @@ def load_dataset(self, data_upload_config: DataUploadConfig) -> None: ) self.loaders[dtype].load_dataset(dataset) - def list_datasets(self, dtype: Optional[str] = None) -> Dict[str, List[str]]: """ List all or specified type of dataset names (dataset-level only) - + Args: dtype: Data type, can be "graph" / "table" / "text". If not specified, returns all types - + Returns: Dictionary with structure {"graph": [...], "table": [...], "text": [...]} Only returns dataset-level names from datasets.yaml, not file-level names - + Raises: ValueError: Unsupported data type """ @@ -250,12 +256,12 @@ def list_datasets(self, dtype: Optional[str] = None) -> Dict[str, List[str]]: # Get dataset-level names from datasets.yaml index result = {"graph": [], "table": [], "text": []} - + for dataset_name, dataset_info in self.datasets_index.items(): - dataset_type = dataset_info.get('type', '').lower() + dataset_type = dataset_info.get("type", "").lower() if dataset_type in self.VALID_DATA_TYPES: result[dataset_type].append(dataset_name) - + if not dtype: return result @@ -265,7 +271,6 @@ def list_datasets(self, dtype: Optional[str] = None) -> Dict[str, List[str]]: ) return {dtype: result[dtype]} - def get_dataset_content( self, dataset_config: DatasetConfig @@ -279,43 +284,58 @@ def get_dataset_content( text → list[str] """ dtype = dataset_config.type.lower() - + if dtype == "graph": return self.graph_loader.get_graph_content_from_raw(dataset_config) - + elif dtype == "table": path = dataset_config.schema.path import pandas as pd + + if not os.path.exists(path): + raise FileNotFoundError(f"表格文件不存在: {path}") return pd.read_csv(path) - + elif dtype == "text": path = dataset_config.schema.path encoding = getattr(dataset_config.schema, "encoding", "utf-8") + + if not os.path.exists(path): + raise FileNotFoundError(f"文本文件不存在: {path}") + # 大文件保护：超过 100MB 拒绝加载，建议使用流式读取 + file_size_mb = os.path.getsize(path) / (1024 * 1024) + if file_size_mb > 100: + raise ValueError( + f"文本文件过大 ({file_size_mb:.1f} MB)，超过 100MB 限制。" + " 建议使用流式读取或分块加载。" + ) with open(path, "r", encoding=encoding) as f: return f.readlines() - + else: raise ValueError(f"Unsupported data type: {dtype}") - def get_dataset_info(self, name: str, dtype: Optional[str] = None) -> Optional[List[DatasetConfig]]: + def get_dataset_info( + self, name: str, dtype: Optional[str] = None + ) -> Optional[List[DatasetConfig]]: """ Get detailed information of a dataset (schema) - + Args: name: Dataset name (dataset-level) dtype: Data type, can be "graph" / "table" / "text". If not specified, searches in all types and returns first match - + Returns: - For graph/table: List[DatasetConfig] with single config - For text: List[DatasetConfig] with multiple file configs - None if not found - + Raises: ValueError: Invalid data type """ dtype = (dtype or "").lower().strip() - + if dtype and dtype not in self.loaders: raise ValueError( f"Invalid data type '{dtype}'. Should be one of {', '.join(self.VALID_DATA_TYPES)}." @@ -327,74 +347,83 @@ def get_dataset_info(self, name: str, dtype: Optional[str] = None) -> Optional[L else: # Unspecified type: fuzzy search ds = next( - (loader.dataset_schemas[name] - for loader in self.loaders.values() - if name in loader.dataset_schemas), - None + ( + loader.dataset_schemas[name] + for loader in self.loaders.values() + if name in loader.dataset_schemas + ), + None, ) return ds - + def get_dataset_original_type(self, dataset_name: str) -> Optional[str]: """ Get the original type of a dataset from datasets.yaml - + Args: dataset_name: Name of the dataset - + Returns: "graph", "text", "table", or None if not found """ if dataset_name in self.datasets_index: - return self.datasets_index[dataset_name].get('type', '').lower() + return self.datasets_index[dataset_name].get("type", "").lower() return None - + def get_converted_graph_dataset(self, dataset_name: str) -> Optional[DatasetConfig]: """ Get the converted graph dataset config for a text dataset. Returns None if the dataset has not been converted to graph. - + This method combines the check and retrieval logic to avoid duplicate file loading. - + Args: dataset_name: Name of the text dataset - + Returns: DatasetConfig if converted graph exists, None otherwise """ # Check if dataset exists in index if dataset_name not in self.datasets_index: return None - + dataset_info = self.datasets_index[dataset_name] - original_type = dataset_info.get('type', '').lower() - - if original_type != 'text': + original_type = dataset_info.get("type", "").lower() + + if original_type != "text": return None - data_path = dataset_info.get('data_path', dataset_name) + data_path = dataset_info.get("data_path", dataset_name) dataset_folder = self.datasets_schema_path / data_path graph_schema_file = dataset_folder / "graph_schemas.yaml" - + if not graph_schema_file.exists(): return None - + try: - schemas = load_data_upload_config(str(graph_schema_file), validate_files=False) + schemas = load_data_upload_config( + str(graph_schema_file), validate_files=False + ) if schemas.datasets: graph_config = schemas.datasets[0] self._resolve_paths(graph_config, graph_schema_file.parent) return graph_config except Exception as e: print(f"[WARNING] Failed to load converted graph for {dataset_name}: {e}") - + return None - + # add gjq: 将图数据加载到Neo4j数据库 - def load_graph_to_neo4j(self, graph_dataset_config: DatasetConfig, dataset_name: str, neo4j_config: Dict[str, Any]): + def load_graph_to_neo4j( + self, + graph_dataset_config: DatasetConfig, + dataset_name: str, + neo4j_config: Dict[str, Any], + ): """ 将图数据加载到Neo4j数据库 - + Args: graph_dataset_config: 图数据集配置（DatasetConfig类型） - graph_dataset_config.schema 是 GraphSchemaConfig 类型 @@ -402,7 +431,7 @@ def load_graph_to_neo4j(self, graph_dataset_config: DatasetConfig, dataset_name: - GraphSchemaConfig.edge 是 List[EdgeSchemaConfig] dataset_name: 数据集名称 neo4j_config: Neo4j配置字典，包含 uri, user, password, enabled 等字段 - + Returns: bool: 加载是否成功 """ @@ -411,57 +440,66 @@ def load_graph_to_neo4j(self, graph_dataset_config: DatasetConfig, dataset_name: if not neo4j_config.get("enabled", False): logger.info("ℹ Neo4j未启用，跳过图数据加载") return False - + # add gjq: 从DatasetConfig中提取GraphSchemaConfig schema_obj = graph_dataset_config.schema # GraphSchemaConfig类型 - + # add gjq: 将schema对象转换为字典格式 # schema_obj.vertex 是 List[VertexSchemaConfig] # schema_obj.edge 是 List[EdgeSchemaConfig] # 需要将这些dataclass对象转换为字典列表 - if hasattr(schema_obj, 'vertex') and hasattr(schema_obj, 'edge'): + if hasattr(schema_obj, "vertex") and hasattr(schema_obj, "edge"): # 直接访问dataclass的属性 vertex_configs = schema_obj.vertex edge_configs = schema_obj.edge - + # 将dataclass对象转换为字典 - vertex_dicts = [asdict(v) if hasattr(v, '__dataclass_fields__') else v for v in vertex_configs] - edge_dicts = [asdict(e) if hasattr(e, '__dataclass_fields__') else e for e in edge_configs] + vertex_dicts = [ + asdict(v) if hasattr(v, "__dataclass_fields__") else v + for v in vertex_configs + ] + edge_dicts = [ + asdict(e) if hasattr(e, "__dataclass_fields__") else e + for e in edge_configs + ] else: # 降级方案：尝试使用__dict__ - schema_dict = schema_obj.__dict__ if hasattr(schema_obj, '__dict__') else schema_obj - vertex_dicts = schema_dict.get('vertex', []) - edge_dicts = schema_dict.get('edge', []) - + schema_dict = ( + schema_obj.__dict__ + if hasattr(schema_obj, "__dict__") + else schema_obj + ) + vertex_dicts = schema_dict.get("vertex", []) + edge_dicts = schema_dict.get("edge", []) + # add gjq: 构建Neo4j加载器需要的schema格式 - loader_schema = { - 'vertex': vertex_dicts, - 'edge': edge_dicts - } - + loader_schema = {"vertex": vertex_dicts, "edge": edge_dicts} + logger.info(f"🔄 开始将图数据集 {dataset_name} 加载到Neo4j...") - logger.info(f" 顶点配置数量: {len(vertex_dicts)}, 边配置数量: {len(edge_dicts)}") - + logger.info( + f" 顶点配置数量: {len(vertex_dicts)}, 边配置数量: {len(edge_dicts)}" + ) + # add gjq: 创建Neo4j加载器 loader = Neo4jGraphLoader( uri=neo4j_config.get("uri", "bolt://localhost:7687"), username=neo4j_config.get("user", "neo4j"), password=neo4j_config.get("password", ""), - schema=loader_schema + schema=loader_schema, ) - + # add gjq: 加载数据（清空现有数据） success = loader.load_all_data(clear_existing=True) - + if success: logger.info(f"✅ 图数据集 {dataset_name} 已成功加载到Neo4j") else: logger.warning(f"⚠️ 图数据集 {dataset_name} 加载到Neo4j失败") - + return success - + except Exception as e: logger.warning(f"⚠️ 加载图数据到Neo4j时发生错误: {e}") logger.exception(e) # add gjq: 打印详细的异常堆栈信息 # 不抛出异常，允许系统在没有Neo4j的情况下继续运行 - return False \ No newline at end of file + return False diff --git a/aag/engine/dependency_resolver.py b/aag/engine/dependency_resolver.py index 7e6cf8a..9be1f1c 100644 --- a/aag/engine/dependency_resolver.py +++ b/aag/engine/dependency_resolver.py @@ -14,27 +14,63 @@ from aag.utils.data_utils import take_sample from aag.error_recovery.error_manager import ErrorRecovery +# 最小安全 builtins（禁止 __import__/open/eval/exec 等危险函数） +_MINIMAL_BUILTINS = { + "len": len, + "sorted": sorted, + "reversed": reversed, + "enumerate": enumerate, + "zip": zip, + "map": map, + "filter": filter, + "sum": sum, + "min": min, + "max": max, + "abs": abs, + "round": round, + "int": int, + "float": float, + "str": str, + "bool": bool, + "list": list, + "dict": dict, + "set": set, + "tuple": tuple, + "range": range, + "any": any, + "all": all, + "isinstance": isinstance, + "type": type, + "print": print, + "json": __import__("json"), + "math": __import__("math"), +} + logger = logging.getLogger(__name__) + class DataDependencyType(Enum): """依赖类型枚举""" - NONE = "none" # 无依赖 - GRAPH = "graph" # 需要构造子图 - PARAMETER = "parameter" # 超参数依赖 - BOTH = "both" # 同时需要子图和参数 + + NONE = "none" # 无依赖 + GRAPH = "graph" # 需要构造子图 + PARAMETER = "parameter" # 超参数依赖 + BOTH = "both" # 同时需要子图和参数 @dataclass class SingleDependencyItem: """单个依赖的数据项（来自 selected_outputs[i]）""" + parent_step_id: int parent_step_output_id: int field_key: str field_type: Any field_desc: str - use_as: str # "graph" | "parameter" - value: Any # 真实数据：来自 StepOutputItem.value[field_key] - reason: str # LLM 给的解释文本 + use_as: str # "graph" | "parameter" + value: Any # 真实数据：来自 StepOutputItem.value[field_key] + reason: str # LLM 给的解释文本 + @dataclass class DataDependencyInfo: @@ -42,6 +78,7 @@ class DataDependencyInfo: 一个父节点的整体依赖信息： 可能包含多个 selected_outputs (multiple SingleDependencyItem) """ + parent_step_id: int parent_question: str dependency_type: DataDependencyType @@ -51,39 +88,43 @@ class DataDependencyInfo: class DataDependencyResolver: """ 依赖解析器 - + 职责： 1. 分析上游节点的输出与当前节点的依赖关系 2. 执行数据转换（子图构造、参数提取） 3. 管理子图缓存 """ - - def __init__(self, reasoner: Reasoner, error_recovery: Optional[ErrorRecovery] = None): + + def __init__( + self, reasoner: Reasoner, error_recovery: Optional[ErrorRecovery] = None + ): self.reasoner = reasoner self.error_recovery = error_recovery self.global_vertices: Optional[List[str]] = None self.global_edges: Optional[List[Tuple[str, str]]] = None - - def set_global_graph(self, graph_nodes: List[str] , graph_edges: List[Tuple[str, str]]): + + def set_global_graph( + self, graph_nodes: List[str], graph_edges: List[Tuple[str, str]] + ): self.global_vertices = graph_nodes self.global_edges = graph_edges - + async def resolve_dependencies( self, step_id: str, step: WorkflowStep, alg_des_info: dict, - data_dependency_parents: List[WorkflowStep] + data_dependency_parents: List[WorkflowStep], ) -> Dict[str, Any]: """ - 分析并转换上游依赖数据 + 分析并转换上游依赖数据 """ if not data_dependency_parents: return { "graph_dependencies": [], "parameter_dependencies": [], "graph_input_adapter_result": None, - "parameter_input_adapter_result": None + "parameter_input_adapter_result": None, } alg_des_doc = None @@ -94,45 +135,14 @@ async def resolve_dependencies( f"```json\n{json.dumps(alg_des_info.get('input_params'), indent=2, ensure_ascii=False)}" ) - logger.info(f"📊 [依赖解析] 当前步骤:{step_id} | 上游节点:{len(data_dependency_parents)}") - - - - # # ---------- Step 1：对每个父节点执行 “分类 + 定位” ---------- - # dependencies: List[DataDependencyInfo] = [] - # parent_step_map = {} - # for parent_step in data_dependency_parents: - # parent_step_map[parent_step.step_id]=parent_step - # if not parent_step.result: - # logger.warning( - # f"⚠️ 父节点 {parent_step.step_id}: {parent_step.question} 没有 result，跳过依赖分析" - # ) - # continue - - # dep_info = await self._classify_and_locate_dependency( - # current_step = step, - # current_algo_desc=alg_des_doc, - # parent_step=parent_step - # ) - # dependencies.append(dep_info) - - - # # Step 2: 把所有依赖项拆分为 graph / parameter 两类 - # graph_items: List[SingleDependencyItem] = [] - # param_items: List[SingleDependencyItem] = [] - - # for dep in dependencies: - # for item in dep.items: - # if item.use_as == "graph": - # graph_items.append(item) - # elif item.use_as == "parameter": - # param_items.append(item) - - - + logger.info( + f"📊 [依赖解析] 当前步骤:{step_id} | 上游节点:{len(data_dependency_parents)}" + ) # ---------- Step 1/2：分类 + 定位（空依赖时自动重试） ---------- max_dependency_retry = 3 - parent_step_map = {parent_step.step_id: parent_step for parent_step in data_dependency_parents} + parent_step_map = { + parent_step.step_id: parent_step for parent_step in data_dependency_parents + } graph_items: List[SingleDependencyItem] = [] param_items: List[SingleDependencyItem] = [] @@ -151,7 +161,7 @@ async def resolve_dependencies( dep_info = await self._classify_and_locate_dependency( current_step=step, current_algo_desc=alg_des_doc, - parent_step=parent_step + parent_step=parent_step, ) dependencies.append(dep_info) @@ -179,26 +189,28 @@ async def resolve_dependencies( logger.warning( f"⚠️ 依赖项为空，已达到最大重试次数（{max_dependency_retry}）" ) - + logger.info( f"📌 依赖整理完成 | graph={len(graph_items)} | parameter={len(param_items)}" - ) + ) # Step 3: 根据任务类型执行不同策略 graph_input_adapter_result = None parameter_input_adapter_result = None # Case A：当前任务是图算法 —— 需要子图优先 - if step.task_type == GraphAnalysisType.GRAPH_ALGORITHM: + if step.task_type == GraphAnalysisType.GRAPH_ALGORITHM: logger.info("🧠 当前任务为图算法 → 交给 Reasoner 写转换代码") # 处理图依赖 ---- if graph_items: conver_graph_llm_info = await self._convert_graph_dependencies( current_question=step.question, graph_items=graph_items, - parent_steps=parent_step_map + parent_steps=parent_step_map, + ) + graph_input_adapter_result = conver_graph_llm_info.get( + "converted_graph" ) - graph_input_adapter_result = conver_graph_llm_info.get("converted_graph") # 处理参数依赖 ---- if param_items: @@ -207,26 +219,28 @@ async def resolve_dependencies( current_question=step.question, alg_des_doc=alg_des_doc, param_items=param_items, - parent_steps=parent_step_map + parent_steps=parent_step_map, + ) + parameter_input_adapter_result = convert_parameter_llm_info.get( + "mapped_params" ) - parameter_input_adapter_result = convert_parameter_llm_info.get("mapped_params") elif step.task_type == GraphAnalysisType.NUMERIC_ANALYSIS: logger.info("🧮 当前任务类型：Numeric Analysis → 返回原始依赖项") else: logger.info("ℹ️ 非图算法/非数值分析 → 不做适配，只返回依赖项") - + return { "graph_dependencies": graph_items, "parameter_dependencies": param_items, "graph_input_adapter_result": graph_input_adapter_result, - "parameter_input_adapter_result": parameter_input_adapter_result + "parameter_input_adapter_result": parameter_input_adapter_result, } - + async def _classify_and_locate_dependency( self, current_step: WorkflowStep, current_algo_desc: str, - parent_step: WorkflowStep + parent_step: WorkflowStep, ) -> DataDependencyInfo: """ 使用 Reasoner 判定依赖类型，并精确定位依赖的上游数据项。 @@ -240,21 +254,23 @@ async def _classify_and_locate_dependency( # Step 1: 调用 LLM 进行依赖判定与定位（带 retry） if self.error_recovery: + async def op_analyze_dependency(error_history): return self.reasoner.analyze_dependency_type_and_locate_dependency_data( current_question=current_step.question, task_type=current_step.task_type, current_algo_desc=current_algo_desc, parent_question=parent_step.question, - parent_outputs_meta=parent_step.get_result_meta() + parent_outputs_meta=parent_step.get_result_meta(), ) - + try: analysis = await self.error_recovery.run( op_analyze_dependency, name=f"analyze_dependency(parent={parent_step.step_id},current={current_step.step_id})", operation_type="dependency_analysis", - location=f"step_{current_step.step_id}" + location=f"step_{current_step.step_id}", + prompt=current_step.question, ) except Exception as e: logger.info(f" Dependency analysis failed after retries: {e}") @@ -265,7 +281,7 @@ async def op_analyze_dependency(error_history): task_type=current_step.task_type, current_algo_desc=current_algo_desc, parent_question=parent_step.question, - parent_outputs_meta=parent_step.get_result_meta() + parent_outputs_meta=parent_step.get_result_meta(), ) logger.info(f"analysis analysis:{analysis}") @@ -275,7 +291,7 @@ async def op_analyze_dependency(error_history): parent_step_id=parent_step.step_id, parent_question=parent_step.question, dependency_type=DataDependencyType.NONE, - items=[] + items=[], ) # Step 2: 解析依赖类型 @@ -288,8 +304,8 @@ async def op_analyze_dependency(error_history): selected_outputs = analysis.get("selected_outputs", []) - # Step 3: 遍历父节点输出 —— 找真实数据源 - parent_outputs: Dict[int, StepOutputItem] = parent_step.result or {} + # Step 3: 遍历父节点输出 —— 找真实数据源 + parent_outputs: Dict[int, StepOutputItem] = parent_step.result or {} dependency_items = [] for sel in selected_outputs: @@ -307,7 +323,7 @@ async def op_analyze_dependency(error_history): if match is None: logger.info(f" 未找到父节点 output_id={output_id}") continue - + value_dict = match.value or {} if field_key not in value_dict: logger.info( @@ -317,7 +333,6 @@ async def op_analyze_dependency(error_history): real_value = value_dict[field_key] - # ---- 获取字段类型/描述 ---- f_type = None f_desc = None @@ -336,7 +351,7 @@ async def op_analyze_dependency(error_history): field_desc=f_desc, use_as=use_as, value=real_value, - reason=reason + reason=reason, ) ) @@ -351,29 +366,27 @@ async def op_analyze_dependency(error_history): dependency_type=dep_type, items=dependency_items, ) - - async def _convert_graph_dependencies( self, current_question: str, graph_items: List[SingleDependencyItem], - parent_steps: Dict[int, WorkflowStep] + parent_steps: Dict[int, WorkflowStep], ) -> Dict[str, Any]: """ - 图依赖转换： - 1. 调用 Reasoner 生成图构造代码（基于 graph_items） - 2. 执行生成的 Python 代码 - 3. 返回 {"description": ..., "converted_graph": {...}, "raw_llm_response": ...} + 图依赖转换： + 1. 调用 Reasoner 生成图构造代码（基于 graph_items） + 2. 执行生成的 Python 代码 + 3. 返回 {"description": ..., "converted_graph": {...}, "raw_llm_response": ...} - graph_items: List[SingleDependencyItem] - parent_steps: { step_id -> WorkflowStep } + graph_items: List[SingleDependencyItem] + parent_steps: { step_id -> WorkflowStep } """ if not graph_items: return { "description": "No graph items were provided.", "converted_graph": None, - "raw_llm_response": None + "raw_llm_response": None, } if not self.global_vertices or not self.global_edges: @@ -382,7 +395,7 @@ async def _convert_graph_dependencies( return { "description": error_msg, "converted_graph": None, - "raw_llm_response": None + "raw_llm_response": None, } # Step 1: 构造传给 LLM 的 dependency_items（带 sample_value） @@ -392,18 +405,21 @@ async def _convert_graph_dependencies( parent_step = parent_steps[item.parent_step_id] raw_data = item.value sample_value = take_sample(raw_data) - dependency_list.append({ - "field_key": item.field_key, - "field_type": item.field_type, - "field_desc": item.field_desc, - "sample_value": sample_value, - "parent_step_id": item.parent_step_id, - "parent_step_question": parent_step.question, - "reason": item.reason - }) + dependency_list.append( + { + "field_key": item.field_key, + "field_type": item.field_type, + "field_desc": item.field_desc, + "sample_value": sample_value, + "parent_step_id": item.parent_step_id, + "parent_step_question": parent_step.question, + "reason": item.reason, + } + ) # Step 1: 调用 LLM 生成图转换代码（带 retry） if self.error_recovery: + async def op_generate_and_execute_graph_code(error_history): # Step 1: 调用 LLM 生成图转换代码 llm_resp = self.reasoner.generate_graph_conversion_code( @@ -422,12 +438,15 @@ async def op_generate_and_execute_graph_code(error_history): field_keys_in_order = [i.field_key for i in graph_items] exec_env = { "global_nodes": self.global_vertices, - "global_edges": self.global_edges + "global_edges": self.global_edges, + "__builtins__": _MINIMAL_BUILTINS, } exec(code, exec_env) fn = exec_env.get("transform_graph") if not fn or not callable(fn): - raise ValueError("Generated code does not define a valid transform_graph function.") + raise ValueError( + "Generated code does not define a valid transform_graph function." + ) args = [upstream_values[k] for k in field_keys_in_order] args += [self.global_vertices, self.global_edges] converted_graph = fn(*args) @@ -438,19 +457,21 @@ async def op_generate_and_execute_graph_code(error_history): op_generate_and_execute_graph_code, name=f"generate_and_execute_graph_code(question={current_question[:50]})", operation_type="generic", - location="graph_conversion" + location="graph_conversion", ) return { "description": description, "converted_graph": converted_graph, - "raw_llm_response": llm_resp + "raw_llm_response": llm_resp, } except Exception as e: - logger.info(f" Graph conversion (generate+execute) failed after retries: {e}") + logger.info( + f" Graph conversion (generate+execute) failed after retries: {e}" + ) return { "description": str(e), "converted_graph": None, - "raw_llm_response": None + "raw_llm_response": None, } else: llm_resp = self.reasoner.generate_graph_conversion_code( @@ -461,7 +482,7 @@ async def op_generate_and_execute_graph_code(error_history): return { "description": "LLM returned empty graph conversion result.", "converted_graph": None, - "raw_llm_response": None + "raw_llm_response": None, } code = llm_resp.get("code") description = llm_resp.get("description", "") @@ -469,14 +490,15 @@ async def op_generate_and_execute_graph_code(error_history): return { "description": "LLM did not return valid code.", "converted_graph": None, - "raw_llm_response": llm_resp + "raw_llm_response": llm_resp, } try: upstream_values = {i.field_key: i.value for i in graph_items} field_keys_in_order = [i.field_key for i in graph_items] exec_env = { "global_nodes": self.global_vertices, - "global_edges": self.global_edges + "global_edges": self.global_edges, + "__builtins__": _MINIMAL_BUILTINS, } exec(code, exec_env) fn = exec_env.get("transform_graph") @@ -484,7 +506,7 @@ async def op_generate_and_execute_graph_code(error_history): return { "description": "Generated code does not define a valid transform_graph function.", "converted_graph": None, - "raw_llm_response": llm_resp + "raw_llm_response": llm_resp, } args = [upstream_values[k] for k in field_keys_in_order] args += [self.global_vertices, self.global_edges] @@ -493,21 +515,20 @@ async def op_generate_and_execute_graph_code(error_history): return { "description": f"Error executing generated code: {e}", "converted_graph": None, - "raw_llm_response": llm_resp + "raw_llm_response": llm_resp, } return { "description": description, "converted_graph": converted_graph, - "raw_llm_response": llm_resp + "raw_llm_response": llm_resp, } - async def _convert_parameter_dependencies( self, current_question: str, alg_des_doc: str, param_items: List[SingleDependencyItem], - parent_steps: Dict[int, WorkflowStep] + parent_steps: Dict[int, WorkflowStep], ) -> Dict[str, Any]: """ 参数依赖适配： @@ -516,11 +537,7 @@ async def _convert_parameter_dependencies( 3. 返回 {"description":..., "mapped_params": {...}} """ if not param_items: - return { - "mapped_params": {}, - "mapping_raw": None, - "reasoning": None - } + return {"mapped_params": {}, "mapping_raw": None, "reasoning": None} # 构造输入 dependency_list = [] @@ -528,27 +545,32 @@ async def _convert_parameter_dependencies( parent_step = parent_steps[item.parent_step_id] raw_data = item.value sample_value = take_sample(raw_data) - dependency_list.append({ - "field_key": item.field_key, - "field_type": item.field_type, - "field_desc": item.field_desc, - "sample_value": sample_value, - "parent_step_id": item.parent_step_id, - "parent_step_question": parent_step.question, - "reason": item.reason - }) + dependency_list.append( + { + "field_key": item.field_key, + "field_type": item.field_type, + "field_desc": item.field_desc, + "sample_value": sample_value, + "parent_step_id": item.parent_step_id, + "parent_step_question": parent_step.question, + "reason": item.reason, + } + ) # Step 1: 调用 LLM 生成参数映射（带 retry） if self.error_recovery: + async def op_map_parameters_and_execute(error_history): mapping_result = self.reasoner.map_parameters( current_question=current_question, current_algo_desc=alg_des_doc, - dependency_items=dependency_list + dependency_items=dependency_list, ) if "mapping" not in mapping_result: raise ValueError( - mapping_result.get("explanation", "LLM returned unexpected format") + mapping_result.get( + "explanation", "LLM returned unexpected format" + ) ) mapping = mapping_result["mapping"] final_params = {} @@ -557,9 +579,13 @@ async def op_map_parameters_and_execute(error_history): parent_step_id = mp.get("parent_step_id") extract_code = mp.get("extract_code") matched_item = next( - (x for x in param_items - if x.parent_step_id == parent_step_id and x.field_key == from_field), - None + ( + x + for x in param_items + if x.parent_step_id == parent_step_id + and x.field_key == from_field + ), + None, ) if matched_item is None: raise ValueError( @@ -569,7 +595,11 @@ async def op_map_parameters_and_execute(error_history): if extract_code is None: final_params[param_name] = real_value continue - exec_env = {"value": real_value, "param_value": None} + exec_env = { + "value": real_value, + "param_value": None, + "__builtins__": _MINIMAL_BUILTINS, + } exec(extract_code, exec_env) if "param_value" not in exec_env: raise ValueError( @@ -583,31 +613,33 @@ async def op_map_parameters_and_execute(error_history): op_map_parameters_and_execute, name=f"map_parameters+execute(question={current_question})", operation_type="generic", - location="parameter_conversion" + location="parameter_conversion", ) return { "mapped_params": final_params, "mapping_raw": mapping_result, - "reasoning": mapping_result.get("explanation", "") + "reasoning": mapping_result.get("explanation", ""), } except Exception as e: logger.info(f" Parameter mapping+execute failed after retries: {e}") return { "mapped_params": {}, "mapping_raw": None, - "reasoning": f"Parameter mapping+execute failed: {e}" + "reasoning": f"Parameter mapping+execute failed: {e}", } else: mapping_result = self.reasoner.map_parameters( current_question=current_question, current_algo_desc=alg_des_doc, - dependency_items=dependency_list + dependency_items=dependency_list, ) if "mapping" not in mapping_result: return { "mapped_params": {}, "mapping_raw": mapping_result, - "reasoning": mapping_result.get("explanation", "LLM returned unexpected format") + "reasoning": mapping_result.get( + "explanation", "LLM returned unexpected format" + ), } mapping = mapping_result["mapping"] final_params = {} @@ -616,9 +648,13 @@ async def op_map_parameters_and_execute(error_history): parent_step_id = mp.get("parent_step_id") extract_code = mp.get("extract_code") matched_item = next( - (x for x in param_items - if x.parent_step_id == parent_step_id and x.field_key == from_field), - None + ( + x + for x in param_items + if x.parent_step_id == parent_step_id + and x.field_key == from_field + ), + None, ) if matched_item is None: logger.info(f" 参数 {param_name} 的依赖字段找不到真实数据源") @@ -627,15 +663,21 @@ async def op_map_parameters_and_execute(error_history): if extract_code is None: final_params[param_name] = real_value continue - exec_env = {"value": real_value, "param_value": None} + exec_env = { + "value": real_value, + "param_value": None, + "__builtins__": _MINIMAL_BUILTINS, + } try: exec(extract_code, exec_env) final_params[param_name] = exec_env["param_value"] except Exception as e: - logger.info(f" 执行 extract_code 失败: {e}\n参数: {param_name}\n代码:\n{extract_code}") + logger.info( + f" 执行 extract_code 失败: {e}\n参数: {param_name}\n代码:\n{extract_code}" + ) continue return { "mapped_params": final_params, "mapping_raw": mapping_result, - "reasoning": mapping_result.get("explanation", "") - } \ No newline at end of file + "reasoning": mapping_result.get("explanation", ""), + } diff --git a/aag/error_recovery/error_manager.py b/aag/error_recovery/error_manager.py index fa420b0..05837fe 100644 --- a/aag/error_recovery/error_manager.py +++ b/aag/error_recovery/error_manager.py @@ -11,7 +11,9 @@ logger = logging.getLogger(__name__) -def prepare_error_info(error: Exception, *, location: Optional[str] = None, hint: Optional[str] = None) -> Dict[str, Any]: +def prepare_error_info( + error: Exception, *, location: Optional[str] = None, hint: Optional[str] = None +) -> Dict[str, Any]: info = { "error_type": type(error).__name__, "error": str(error), @@ -29,7 +31,9 @@ def __init__(self, *, trace_maxlen: int = 200): self.trace = PromptTraceBuffer(maxlen=trace_maxlen) # ---------- Trace ---------- - def record_prompt(self, fn_name: str, base_prompt: str, meta: Optional[Dict[str, Any]] = None) -> None: + def record_prompt( + self, fn_name: str, base_prompt: str, meta: Optional[Dict[str, Any]] = None + ) -> None: self.trace.record(fn_name, base_prompt, meta=meta) def get_last_base_prompt(self, fn_name: str) -> Optional[str]: @@ -64,7 +68,18 @@ async def run( name: str, operation_type: str = "generic", location: Optional[str] = None, + prompt: Optional[str] = None, ) -> Any: + """ + 带重试的操作执行器。 + + 参数： + operation: 接收 error_history 并返回结果的异步操作。 + name: 操作名称，用于日志和 prompt 跟踪。 + operation_type: 操作类型，用于策略选择。 + location: 错误发生位置标识。 + prompt: 操作的原始 prompt，用于记录和增强重试。 + """ policy = get_policy(operation_type) max_attempts = policy.max_attempts max_error_history = policy.max_error_history @@ -72,11 +87,35 @@ async def run( error_history: List[Dict[str, Any]] = [] last_exc: Optional[Exception] = None - for attempt in range(max_attempts+1): + for attempt in range(max_attempts + 1): + # 首次尝试前记录原始 prompt；重试时构建增强 prompt + if attempt == 0 and prompt: + self.record_prompt( + name, + prompt, + meta={"operation_type": operation_type, "location": location}, + ) + elif attempt > 0: + enhanced = self.build_enhanced_prompt( + fn_name=name, + error_history=error_history, + operation_type=operation_type, + ) + if enhanced: + error_history.append( + { + "type": "enhanced_prompt", + "content": enhanced, + "attempt": attempt + 1, + } + ) + try: result = await operation(error_history) if attempt > 1: - logger.info("✅ %s succeeded on attempt %d/%d", name, attempt, max_attempts) + logger.info( + "✅ %s succeeded on attempt %d/%d", name, attempt, max_attempts + ) return result except Exception as e: last_exc = e @@ -85,7 +124,13 @@ async def run( error_history[:] = error_history[-max_error_history:] if attempt < max_attempts: - logger.warning("⚠️ %s failed on attempt %d/%d: %s", name, attempt, max_attempts, err["error"]) + logger.warning( + "⚠️ %s failed on attempt %d/%d: %s", + name, + attempt, + max_attempts, + err["error"], + ) else: logger.error("❌ %s failed after %d attempts", name, max_attempts) @@ -154,7 +199,9 @@ def apply_cross_step_recovery( for tid in target_step_ids: try: dag.set_pending(tid) - logger.info(f"🔄 Cross-step recovery: reset step {tid} -> pending ({reason})") + logger.info( + f"🔄 Cross-step recovery: reset step {tid} -> pending ({reason})" + ) except Exception as e: logger.warning(f"⚠️ Failed to reset step {tid} to pending: {e}") @@ -171,12 +218,15 @@ def global_recover( Decide + apply cross-step recovery in one call. Returns the target step ids that were reset. """ - targets = self.decide_cross_step_recovery(dag, step_id, error_info, context=context) + targets = self.decide_cross_step_recovery( + dag, step_id, error_info, context=context + ) if not targets: return [] self.apply_cross_step_recovery( dag, targets, - reason=reason or f"cross_step_recovery(error_type={error_info.get('error_type')})", + reason=reason + or f"cross_step_recovery(error_type={error_info.get('error_type')})", ) - return targets \ No newline at end of file + return targets diff --git a/aag/expert_search_engine/data_process/openai_extractor/openai_config.json b/aag/expert_search_engine/data_process/openai_extractor/openai_config.json deleted file mode 100644 index 39866d3..0000000 --- a/aag/expert_search_engine/data_process/openai_extractor/openai_config.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "base_url": "https://gitaigc.com/v1/", - "api_key": "sk-wkQq5dWpyYpBtl8PYH0jI85ct0TsobEsXg6AUwEI5OLrfu8K", - "model": "gpt-4o" -} \ No newline at end of file diff --git a/aag/expert_search_engine/data_process/openai_extractor/openai_config.json.example b/aag/expert_search_engine/data_process/openai_extractor/openai_config.json.example new file mode 100644 index 0000000..73e871e --- /dev/null +++ b/aag/expert_search_engine/data_process/openai_extractor/openai_config.json.example @@ -0,0 +1,6 @@ +{ + "_comment": "将 api_key 设置为你的实际密钥，或通过环境变量 OPENAI_API_KEY 注入", + "base_url": "https://gitaigc.com/v1/", + "api_key": "YOUR_API_KEY", + "model": "gpt-4o" +} diff --git a/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json b/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json deleted file mode 100644 index c26ddef..0000000 --- a/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json +++ /dev/null @@ -1,5 +0,0 @@ -{ - "base_url": "https://gitaigc.com/v1/", - "api_key": "sk-McYA8mFqvB7XJh9PaTkpWcXjHftRmbAtzt0FP4OQHWbLdqj1", - "model": "gpt-4o" -} \ No newline at end of file diff --git a/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json.example b/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json.example new file mode 100644 index 0000000..73e871e --- /dev/null +++ b/aag/expert_search_engine/data_process/openai_extractor/openai_config_4o.json.example @@ -0,0 +1,6 @@ +{ + "_comment": "将 api_key 设置为你的实际密钥，或通过环境变量 OPENAI_API_KEY 注入", + "base_url": "https://gitaigc.com/v1/", + "api_key": "YOUR_API_KEY", + "model": "gpt-4o" +} diff --git a/aag/expert_search_engine/database/nebulagraph.py b/aag/expert_search_engine/database/nebulagraph.py index 973a42d..bba739f 100644 --- a/aag/expert_search_engine/database/nebulagraph.py +++ b/aag/expert_search_engine/database/nebulagraph.py @@ -1,8 +1,10 @@ import os + # from typing import Any, Callable, Dict, List, Optional, Set, Tuple from typing import Any, Dict, List, Optional, Tuple from llama_index.core import StorageContext from llama_index.core import KnowledgeGraphIndex + # from llama_index.core.graph_stores.nebula import NebulaGraphStore from llama_index.legacy.graph_stores.nebulagraph import NebulaGraphStore @@ -14,6 +16,7 @@ from nebula3.Config import Config from nebula3.data.ResultSet import ResultSet from nebula3.data.DataObject import ValueWrapper + # from nebula3.common import * import json @@ -31,6 +34,7 @@ from llama_index.core import load_index_from_storage from llama_index.core.utils import print_text import re + # from utils.utils import create_dir from llama_index.core.schema import ( # BaseNode, @@ -41,7 +45,8 @@ ) from llama_index.core.retrievers import ( - KnowledgeGraphRAGRetriever, ) + KnowledgeGraphRAGRetriever, +) from llama_index.embeddings.huggingface import HuggingFaceEmbedding from aag.utils.extract_subgraph import filter_pr_rels @@ -52,16 +57,15 @@ class NebulaClient: - def __init__(self): config = Config() config.max_connection_pool_size = 10 self.connection_pool = ConnectionPool() - ok = self.connection_pool.init([('127.0.0.1', 9669)], config) + ok = self.connection_pool.init([("127.0.0.1", 9669)], config) assert ok - self.session = self.connection_pool.get_session('root', 'nebula') + self.session = self.connection_pool.get_session("root", "nebula") def __del__(self): if self.connection_pool: @@ -71,16 +75,17 @@ def __del__(self): def create_space(self, space_name): self.session.execute( - f'CREATE SPACE IF NOT EXISTS {space_name}(vid_type=FIXED_STRING(256), partition_num=1, replica_factor=1);' + f"CREATE SPACE IF NOT EXISTS {space_name}(vid_type=FIXED_STRING(256), partition_num=1, replica_factor=1);" ) time.sleep(10) self.session.execute( - f'USE {space_name}; CREATE TAG IF NOT EXISTS entity(name string);') + f"USE {space_name}; CREATE TAG IF NOT EXISTS entity(name string);" + ) self.session.execute( - f'USE {space_name}; CREATE EDGE IF NOT EXISTS relationship(relationship string);' + f"USE {space_name}; CREATE EDGE IF NOT EXISTS relationship(relationship string);" ) self.session.execute( - f'USE {space_name}; CREATE TAG INDEX IF NOT EXISTS entity_index ON entity(name(256));' + f"USE {space_name}; CREATE TAG INDEX IF NOT EXISTS entity_index ON entity(name(256));" ) time.sleep(10) @@ -88,104 +93,112 @@ def drop_space(self, space_name): if not isinstance(space_name, list): space_name = [space_name] for space in space_name: - self.session.execute(f'drop space {space}') + self.session.execute(f"drop space {space}") def info(self, space_name): result = self.session.execute( - f'use {space_name}; submit job stats; show stats;') + f"use {space_name}; submit job stats; show stats;" + ) print(result) print_resp(result) def count_edges(self, space_name): result = self.session.execute( - f'use {space_name}; MATCH (m)-[e]->(n) RETURN COUNT(*);') + f"use {space_name}; MATCH (m)-[e]->(n) RETURN COUNT(*);" + ) print_resp(result) def show_space(self): - result = self.session.execute('SHOW SPACES;') + result = self.session.execute("SHOW SPACES;") print_resp(result) def show_edges(self, space_name, limits): result = self.session.execute( - f'use {space_name}; MATCH ()-[e]->() RETURN e LIMIT {limits};') + f"use {space_name}; MATCH ()-[e]->() RETURN e LIMIT {limits};" + ) print_resp(result) def clear(self, space_name): - query = f'CLEAR SPACE {space_name};' + query = f"CLEAR SPACE {space_name};" self.session.execute(query) def get_triplets(self, space_name): result = self.session.execute( - f'use {space_name}; MATCH (n1)-[e]->(n2) RETURN n1, e, n2;') + f"use {space_name}; MATCH (n1)-[e]->(n2) RETURN n1, e, n2;" + ) def show_triplets(self, space_name, file_path=None): result = self.session.execute( - f'use {space_name}; MATCH (n1)-[e]->(n2) RETURN n1, e, n2;') + f"use {space_name}; MATCH (n1)-[e]->(n2) RETURN n1, e, n2;" + ) # print(result, len(result)) # assert False if not file_path: - file_path = space_name + '_triplets.txt' + file_path = space_name + "_triplets.txt" - json_path = file_path + '.json' + json_path = file_path + ".json" all_triples = [] - with open(file_path, 'w', encoding='utf-8') as file: + with open(file_path, "w", encoding="utf-8") as file: if result.row_size() > 0: for row in result.rows(): values = row.values - head = '' - relation = '' - tail = '' + head = "" + relation = "" + tail = "" for value in values: if value.field == 9: # 对应 Vertex vertex = value.get_vVal() if not head: - head = vertex.vid.get_sVal().decode('utf-8') + head = vertex.vid.get_sVal().decode("utf-8") else: - tail = vertex.vid.get_sVal().decode('utf-8') + tail = vertex.vid.get_sVal().decode("utf-8") elif value.field == 10: # 对应 Edge edge = value.get_eVal() - relation = edge.props.get( - b'relationship').get_sVal().decode('utf-8') + relation = ( + edge.props.get(b"relationship") + .get_sVal() + .decode("utf-8") + ) triplet = [head, relation, tail] all_triples.append(triplet) file.write(f"{triplet}\n") print( - f'triplet write to {file_path}, tot {len(result.rows())} triplets.' + f"triplet write to {file_path}, tot {len(result.rows())} triplets." ) else: - print('No data found.') - file.write('No data found.\n') + print("No data found.") + file.write("No data found.\n") all_triples = set(tuple(triplet) for triplet in all_triples) - print( - f'after filter the sample triplet last {len(all_triples)} triplets.' - ) + print(f"after filter the sample triplet last {len(all_triples)} triplets.") all_triples = [list(triplet) for triplet in all_triples] - with open(json_path, 'w', encoding='utf-8') as file: + with open(json_path, "w", encoding="utf-8") as file: json.dump(all_triples, file, ensure_ascii=False, indent=4) - print(f'save {len(all_triples)} triples to {json_path}.') + print(f"save {len(all_triples)} triples to {json_path}.") # 写一个 nebula 的类，将这个文件里（/home/chency/GraphLLM/graphllm/graph_data/AMLSim/load_data_into_nebulagraph.py）关于 nebulagraph的操作 提炼，写一个通用的nebula操作类 -# 需要完成的功能： 指定spacename， 创建tag（作为参数），创建edge_type(作为参数)， 插入点，插入边， 删除space， 提取整张图，提取所有顶点， 根据顶点集合返回每个顶点的k-hop子图 +# 需要完成的功能： 指定spacename， 创建tag（作为参数），创建edge_type(作为参数)， 插入点，插入边， 删除space， 提取整张图，提取所有顶点， 根据顶点集合返回每个顶点的k-hop子图 class NebulaGraphClient: - def __init__(self, - host='127.0.0.1', - port=9669, - username='root', - password='nebula', - vid_type='INT64', - partition_num=3, - replica_factor=1): + def __init__( + self, + host="127.0.0.1", + port=9669, + username="root", + password=os.getenv("NEBULA_PASSWORD", ""), + vid_type="INT64", + partition_num=3, + replica_factor=1, + ): self.host = host self.port = port self.username = username @@ -193,12 +206,12 @@ def __init__(self, self.vid_type = vid_type self.partition_num = partition_num self.replica_factor = replica_factor - + self.current_space: Optional[str] = None self.connection_pool = None self.session = None - self.vertices = [] - self.edges = [] + self.vertices = [] + self.edges = [] self._connect() def _connect(self): @@ -214,10 +227,10 @@ def use_space(self, space_name: str): raise RuntimeError(f"切换空间失败: {resp.error_msg()}") self.current_space = space_name - def create_space(self, space_name: str): resp = self.session.execute( - f"CREATE SPACE IF NOT EXISTS {space_name} (partition_num = {self.partition_num}, replica_factor = {self.replica_factor}, vid_type = {self.vid_type})") + f"CREATE SPACE IF NOT EXISTS {space_name} (partition_num = {self.partition_num}, replica_factor = {self.replica_factor}, vid_type = {self.vid_type})" + ) if not resp.is_succeeded(): raise RuntimeError(f"创建空间失败: {resp.error_msg()}") time.sleep(3) @@ -239,8 +252,10 @@ def create_tag(self, tag_name, fields): fields: list of (name, type) tuples, e.g. [('name', 'string'), ('age', 'int')] """ self._ensure_space_selected() - fields_str = ', '.join([f'{name} {ftype}' for name, ftype in fields]) - resp = self.session.execute(f"CREATE TAG IF NOT EXISTS {tag_name}({fields_str})") + fields_str = ", ".join([f"{name} {ftype}" for name, ftype in fields]) + resp = self.session.execute( + f"CREATE TAG IF NOT EXISTS {tag_name}({fields_str})" + ) if not resp.is_succeeded(): raise RuntimeError(f"创建Tag失败: {resp.error_msg()}") time.sleep(1) @@ -250,8 +265,10 @@ def create_edge_type(self, edge_name, fields): fields: list of (name, type) tuples """ self._ensure_space_selected() - fields_str = ', '.join([f'{name} {ftype}' for name, ftype in fields]) - resp = self.session.execute(f"CREATE EDGE IF NOT EXISTS {edge_name}({fields_str})") + fields_str = ", ".join([f"{name} {ftype}" for name, ftype in fields]) + resp = self.session.execute( + f"CREATE EDGE IF NOT EXISTS {edge_name}({fields_str})" + ) if not resp.is_succeeded(): raise RuntimeError(f"创建Edge Type失败: {resp.error_msg()}") time.sleep(40) @@ -263,28 +280,30 @@ def insert_vertex(self, tag_name, vid, prop_names, prop_values): """ if not self.current_space: raise RuntimeError("未选择任何图空间，请先调用 use_space(space_name)") - values_str = ', '.join([self._format_value(v) for v in prop_values]) - stmt = f'INSERT VERTEX {tag_name}({", ".join(prop_names)}) VALUES {vid}:({values_str})' + values_str = ", ".join([self._format_value(v) for v in prop_values]) + stmt = f"INSERT VERTEX {tag_name}({', '.join(prop_names)}) VALUES {vid}:({values_str})" resp = self.session.execute(stmt) if not resp.is_succeeded(): raise RuntimeError(f"插入顶点失败: {resp.error_msg()}\n语句: {stmt}") - def insert_edge(self, edge_name, src_vid, dst_vid, prop_names, prop_values, rank=None): + def insert_edge( + self, edge_name, src_vid, dst_vid, prop_names, prop_values, rank=None + ): """ prop_names: list of property names prop_values: list of property values (顺序与prop_names一致) rank: 可选，边的rank """ self._ensure_space_selected() - values_str = ', '.join([self._format_value(v) for v in prop_values]) + values_str = ", ".join([self._format_value(v) for v in prop_values]) if rank is not None: - stmt = f'INSERT EDGE {edge_name}({", ".join(prop_names)}) VALUES {src_vid} -> {dst_vid}@{rank}:({values_str})' + stmt = f"INSERT EDGE {edge_name}({', '.join(prop_names)}) VALUES {src_vid} -> {dst_vid}@{rank}:({values_str})" else: - stmt = f'INSERT EDGE {edge_name}({", ".join(prop_names)}) VALUES {src_vid} -> {dst_vid}:({values_str})' + stmt = f"INSERT EDGE {edge_name}({', '.join(prop_names)}) VALUES {src_vid} -> {dst_vid}:({values_str})" resp = self.session.execute(stmt) if not resp.is_succeeded(): raise RuntimeError(f"插入边失败: {resp.error_msg()}\n语句: {stmt}") - + # =============================== # 数据读取 # =============================== @@ -330,7 +349,7 @@ def get_all_vertices(self): props[key] = str(v) # fallback self.vertices.append(VertexData(vid=vid, properties=props)) return self.vertices - + def get_all_edges(self): """提取所有边和属性""" self._ensure_space_selected() @@ -339,8 +358,8 @@ def get_all_edges(self): if edge_resp.is_succeeded(): for row in edge_resp.rows(): e_val = row.values[0].get_eVal() - src = str(e_val.src.get_iVal()) # - dst = str(e_val.dst.get_iVal()) # + src = str(e_val.src.get_iVal()) # + dst = str(e_val.dst.get_iVal()) # edge_type = e_val.name.decode("utf-8") rank = e_val.ranking props = {} @@ -363,9 +382,11 @@ def get_all_edges(self): props[key] = v.as_datetime() else: props[key] = str(v) # fallback - self.edges.append(EdgeData(src=src, dst=dst, rank=rank, properties=props)) + self.edges.append( + EdgeData(src=src, dst=dst, rank=rank, properties=props) + ) return self.edges - + def get_k_hop_subgraph(self, vertex_ids, k=1): """ 根据顶点集合返回每个顶点的k-hop子图 @@ -377,7 +398,7 @@ def get_k_hop_subgraph(self, vertex_ids, k=1): result = {} for vid in vertex_ids: # 根据vid_type格式化查询条件 - if self.vid_type == 'INT64': + if self.vid_type == "INT64": query = f"MATCH p=(v)-[*1..{k}]-(n) WHERE id(v)=={vid} RETURN p" else: query = f'MATCH p=(v)-[*1..{k}]-(n) WHERE id(v)=="{vid}" RETURN p' @@ -392,19 +413,19 @@ def get_k_hop_subgraph(self, vertex_ids, k=1): def get_vertex_by_id(self, vid): """根据顶点ID获取顶点信息""" self._ensure_space_selected() - if self.vid_type == 'INT64': + if self.vid_type == "INT64": query = f"MATCH (v) WHERE id(v)=={vid} RETURN v" else: query = f'MATCH (v) WHERE id(v)=="{vid}" RETURN v' - + resp = self.session.execute(query) if resp.is_succeeded() and resp.row_size() > 0: v_val = resp.rows()[0].values[0].get_vVal() - if self.vid_type == 'INT64': + if self.vid_type == "INT64": vid_actual = int(v_val.vid.get_iVal()) else: - vid_actual = v_val.vid.get_sVal().decode('utf-8') - + vid_actual = v_val.vid.get_sVal().decode("utf-8") + props = {} for tag in v_val.tags: for k, v_raw in tag.props.items(): @@ -423,7 +444,7 @@ def get_vertex_by_id(self, vid): return VertexData(vid=vid_actual, properties=props) return None - def get_neighbors(self, vid, direction='both'): + def get_neighbors(self, vid, direction="both"): """ 获取指定顶点的邻居 :param vid: 顶点ID @@ -431,22 +452,22 @@ def get_neighbors(self, vid, direction='both'): :return: list of neighbor vertex IDs """ self._ensure_space_selected() - if direction == 'in': - if self.vid_type == 'INT64': + if direction == "in": + if self.vid_type == "INT64": query = f"MATCH (n)-[e]->(v) WHERE id(v)=={vid} RETURN id(n)" else: query = f'MATCH (n)-[e]->(v) WHERE id(v)=="{vid}" RETURN id(n)' - elif direction == 'out': - if self.vid_type == 'INT64': + elif direction == "out": + if self.vid_type == "INT64": query = f"MATCH (v)-[e]->(n) WHERE id(v)=={vid} RETURN id(n)" else: query = f'MATCH (v)-[e]->(n) WHERE id(v)=="{vid}" RETURN id(n)' else: # both - if self.vid_type == 'INT64': + if self.vid_type == "INT64": query = f"MATCH (v)-[e]-(n) WHERE id(v)=={vid} RETURN id(n)" else: query = f'MATCH (v)-[e]-(n) WHERE id(v)=="{vid}" RETURN id(n)' - + resp = self.session.execute(query) neighbors = [] if resp.is_succeeded(): @@ -482,20 +503,19 @@ def execute_query(self, query): def _format_value(self, v): if v is None: - return 'NULL' + return "NULL" elif isinstance(v, str): return f'"{v}"' elif isinstance(v, bool): return str(v).lower() else: return str(v) - + def _ensure_space_selected(self): """统一检查空间是否已选择""" if not self.current_space: raise RuntimeError("未选择任何图空间，请先调用 use_space(space_name)") - def close(self): if self.session: self.session.release() @@ -506,16 +526,16 @@ def __del__(self): self.close() - class NebulaDB: - - def __init__(self, - space_name, - log_file='./database/nebula.log', - server_ip='127.0.0.1', - server_port='9669', - create=False, - verbose=False): + def __init__( + self, + space_name, + log_file="./database/nebula.log", + server_ip="127.0.0.1", + server_port="9669", + create=False, + verbose=False, + ): # verbose=False, retriever=False, llm_env=None): self.log_file = log_file self.server_ip = server_ip @@ -526,9 +546,9 @@ def __init__(self, os.environ["NEBULA_ADDRESS"] = f"{self.server_ip}:{self.server_port}" self.space_name = space_name - self.edge_types = ['relationship'] - self.rel_prop_names = ['relationship'] - self.tags = ['entity'] + self.edge_types = ["relationship"] + self.rel_prop_names = ["relationship"] + self.tags = ["entity"] self.client = NebulaClient() self.verbose = verbose self.store: NebulaGraphStore = None @@ -537,7 +557,7 @@ def __init__(self, self.store, self.storage_context = self.init_nebula_store() except Exception: print( - f'please use NebulaClient().create() to create space {self.space_name}!!!\n\n\n' + f"please use NebulaClient().create() to create space {self.space_name}!!!\n\n\n" ) self.graph_schema = self.store.get_schema(refresh=None) @@ -554,8 +574,7 @@ def init_nebula_store(self): rel_prop_names=self.rel_prop_names, tags=self.tags, ) - storage_context = StorageContext.from_defaults( - graph_store=nebula_store) + storage_context = StorageContext.from_defaults(graph_store=nebula_store) return nebula_store, storage_context def upsert_triplet(self, triplet: Tuple[str, str, str]): @@ -574,7 +593,7 @@ def get_triplets(self): return self.client.get_triplets(self.space_name) def get_all_entities(self, triplets_file=""): - entities_file = f'/home/chency/NeutronRAG/external_corpus/all_processed_corpus/entity_data/{self.space_name}_entities.json' + entities_file = f"/home/chency/NeutronRAG/external_corpus/all_processed_corpus/entity_data/{self.space_name}_entities.json" if file_exist(entities_file): print(f"load entities from {entities_file}") @@ -598,18 +617,19 @@ def get_all_entities(self, triplets_file=""): save_response(entities, entities_file) # print(f'triplets: {len(all_triplets)}, entities: {len(entities)}') - print(f'entities: {len(entities)}') + print(f"entities: {len(entities)}") return set(entities) - def process_docs(self, - documents, - triplets_per_chunk=10, - include_embeddings=True, - data_dir='./storage_graph', - extract_fn=None, - cache=True): - + def process_docs( + self, + documents, + triplets_per_chunk=10, + include_embeddings=True, + data_dir="./storage_graph", + extract_fn=None, + cache=True, + ): # TODO: use rebel to extract the kg elements. # filter documents @@ -636,7 +656,8 @@ def process_docs(self, if cache: try: storage_context = StorageContext.from_defaults( - persist_dir=data_dir, graph_store=self.store) + persist_dir=data_dir, graph_store=self.store + ) kg_index = load_index_from_storage( storage_context=storage_context, # service_context=service_context, @@ -676,23 +697,31 @@ def process_docs(self, return kg_index def get_rel_map(self, entities, depth=2, limit=30): - rel_map: Optional[Dict] = self.store.get_rel_map(entities, - depth=depth, - limit=limit) + rel_map: Optional[Dict] = self.store.get_rel_map( + entities, depth=depth, limit=limit + ) return rel_map - def set_retriever(self, graph_traversal_depth=2, llm_env=None, limit=30, max_entities=5, max_synonyms=0): + def set_retriever( + self, + graph_traversal_depth=2, + llm_env=None, + limit=30, + max_entities=5, + max_synonyms=0, + ): self.retriever = KnowledgeGraphRAGRetriever( storage_context=self.storage_context, graph_traversal_depth=graph_traversal_depth, max_entities=max_entities, max_synonyms=max_synonyms, - retriever_mode='keyword', + retriever_mode="keyword", verbose=False, entity_extract_template=llm_env.keyword_extract_prompt_template, synonym_expand_template=llm_env.synonym_expand_prompt_template, # clean_kg_sequences_fn=self.clean_kg_sequences, - max_knowledge_sequence=limit) + max_knowledge_sequence=limit, + ) return self.retriever def get_entities(self, query_str: str) -> List[str]: @@ -700,47 +729,47 @@ def get_entities(self, query_str: str) -> List[str]: return self.retriever._get_entities(query_str) def _get_knowledge_sequence( - self, - entities: List[str]) -> Tuple[List[str], Optional[Dict[Any, Any]]]: + self, entities: List[str] + ) -> Tuple[List[str], Optional[Dict[Any, Any]]]: return self.retriever._get_knowledge_sequence(entities) def _build_nodes( - self, - knowledge_sequence: List[str], - rel_map: Optional[Dict[Any, Any]] = None) -> List[NodeWithScore]: - + self, knowledge_sequence: List[str], rel_map: Optional[Dict[Any, Any]] = None + ) -> List[NodeWithScore]: return self.retriever._build_nodes(knowledge_sequence, rel_map) def get_knowledge_sequence(self, rel_map): knowledge_sequence = [] if rel_map: - knowledge_sequence.extend([ - str(rel_obj) for rel_objs in rel_map.values() - for rel_obj in rel_objs - ]) + knowledge_sequence.extend( + [str(rel_obj) for rel_objs in rel_map.values() for rel_obj in rel_objs] + ) else: print("> No knowledge sequence extracted from entities.") return [] return knowledge_sequence - def clean_sequence(self, - sequence, - name_pattern=r'(?<=\{name: )([^{}]+)(?=\})', - edge_pattern=r'(?<=\{relationship: )([^{}]+)(?=\})'): - ''' + def clean_sequence( + self, + sequence, + name_pattern=r"(?<=\{name: )([^{}]+)(?=\})", + edge_pattern=r"(?<=\{relationship: )([^{}]+)(?=\})", + ): + """ kg result: 'James{name: James} -[relationship:{relationship: Joined}]-> Michael jordan{name: Michael jordan}' clean the kg result above to James -Joined-> Michael jordan - ''' + """ names = re.findall(name_pattern, sequence) edges = re.findall(edge_pattern, sequence) - assert len(names) == sequence.count('{name:') - assert len(edges) == sequence.count('{relationship:') + assert len(names) == sequence.count("{name:") + assert len(edges) == sequence.count("{relationship:") for name in names: - sequence = sequence.replace(f'{{name: {name}}}', '') + sequence = sequence.replace(f"{{name: {name}}}", "") for edge in edges: sequence = sequence.replace( - f'[relationship:{{relationship: {edge}}}]', f'{edge}') + f"[relationship:{{relationship: {edge}}}]", f"{edge}" + ) return sequence def clean_kg_sequences(self, knowledge_sequence): @@ -750,19 +779,16 @@ def clean_kg_sequences(self, knowledge_sequence): return clean_knowledge_sequence def clean_rel_map(self, rel_map): - name_pattern = r'(?<=\{name: )([^{}]+)(?=\})' + name_pattern = r"(?<=\{name: )([^{}]+)(?=\})" clean_rel_map = {} for entity, sequences in rel_map.items(): name = re.findall(name_pattern, entity)[0] - clean_ent = entity.replace(f'{{name: {name}}}', '') + clean_ent = entity.replace(f"{{name: {name}}}", "") clean_seq = [self.clean_sequence(seq) for seq in sequences] clean_rel_map[clean_ent] = clean_seq return clean_rel_map - def build_nodes(self, - rel_map, - knowledge_sequence, - depth=2) -> List[NodeWithScore]: + def build_nodes(self, rel_map, knowledge_sequence, depth=2) -> List[NodeWithScore]: """Build nodes from knowledge sequence.""" new_line_char = "\n" context_string = ( @@ -772,7 +798,8 @@ def build_nodes(self, f"`subject -[predicate]->, object, <-[predicate_next_hop]-," f" object_next_hop ...`" f" extracted based on key entities as subject:\n" - f"{new_line_char.join(knowledge_sequence)}") + f"{new_line_char.join(knowledge_sequence)}" + ) if self.verbose: print_text(f"Graph RAG context:\n{context_string}\n", color="blue") @@ -785,13 +812,15 @@ def build_nodes(self, if self.graph_schema != "": rel_node_info["kg_schema"] = {"schema": self.graph_schema} metadata_keys.append("kg_schema") - node = NodeWithScore(node=TextNode( - text=context_string, - score=1.0, - metadata=rel_node_info, - excluded_embed_metadata_keys=metadata_keys, - excluded_llm_metadata_keys=metadata_keys, - )) + node = NodeWithScore( + node=TextNode( + text=context_string, + score=1.0, + metadata=rel_node_info, + excluded_embed_metadata_keys=metadata_keys, + excluded_llm_metadata_keys=metadata_keys, + ) + ) return [node] def drop(self): @@ -821,37 +850,33 @@ def execute(self, query): def two_hop_parse_triplets(self, query): # 定义正则表达式模式 - two_hop_pattern1 = re.compile(r'(.+?) <-(.+?)- (.+?) -(.+?)-> (.+)') - two_hop_pattern2 = re.compile(r'(.+?) <-(.+?)- (.+?) <-(.+?)- (.+)') - two_hop_pattern3 = re.compile(r'(.+?) -(.+?)-> (.+?) -(.+?)-> (.+)') - two_hop_pattern4 = re.compile(r'(.+?) -(.+?)-> (.+?) <-(.+?)- (.+)') + two_hop_pattern1 = re.compile(r"(.+?) <-(.+?)- (.+?) -(.+?)-> (.+)") + two_hop_pattern2 = re.compile(r"(.+?) <-(.+?)- (.+?) <-(.+?)- (.+)") + two_hop_pattern3 = re.compile(r"(.+?) -(.+?)-> (.+?) -(.+?)-> (.+)") + two_hop_pattern4 = re.compile(r"(.+?) -(.+?)-> (.+?) <-(.+?)- (.+)") - one_hop_pattern5 = re.compile(r'(.+?) -(.+?)-> (.+)') - one_hop_pattern6 = re.compile(r'(.+?) <-(.+?)- (.+)') + one_hop_pattern5 = re.compile(r"(.+?) -(.+?)-> (.+)") + one_hop_pattern6 = re.compile(r"(.+?) <-(.+?)- (.+)") match = two_hop_pattern1.match(query) if match: entity1, relation1, entity2, relation2, entity3 = match.groups() - return [(entity2, relation1, entity1), - (entity2, relation2, entity3)] + return [(entity2, relation1, entity1), (entity2, relation2, entity3)] match = two_hop_pattern2.match(query) if match: entity1, relation1, entity2, relation2, entity3 = match.groups() - return [(entity2, relation1, entity1), - (entity3, relation2, entity2)] + return [(entity2, relation1, entity1), (entity3, relation2, entity2)] match = two_hop_pattern3.match(query) if match: entity1, relation1, entity2, relation2, entity3 = match.groups() - return [(entity1, relation1, entity2), - (entity2, relation2, entity3)] + return [(entity1, relation1, entity2), (entity2, relation2, entity3)] match = two_hop_pattern4.match(query) if match: entity1, relation1, entity2, relation2, entity3 = match.groups() - return [(entity1, relation1, entity2), - (entity3, relation2, entity2)] + return [(entity1, relation1, entity2), (entity3, relation2, entity2)] match = one_hop_pattern5.match(query) if match: @@ -890,61 +915,60 @@ def test_show_triplets(db: NebulaDB): def test_rel_map(db: NebulaDB): - entities = ['Lebron james', 'James'] - entities = ['Lebron james'] - entities = ['Lakers', 'James'] - entities = ['James', 'Lakers'] + entities = ["Lebron james", "James"] + entities = ["Lebron james"] + entities = ["Lakers", "James"] + entities = ["James", "Lakers"] rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) - print_text(f"rel_map: {rel_map}\n", color='yellow') - print_text(f"rel_map has {len(rel_map)} entities, {rel_map.keys()}\n", - color='red') + print_text(f"rel_map: {rel_map}\n", color="yellow") + print_text(f"rel_map has {len(rel_map)} entities, {rel_map.keys()}\n", color="red") def test_clean(db: NebulaDB): - entities = ['James', 'Lakers'] + entities = ["James", "Lakers"] rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) - print_text(f"rel_map: {rel_map}\n", color='yellow') + print_text(f"rel_map: {rel_map}\n", color="yellow") knowledge_sequence = db.get_knowledge_sequence(rel_map) - print_text(f"\nknowledge_sequence: {knowledge_sequence}\n", color='yellow') + print_text(f"\nknowledge_sequence: {knowledge_sequence}\n", color="yellow") clean_knowledge_sequence = db.clean_kg_sequences(knowledge_sequence) - print_text(f"\nclean_knowledge_sequence: {clean_knowledge_sequence}\n", - color='yellow') + print_text( + f"\nclean_knowledge_sequence: {clean_knowledge_sequence}\n", color="yellow" + ) clean_rel_map = db.clean_rel_map(rel_map) - print_text(f"\nclean_rel_map: {clean_rel_map}\n", color='yellow') + print_text(f"\nclean_rel_map: {clean_rel_map}\n", color="yellow") def test_build_nodes(db: NebulaDB): - entities = ['Lakers'] + entities = ["Lakers"] rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) rel_map1 = db.clean_rel_map(rel_map) kg_seq1 = [seq for _, seqs in rel_map1.items() for seq in seqs] - print_text(f"\nkg_seq1: {kg_seq1}\n", color='yellow') + print_text(f"\nkg_seq1: {kg_seq1}\n", color="yellow") - entities = ['James'] + entities = ["James"] rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) rel_map2 = db.clean_rel_map(rel_map) kg_seq2 = db.get_knowledge_sequence(rel_map) kg_seq2 = db.clean_kg_sequences(kg_seq2) - print_text(f"\nkg_seq2: {kg_seq2}\n", color='yellow') + print_text(f"\nkg_seq2: {kg_seq2}\n", color="yellow") assert len(set(rel_map1.keys()) & set(rel_map2.keys())) == 0 all_rel_map = {**rel_map1, **rel_map2} all_kg_seq = kg_seq1 + kg_seq2 - print(f'all_rel_map: {len(all_rel_map)}') - print(f'all_kg_seq: {len(all_kg_seq)}') + print(f"all_rel_map: {len(all_rel_map)}") + print(f"all_kg_seq: {len(all_kg_seq)}") # print_text(f"\nall_rel_map: {all_rel_map}\n", color='yellow') # print_text(f"\nall_kg_seq: {all_kg_seq}\n", color='yellow') nodes = db.build_nodes(all_rel_map, all_kg_seq) - print_text(f"\nnodes: {nodes}\n", color='blue') + print_text(f"\nnodes: {nodes}\n", color="blue") def test_parse_2_hop_rel(db: NebulaDB): - # if not questions: # questions = """Elon musk <-Led by- Team -Play-> Next time # Elon musk <-Reports to- Linda yaccarino -Is-> The new ceo of twitter @@ -957,13 +981,13 @@ def test_parse_2_hop_rel(db: NebulaDB): # Shaw -Was tempted by-> Fascism <-Suffered under- Pinocchio # """ - entities = ['Twitter', 'Elon musk'] + entities = ["Twitter", "Elon musk"] rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) - print_text(f"\nrel_map: {rel_map}\n", color='green') + print_text(f"\nrel_map: {rel_map}\n", color="green") clean_map = db.clean_rel_map(rel_map) - print_text(f"\nclean_map: {clean_map}\n", color='yellow') + print_text(f"\nclean_map: {clean_map}\n", color="yellow") knowledge_sequence = db.get_knowledge_sequence(clean_map) - print_text(f"\knowledge_sequence: {knowledge_sequence}\n", color='blue') + print_text(f"\knowledge_sequence: {knowledge_sequence}\n", color="blue") # for keyword, rels in clean_map.items(): # print('###', keyword) # for rel in rels: @@ -990,7 +1014,7 @@ def query_time( # start_time = time.time() rel_map = nebula_db.get_rel_map(en, limit=limit) for k, v in rel_map.items(): - print('############', k) + print("############", k) for rel in v: print(rel) print() @@ -999,7 +1023,7 @@ def query_time( # print(clean_rel_map) for k, v in clean_rel_map.items(): - print('############', k) + print("############", k) for rel in v: print(rel) @@ -1016,7 +1040,6 @@ def query_time( def test_query_time(nebula_db: NebulaDB): - # ret = nebula_db.execute(""" # # MATCH (talent:entity {name: 'Talent'})-[:relationship*2]-(neighbor) # # RETURN DISTINCT neighbor @@ -1037,12 +1060,22 @@ def test_query_time(nebula_db: NebulaDB): # ] q1_entity = [ - 'Top two contestants', 'Finale', "America's Got Talent", 'Talent', - 'Top', 'Contestants', 'Season 17' + "Top two contestants", + "Finale", + "America's Got Talent", + "Talent", + "Top", + "Contestants", + "Season 17", ] q1_entity = [ - 'Top two contestants', 'Finale', "America's Got Talent", 'Top', - 'Contestants', 'Season 17', 'Talent' + "Top two contestants", + "Finale", + "America's Got Talent", + "Top", + "Contestants", + "Season 17", + "Talent", ] query_time(nebula_db, [q1_entity], limit=30) @@ -1063,11 +1096,18 @@ def test_simple_pruning(nebula_db: NebulaDB): # ] q1_entity = [ - 'Top two contestants', 'Finale', "America's Got Talent", 'Talent', - 'Top', 'Contestants', 'Season 17' + "Top two contestants", + "Finale", + "America's Got Talent", + "Talent", + "Top", + "Contestants", + "Season 17", ] - question = "Who were the top two contestants in the America's Got Talent Season 17 finale?" + question = ( + "Who were the top two contestants in the America's Got Talent Season 17 finale?" + ) # q1_entity = ['Wallis', 'Susie', 'Susie wallis', 'Film', 'Actress', 'Movie star', 'Motion picture', 'Flick', 'Movie', 'Searches', 'Star', 'Susie searches', 'Picture', 'Motion'] @@ -1096,25 +1136,23 @@ def test_simple_pruning(nebula_db: NebulaDB): def test_ppr(db: NebulaDB): - entities = ['Twitter', 'Elon musk'] - question = 'What relation between Twitter and Elon musk' + entities = ["Twitter", "Elon musk"] + question = "What relation between Twitter and Elon musk" # rel_map = db.get_rel_map(entities=entities, depth=2, limit=30) rel_map = db.get_rel_map(entities=entities, depth=2, limit=30000) clean_map = db.clean_rel_map(rel_map) - print_text(f"\nclean_map: {clean_map}\n", color='yellow') + print_text(f"\nclean_map: {clean_map}\n", color="yellow") all_rels = [] for rels in clean_map.values(): all_rels += rels - print('relations:', len(all_rels)) + print("relations:", len(all_rels)) triplets, rel_to_entities = db.two_hop_parse_multi_triplets(all_rels) - filter_rels, filter_triplets = filter_pr_rels(question, - entities, - triplets, - rel_to_entities, - max_ent=5) + filter_rels, filter_triplets = filter_pr_rels( + question, entities, triplets, rel_to_entities, max_ent=5 + ) # for rel in filter_rels: # print(rel) @@ -1127,9 +1165,7 @@ def test_ppr(db: NebulaDB): return filter_rels, filter_triplets - -if __name__ == '__main__': - +if __name__ == "__main__": # space_name = 'integrationrgb' client = NebulaClient() client.show_space() @@ -1149,7 +1185,7 @@ def test_ppr(db: NebulaDB): client.show_space() # client.clear('rgb') - space_name = 'multihop_ccy' + space_name = "multihop_ccy" # space_name = 'kelm_1m' # space_name = 'rgb_llama2_70b' # space_name = 'newrgb' @@ -1160,7 +1196,7 @@ def test_ppr(db: NebulaDB): # db.count_edges() # db.show_space() - space_name = 'multihop_ccy' + space_name = "multihop_ccy" # client.drop_space('hotpotqa') # client.show_space() # db.show_space() @@ -1180,4 +1216,3 @@ def test_ppr(db: NebulaDB): print(f"vertices number: {len(vertices)}, edge number: {len(edges)}") print(vertices[0]) print(edges[0]) - \ No newline at end of file diff --git a/aag/rag_engine/graph_query/graph_query.py b/aag/rag_engine/graph_query/graph_query.py index 9116d8b..777c4e1 100644 --- a/aag/rag_engine/graph_query/graph_query.py +++ b/aag/rag_engine/graph_query/graph_query.py @@ -5,11 +5,13 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple from neo4j import GraphDatabase from neo4j.exceptions import ClientError, DatabaseError, TransientError +import os import re import time JsonDict = Dict[str, Any] + @dataclass class Neo4jConfig: uri: str @@ -33,10 +35,33 @@ class Neo4jGraphClient: # Neo4j reserved keywords (partial list) RESERVED_KEYWORDS = { - 'MATCH', 'RETURN', 'WHERE', 'CREATE', 'DELETE', 'SET', - 'MERGE', 'WITH', 'UNWIND', 'CASE', 'WHEN', 'THEN', - 'ELSE', 'END', 'ORDER', 'BY', 'SKIP', 'LIMIT', 'AS', - 'AND', 'OR', 'NOT', 'IN', 'IS', 'NULL', 'TRUE', 'FALSE' + "MATCH", + "RETURN", + "WHERE", + "CREATE", + "DELETE", + "SET", + "MERGE", + "WITH", + "UNWIND", + "CASE", + "WHEN", + "THEN", + "ELSE", + "END", + "ORDER", + "BY", + "SKIP", + "LIMIT", + "AS", + "AND", + "OR", + "NOT", + "IN", + "IS", + "NULL", + "TRUE", + "FALSE", } def __init__(self, config: Neo4jConfig): @@ -47,11 +72,11 @@ def __init__(self, config: Neo4jConfig): config: Neo4j connection settings. """ self._driver = GraphDatabase.driver( - config.uri, + config.uri, auth=(config.user, config.password), max_connection_lifetime=3600, # Max connection lifetime: 1 hour - max_connection_pool_size=50, # Connection pool size - connection_acquisition_timeout=60.0 # Connection acquisition timeout (seconds) + max_connection_pool_size=50, # Connection pool size + connection_acquisition_timeout=60.0, # Connection acquisition timeout (seconds) ) self._db = config.database @@ -76,7 +101,7 @@ def run( *, read: bool = True, max_retries: int = 3, - show_query: bool = True + show_query: bool = True, ) -> List[JsonDict]: """ Run a Cypher query with automatic retries on transient failures. @@ -95,11 +120,11 @@ def run( RuntimeError: Query failed after retries or on non-transient errors. """ params = params or {} - + # Optionally print filled query (debug) if show_query: self._print_filled_query(cypher, params) - + # Retry loop for attempt in range(max_retries): try: @@ -112,7 +137,7 @@ def run( return session.write_transaction( lambda tx: [r.data() for r in tx.run(cypher, params)] ) - + except TransientError as e: # Transient error: retry if attempt == max_retries - 1: @@ -120,17 +145,17 @@ def run( f"Neo4j TransientError after {max_retries} retries: {e}\n" f"Cypher: {cypher}\nParams: {params}" ) from e - - wait_time = 2 ** attempt # Exponential backoff - print(f"⚠️ TransientError, retrying in {wait_time}s ({attempt + 1}/{max_retries})...") + + wait_time = 2**attempt # Exponential backoff + print( + f"⚠️ TransientError, retrying in {wait_time}s ({attempt + 1}/{max_retries})..." + ) time.sleep(wait_time) - + except (ClientError, DatabaseError) as e: # Client or database error: do not retry raise RuntimeError( - f"Neo4j Error: {e}\n" - f"Cypher: {cypher}\n" - f"Params: {params}" + f"Neo4j Error: {e}\nCypher: {cypher}\nParams: {params}" ) from e def _print_filled_query(self, cypher: str, params: Dict) -> None: @@ -139,28 +164,33 @@ def _print_filled_query(self, cypher: str, params: Dict) -> None: Note: Display only; execution still uses parameterized queries via the driver. """ - print("\n" + "="*80) + print("\n" + "=" * 80) print("📝 执行的 Cypher 查询语句（参数已填充）:") - print("-"*80) - + print("-" * 80) + filled_cypher = cypher if params: import json + # Sort param names by length descending to avoid partial $id vs $id2 replacements - sorted_params = sorted(params.items(), key=lambda x: len(x[0]), reverse=True) - + sorted_params = sorted( + params.items(), key=lambda x: len(x[0]), reverse=True + ) + for key, value in sorted_params: # CRITICAL: Skip tuple params (DSL internal form; must not appear in final Cypher) # Tuple form like (">", 500000) should already be expanded in WHERE building if isinstance(value, tuple): # Should not happen; indicates a bug if it does - print(f"⚠️ WARNING: 参数 ${key} 是元组格式 {value}，这不应该出现在最终查询中！") + print( + f"⚠️ WARNING: 参数 ${key} 是元组格式 {value}，这不应该出现在最终查询中！" + ) continue - + # Format value by type if isinstance(value, str): # Escape single quotes - formatted_value = f"'{value.replace(chr(39), chr(39)+chr(39))}'" + formatted_value = f"'{value.replace(chr(39), chr(39) + chr(39))}'" elif isinstance(value, (int, float)): formatted_value = str(value) elif isinstance(value, bool): @@ -176,22 +206,22 @@ def _print_filled_query(self, cypher: str, params: Dict) -> None: formatted_value = json.dumps(value, ensure_ascii=False) else: formatted_value = str(value) - + # Regex for full token match (avoid $id matching inside $id2) filled_cypher = re.sub( - r'\$' + re.escape(key) + r'\b', # Word boundary + r"\$" + re.escape(key) + r"\b", # Word boundary formatted_value, - filled_cypher + filled_cypher, ) - + print(filled_cypher) - print("="*80 + "\n") + print("=" * 80 + "\n") # ===== Schema introspection ===== def get_schema(self) -> Dict: """ Fetch graph schema metadata (extended). - + Returns: { "node_labels": { @@ -209,63 +239,72 @@ def get_schema(self) -> Dict: "patterns": [pattern_strings] } """ - schema = { - "node_labels": {}, - "relationship_types": {}, - "patterns": [] - } - + schema = {"node_labels": {}, "relationship_types": {}, "patterns": []} + # 1. All node labels and properties (with sample values) - labels_result = self.run("CALL db.labels() YIELD label RETURN label", show_query=False) + labels_result = self.run( + "CALL db.labels() YIELD label RETURN label", show_query=False + ) valid_labels = [item["label"] for item in labels_result] - + for label in valid_labels: if not label or not self._is_valid_identifier(label): continue - + # Properties and sample values for this label - props_result = self.run(f""" + props_result = self.run( + f""" MATCH (n:`{label}`) WITH n LIMIT 1 UNWIND keys(n) AS key RETURN key, n[key] AS sample_value ORDER BY key - """, show_query=False) - + """, + show_query=False, + ) + if props_result: schema["node_labels"][label] = { "properties": [p["key"] for p in props_result], - "sample_values": {p["key"]: p["sample_value"] for p in props_result} + "sample_values": { + p["key"]: p["sample_value"] for p in props_result + }, } - + # 2. All relationship types and properties (with sample values) rels_result = self.run( "CALL db.relationshipTypes() YIELD relationshipType RETURN relationshipType", - show_query=False + show_query=False, ) valid_rels = [item["relationshipType"] for item in rels_result] - + for rel_type in valid_rels: if not rel_type or not self._is_valid_identifier(rel_type): continue - + # Properties and sample values for this relationship type - props_result = self.run(f""" + props_result = self.run( + f""" MATCH ()-[r:`{rel_type}`]->() WITH r LIMIT 1 UNWIND keys(r) AS key RETURN key, r[key] AS sample_value ORDER BY key - """, show_query=False) - + """, + show_query=False, + ) + if props_result: schema["relationship_types"][rel_type] = { "properties": [p["key"] for p in props_result], - "sample_values": {p["key"]: p["sample_value"] for p in props_result} + "sample_values": { + p["key"]: p["sample_value"] for p in props_result + }, } - + # 3. Relationship patterns (start_label, rel, end_label) - patterns = self.run(""" + patterns = self.run( + """ MATCH (a)-[r]->(b) WITH labels(a)[0] AS start_label, type(r) AS rel_type, @@ -275,20 +314,22 @@ def get_schema(self) -> Dict: AND end_label IS NOT NULL RETURN DISTINCT start_label, rel_type, end_label LIMIT 100 - """, show_query=False) - + """, + show_query=False, + ) + schema["patterns"] = [ f"({p['start_label']})-[:{p['rel_type']}]->({p['end_label']})" for p in patterns ] - + return schema # ===== Validation helpers ===== @staticmethod def _is_valid_identifier(name: str) -> bool: """Return True if name is a simple alphanumeric/underscore identifier.""" - return bool(re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', name)) + return bool(re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", name)) @classmethod def _sanitize_label(cls, label: Optional[str]) -> str: @@ -303,19 +344,21 @@ def _sanitize_label(cls, label: Optional[str]) -> str: """ if not label: return "" - + # Length check if len(label) > 255: raise ValueError(f"Label too long (max 255): {label}") - + # Format check - if not re.match(r'^[A-Za-z][A-Za-z0-9_]*$', label): - raise ValueError(f"Invalid label format (must start with letter, contain only alphanumeric and underscore): {label}") - + if not re.match(r"^[A-Za-z][A-Za-z0-9_]*$", label): + raise ValueError( + f"Invalid label format (must start with letter, contain only alphanumeric and underscore): {label}" + ) + # Reserved keyword check if label.upper() in cls.RESERVED_KEYWORDS: raise ValueError(f"Reserved keyword cannot be used as label: {label}") - + return label @classmethod @@ -330,17 +373,19 @@ def _sanitize_rel_type(cls, rel_type: Optional[str]) -> str: """ if not rel_type: return "" - + if len(rel_type) > 255: raise ValueError(f"Relationship type too long (max 255): {rel_type}") - + # Format: letter first, then alphanumeric/underscore - if not re.match(r'^[A-Za-z][A-Za-z0-9_]*$', rel_type): + if not re.match(r"^[A-Za-z][A-Za-z0-9_]*$", rel_type): raise ValueError(f"Invalid relationship type format: {rel_type}") - + if rel_type.upper() in cls.RESERVED_KEYWORDS: - raise ValueError(f"Reserved keyword cannot be used as relationship type: {rel_type}") - + raise ValueError( + f"Reserved keyword cannot be used as relationship type: {rel_type}" + ) + return rel_type @staticmethod @@ -355,23 +400,20 @@ def _sanitize_property_key(key: str) -> str: """ if not key: raise ValueError("Property key cannot be empty") - + if len(key) > 255: raise ValueError(f"Property key too long (max 255): {key}") - - if not re.match(r'^[a-zA-Z][a-zA-Z0-9_]*$', key): + + if not re.match(r"^[a-zA-Z][a-zA-Z0-9_]*$", key): raise ValueError(f"Invalid property key format: {key}") - + return key # ========================================================= # 1. Lookup by internal ID / unique business key # ========================================================= def get_node_by_internal_id( - self, - internal_id: int, - *, - return_props: bool = True + self, internal_id: int, *, return_props: bool = True ) -> Optional[JsonDict]: """ Look up a node by Neo4j internal id(n). @@ -387,7 +429,7 @@ def get_node_by_internal_id( """ cypher = "MATCH (n) WHERE id(n) = $id RETURN n AS node" res = self.run(cypher, {"id": internal_id}) - + if not res: return None return res[0]["node"] if return_props else res[0] @@ -398,7 +440,7 @@ def get_node_by_unique_key( key: str, value: Any, *, - return_fields: Optional[List[str]] = None + return_fields: Optional[List[str]] = None, ) -> Optional[JsonDict]: """ Find a node by label and a unique key property. @@ -422,7 +464,7 @@ def get_node_by_unique_key( """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + # Build RETURN clause if return_fields: # Projected fields @@ -434,13 +476,13 @@ def get_node_by_unique_key( else: # Whole node return_clause = "RETURN n AS node" - + cypher = f"MATCH (n:`{label}` {{`{key}`: $value}}) {return_clause} LIMIT 1" res = self.run(cypher, {"value": value}) - + if not res: return None - + # Dict of fields vs. full node if return_fields: return res[0] @@ -456,7 +498,7 @@ def filter_nodes_by_properties( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Filter multiple nodes by property predicates (AND-combined). @@ -522,25 +564,35 @@ def filter_nodes_by_properties( - Conditions are ANDed; for OR use filter_query or custom Cypher. """ label = self._sanitize_label(label) - + if not conditions: raise ValueError("conditions cannot be empty") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + # Build WHERE clause where_parts = [] params = {} for i, (key, value) in enumerate(conditions.items()): key = self._sanitize_property_key(key) param_name = f"cond_{i}" - + # Fix 2: (operator, value) tuple or [operator, value] list - if (isinstance(value, (tuple, list)) and len(value) == 2): + if isinstance(value, (tuple, list)) and len(value) == 2: operator, actual_value = value # Validate operator - valid_operators = ["=", ">", "<", ">=", "<=", "!=", "IN", "CONTAINS", "STARTS WITH"] + valid_operators = [ + "=", + ">", + "<", + ">=", + "<=", + "!=", + "IN", + "CONTAINS", + "STARTS WITH", + ] if operator.upper() in ["IN", "CONTAINS"]: where_parts.append(f"a.`{key}` {operator.upper()} ${param_name}") elif operator.upper() == "STARTS WITH": @@ -554,9 +606,9 @@ def filter_nodes_by_properties( # Equality where_parts.append(f"a.`{key}` = ${param_name}") params[param_name] = value - + where_clause = "WHERE " + " AND ".join(where_parts) - + # Build RETURN clause if return_fields: # Projected fields @@ -568,19 +620,19 @@ def filter_nodes_by_properties( else: # Whole node return_clause = "RETURN a AS node" - + # Optional ORDER BY order_clause = "" if order_by: order_by = self._sanitize_property_key(order_by) order_clause = f"ORDER BY a.`{order_by}` {order_direction}" - + # Optional LIMIT limit_clause = "" if limit is not None: limit_clause = "LIMIT $limit" params["limit"] = limit - + # Full query cypher = f""" MATCH (a:`{label}`) @@ -589,7 +641,7 @@ def filter_nodes_by_properties( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # add gjq @@ -605,7 +657,7 @@ def filter_relationships( aggregate_field: Optional[str] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Filter relationships by relationship properties; optional aggregates. @@ -673,18 +725,22 @@ def filter_relationships( - rel_conditions are ANDed; for OR use filter_query or custom Cypher. """ rel_type = self._sanitize_rel_type(rel_type) - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + # Node patterns - start_pattern = f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" - end_pattern = f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" - + start_pattern = ( + f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" + ) + end_pattern = ( + f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" + ) + # WHERE clause where_parts = [] params = {} - + if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): # Keys like tran_timestamp_start / _end → strip suffix for real prop name @@ -694,33 +750,37 @@ def filter_relationships( actual_key = self._sanitize_property_key(actual_key) else: actual_key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name = f"rel_cond_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"t.`{actual_key}` {operator.upper()} ${param_name}") + where_parts.append( + f"t.`{actual_key}` {operator.upper()} ${param_name}" + ) elif operator.upper() == "STARTS WITH": - where_parts.append(f"t.`{actual_key}` STARTS WITH ${param_name}") + where_parts.append( + f"t.`{actual_key}` STARTS WITH ${param_name}" + ) else: where_parts.append(f"t.`{actual_key}` {operator} ${param_name}") - + params[param_name] = value else: # Equality param_name = f"rel_cond_{i}" where_parts.append(f"t.`{actual_key}` = ${param_name}") params[param_name] = condition - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # Aggregate branch if aggregate: agg_type = aggregate.upper() if agg_type not in ["COUNT", "SUM", "AVG", "MAX", "MIN"]: raise ValueError(f"Invalid aggregate type: {aggregate}") - + if agg_type == "COUNT": agg_expr = "COUNT(t)" else: @@ -728,18 +788,20 @@ def filter_relationships( raise ValueError(f"aggregate_field is required for {agg_type}") agg_field = self._sanitize_property_key(aggregate_field) agg_expr = f"{agg_type}(t.`{agg_field}`)" - + cypher = f""" MATCH {start_pattern}-[t:`{rel_type}`]->{end_pattern} {where_clause} RETURN {agg_expr} AS value """ - + result = self.run(cypher, params) if result: - return [{"aggregate_type": agg_type.lower(), "value": result[0]["value"]}] + return [ + {"aggregate_type": agg_type.lower(), "value": result[0]["value"]} + ] return [{"aggregate_type": agg_type.lower(), "value": 0}] - + # Row-returning query if return_fields: # RETURN list @@ -761,24 +823,24 @@ def filter_relationships( # Relationship property (bare name) prop = self._sanitize_property_key(field) return_parts.append(f"t.`{prop}` AS {prop}") - + return_clause = "RETURN " + ", ".join(return_parts) else: # Full from / rel / to return_clause = "RETURN from, t AS relationship, to" - + # Optional ORDER BY order_clause = "" if order_by: order_by = self._sanitize_property_key(order_by) order_clause = f"ORDER BY t.`{order_by}` {order_direction}" - + # Optional LIMIT limit_clause = "" if limit is not None: limit_clause = "LIMIT $limit" params["limit"] = limit - + # Full Cypher cypher = f""" MATCH {start_pattern}-[t:`{rel_type}`]->{end_pattern} @@ -787,7 +849,7 @@ def filter_relationships( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # add gjq @@ -805,7 +867,7 @@ def aggregation_query( where: Optional[str] = None, order_by: Optional[str] = None, order_direction: str = "DESC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Aggregations with grouping by node or by a node property. @@ -875,12 +937,12 @@ def aggregation_query( agg_type = aggregate_type.upper() if agg_type not in ["COUNT", "SUM", "AVG", "MAX", "MIN"]: raise ValueError(f"Invalid aggregate type: {aggregate_type}") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + params = {} - + # Aggregate expression if agg_type == "COUNT": if rel_type: @@ -898,18 +960,18 @@ def aggregation_query( else: agg_expr = f"{agg_type}(n.`{agg_field}`) AS total" agg_alias = "total" - + # Case 1: group by property on nodes only (no rel_type) if group_by_property and not rel_type: label = self._sanitize_label(node_label) if node_label else "" label_pattern = f":`{label}`" if label else "" prop = self._sanitize_property_key(group_by_property) - + where_clause = f"WHERE {where}" if where else "" limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH (n{label_pattern}) {where_clause} @@ -917,13 +979,13 @@ def aggregation_query( ORDER BY {order_by or agg_alias} {order_direction} {limit_clause} """ - + # Case 2: group by node + relationship aggregation elif group_by_node and rel_type: label = self._sanitize_label(node_label) if node_label else "" label_pattern = f":`{label}`" if label else "" rt = self._sanitize_rel_type(rel_type) - + # Relationship pattern if direction == "out": if group_by_node == "start": @@ -937,7 +999,7 @@ def aggregation_query( pattern = f"()<-[r:`{rt}`]-(n{label_pattern})" else: pattern = f"(n{label_pattern})-[r:`{rt}`]-()" - + # RETURN parts return_parts = [] if return_fields: @@ -946,14 +1008,14 @@ def aggregation_query( return_parts.append(f"n.`{field}` AS {field}") else: return_parts.append("n.node_key AS node_key") - + return_clause = ", ".join(return_parts) + f", {agg_expr}" - + where_clause = f"WHERE {where}" if where else "" limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH {pattern} {where_clause} @@ -961,10 +1023,12 @@ def aggregation_query( ORDER BY {order_by or agg_alias} {order_direction} {limit_clause} """ - + else: - raise ValueError("Must specify either group_by_property or (group_by_node + rel_type)") - + raise ValueError( + "Must specify either group_by_property or (group_by_node + rel_type)" + ) + return self.run(cypher, params) # ========================================================= @@ -987,7 +1051,7 @@ def neighbors_n_hop( limit: Optional[int] = None, return_distinct: bool = False, exclude_start: bool = False, - return_path_length: bool = False + return_path_length: bool = False, ) -> List[JsonDict]: """ N-hop neighbors from a start node (filters, projection, path stats). @@ -1058,19 +1122,19 @@ def neighbors_n_hop( """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + if not (1 <= hops <= 10): raise ValueError("hops must be between 1 and 10") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # Fast path: single hop + projection + no path length → skip path aggregation if hops == 1 and return_fields and not return_path_length: # One-hop pattern @@ -1080,7 +1144,7 @@ def neighbors_n_hop( pattern = f"(start)<-[r{rel}]-(nbr)" else: pattern = f"(start)-[r{rel}]-(nbr)" - + # Optional WHERE where_parts = [] if where: @@ -1088,7 +1152,7 @@ def neighbors_n_hop( if exclude_start: where_parts.append("nbr <> start") where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # RETURN return_parts = [] for field in return_fields: @@ -1104,15 +1168,19 @@ def neighbors_n_hop( # Treat bare name as neighbor property prop = self._sanitize_property_key(field) return_parts.append(f"nbr.`{prop}` AS {prop}") - - return_clause = "RETURN " + (" DISTINCT " if return_distinct else " ") + ", ".join(return_parts) - + + return_clause = ( + "RETURN " + + (" DISTINCT " if return_distinct else " ") + + ", ".join(return_parts) + ) + # Optional ORDER BY order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (start:`{label}` {{`{key}`: $value}}) MATCH {pattern} @@ -1130,7 +1198,7 @@ def neighbors_n_hop( pattern = f"(start)<-[r{rel}*1..{hops}]-(nbr)" else: pattern = f"(start)-[r{rel}*1..{hops}]-(nbr)" - + # Optional WHERE where_parts = [] if where: @@ -1138,7 +1206,7 @@ def neighbors_n_hop( if exclude_start: where_parts.append("nbr <> start") where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # RETURN if return_path_length: # Include hop count as path_length @@ -1160,7 +1228,7 @@ def neighbors_n_hop( firstRel AS rel, type(firstRel) AS relType """ - + # Optional ORDER BY if order_by: if order_by == "path_length" and return_path_length: @@ -1169,10 +1237,10 @@ def neighbors_n_hop( order_clause = f"ORDER BY {order_by} {order_direction}" else: order_clause = "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (start:`{label}` {{`{key}`: $value}}) MATCH p = {pattern} @@ -1183,11 +1251,11 @@ def neighbors_n_hop( {order_clause} {limit_clause} """ - + params = {"value": value} if limit is not None: params["limit"] = limit - + return self.run(cypher, params) # ========================================================= @@ -1204,7 +1272,7 @@ def common_neighbors( order_by: Optional[str] = None, order_direction: str = "ASC", limit: Optional[int] = None, - aggregate: bool = False + aggregate: bool = False, ) -> List[JsonDict]: """ One-hop common neighbors of two nodes; optional count-based ordering. @@ -1229,21 +1297,21 @@ def common_neighbors( """ (la, ka, va) = a (lb, kb, vb) = b - + la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # Path patterns A–C and B–C if direction == "out": pat_a = f"(A)-[rA{rel}]->(C)" @@ -1254,20 +1322,24 @@ def common_neighbors( else: pat_a = f"(A)-[rA{rel}]-(C)" pat_b = f"(B)-[rB{rel}]-(C)" - + # Optional WHERE where_clause = f"WHERE {where}" if where else "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + # Aggregate: COUNT per common neighbor if aggregate: if order_by == "count": order_clause = f"ORDER BY count {order_direction}" else: - order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "ORDER BY count DESC" - + order_clause = ( + f"ORDER BY {order_by} {order_direction}" + if order_by + else "ORDER BY count DESC" + ) + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1281,7 +1353,7 @@ def common_neighbors( else: # Full relA / relB per pair order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1292,11 +1364,11 @@ def common_neighbors( {order_clause} {limit_clause} """ - + params = {"va": va, "vb": vb} if limit is not None: params["limit"] = limit - + return self.run(cypher, params) def common_neighbors_with_rel_filter( @@ -1311,7 +1383,7 @@ def common_neighbors_with_rel_filter( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Common neighbors with identical predicates on edges A–C and B–C. @@ -1367,21 +1439,21 @@ def common_neighbors_with_rel_filter( """ (la, ka, va) = a (lb, kb, vb) = b - + la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # A–C / B–C patterns if direction == "out": pat_a = f"(A)-[rA{rel}]->(C)" @@ -1392,31 +1464,35 @@ def common_neighbors_with_rel_filter( else: pat_a = f"(A)-[rA{rel}]-(C)" pat_b = f"(B)-[rB{rel}]-(C)" - + # WHERE where_parts = [] params = {"va": va, "vb": vb} - + # Duplicate rel_conditions onto rA and rB if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name_a = f"rel_cond_a_{i}" param_name_b = f"rel_cond_b_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"rA.`{key}` {operator.upper()} ${param_name_a}") - where_parts.append(f"rB.`{key}` {operator.upper()} ${param_name_b}") + where_parts.append( + f"rA.`{key}` {operator.upper()} ${param_name_a}" + ) + where_parts.append( + f"rB.`{key}` {operator.upper()} ${param_name_b}" + ) elif operator.upper() == "STARTS WITH": where_parts.append(f"rA.`{key}` STARTS WITH ${param_name_a}") where_parts.append(f"rB.`{key}` STARTS WITH ${param_name_b}") else: where_parts.append(f"rA.`{key}` {operator} ${param_name_a}") where_parts.append(f"rB.`{key}` {operator} ${param_name_b}") - + params[param_name_a] = value params[param_name_b] = value else: @@ -1427,13 +1503,13 @@ def common_neighbors_with_rel_filter( where_parts.append(f"rB.`{key}` = ${param_name_b}") params[param_name_a] = condition params[param_name_b] = condition - + # Predicate on C if neighbor_where: where_parts.append(f"({neighbor_where})") - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # RETURN if return_fields: return_parts = [] @@ -1452,15 +1528,15 @@ def common_neighbors_with_rel_filter( return_clause = "RETURN " + ", ".join(return_parts) else: return_clause = "RETURN C AS commonNeighbor, rA AS relA, rB AS relB" - + # Optional ORDER BY order_clause = f"ORDER BY {order_by} {order_direction}" if order_by else "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}) MATCH (B:`{lb}` {{`{kb}`: $vb}}) @@ -1471,7 +1547,7 @@ def common_neighbors_with_rel_filter( {order_clause} {limit_clause} """ - + return self.run(cypher, params) # ========================================================= @@ -1487,7 +1563,7 @@ def filter_query( node_where: Optional[str] = None, rel_where: Optional[str] = None, params: Optional[JsonDict] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ One hop from `start` with optional rel/node predicates. @@ -1508,16 +1584,16 @@ def filter_query( (sl, sk, sv) = start sl = self._sanitize_label(sl) sk = self._sanitize_property_key(sk) - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + tl = self._sanitize_label(node_label) if node_label else "" tlabel = f":`{tl}`" if tl else "" - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # s–(r)–n pattern if direction == "out": pat = f"(s)-[r{rel}]->(n{tlabel})" @@ -1525,7 +1601,7 @@ def filter_query( pat = f"(s)<-[r{rel}]-(n{tlabel})" else: pat = f"(s)-[r{rel}]-(n{tlabel})" - + # WHERE fragments where_parts = [] if rel_where: @@ -1533,10 +1609,10 @@ def filter_query( if node_where: where_parts.append(f"({node_where})") where_clause = ("WHERE " + " AND ".join(where_parts)) if where_parts else "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (s:`{sl}` {{`{sk}`: $sv}}) MATCH {pat} @@ -1544,13 +1620,13 @@ def filter_query( RETURN s AS start, r AS rel, n AS node {limit_clause} """ - + p = {"sv": sv} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) # ========================================================= @@ -1564,7 +1640,7 @@ def subgraph_extract( rel_type: Optional[str] = None, direction: str = "both", where: Optional[str] = None, - limit_paths: int = 200 + limit_paths: int = 200, ) -> JsonDict: """ Ego network around `center` up to `hops` (distinct nodes and rels from paths). @@ -1582,13 +1658,13 @@ def subgraph_extract( (cl, ck, cv) = center cl = self._sanitize_label(cl) ck = self._sanitize_property_key(ck) - + if not (1 <= hops <= 5): raise ValueError("hops must be between 1 and 5") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # Variable-length from center if direction == "out": pat = f"(c)-[r{rel}*1..{hops}]->(n)" @@ -1596,10 +1672,10 @@ def subgraph_extract( pat = f"(c)<-[r{rel}*1..{hops}]-(n)" else: pat = f"(c)-[r{rel}*1..{hops}]-(n)" - + # Optional WHERE where_clause = f"WHERE {where}" if where else "" - + cypher = f""" MATCH (c:`{cl}` {{`{ck}`: $cv}}) MATCH p = {pat} @@ -1611,15 +1687,15 @@ def subgraph_extract( WITH collect(DISTINCT nn) AS nodes, collect(DISTINCT rr) AS relationships RETURN nodes, relationships """ - + res = self.run(cypher, {"cv": cv, "limit_paths": limit_paths}) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), - "relationships": res[0].get("relationships", []) + "relationships": res[0].get("relationships", []), } - + return {"nodes": [], "relationships": []} def subgraph_extract_by_nodes( @@ -1631,7 +1707,7 @@ def subgraph_extract_by_nodes( include_internal: bool = True, rel_type: Optional[str] = None, direction: str = "both", - where: Optional[str] = None + where: Optional[str] = None, ) -> JsonDict: """ Induced subgraph on a set of nodes keyed by (label, key, values). @@ -1676,17 +1752,17 @@ def subgraph_extract_by_nodes( """ label = self._sanitize_label(label) key = self._sanitize_property_key(key) - + if not values or len(values) == 0: raise ValueError("values list cannot be empty") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # Relationship type token rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # n1-(r)-n2 pattern if direction == "out": pattern = f"(n1)-[r{rel}]->(n2)" @@ -1694,21 +1770,21 @@ def subgraph_extract_by_nodes( pattern = f"(n1)<-[r{rel}]-(n2)" else: pattern = f"(n1)-[r{rel}]-(n2)" - + # Both endpoints in values where_parts = [] - + where_parts.append("n1.`" + key + "` IN $values") where_parts.append("n2.`" + key + "` IN $values") - + if not include_internal: where_parts.append("n1 <> n2") - + if where: where_parts.append(f"({where})") - + where_clause = "WHERE " + " AND ".join(where_parts) - + cypher = f""" MATCH (n1:`{label}`) WHERE n1.`{key}` IN $values @@ -1723,22 +1799,22 @@ def subgraph_extract_by_nodes( size(allNodes) AS node_count, size(allRels) AS relationship_count """ - + res = self.run(cypher, {"values": values}) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), "relationships": res[0].get("relationships", []), "node_count": res[0].get("node_count", 0), - "relationship_count": res[0].get("relationship_count", 0) + "relationship_count": res[0].get("relationship_count", 0), } - + return { "nodes": [], "relationships": [], "node_count": 0, - "relationship_count": 0 + "relationship_count": 0, } def subgraph_extract_by_rel_filter( @@ -1749,7 +1825,7 @@ def subgraph_extract_by_rel_filter( start_label: Optional[str] = None, end_label: Optional[str] = None, direction: str = "both", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> JsonDict: """ Subgraph from relationships matching property predicates (endpoints + edges). @@ -1796,14 +1872,18 @@ def subgraph_extract_by_rel_filter( - relationship_count reflects collected edges after LIMIT. """ rel_type = self._sanitize_rel_type(rel_type) - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + # Endpoint patterns - start_pattern = f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" - end_pattern = f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" - + start_pattern = ( + f"(from:`{self._sanitize_label(start_label)}`)" if start_label else "(from)" + ) + end_pattern = ( + f"(to:`{self._sanitize_label(end_label)}`)" if end_label else "(to)" + ) + # Directional relationship pattern if direction == "out": rel_pattern = f"{start_pattern}-[r:`{rel_type}`]->{end_pattern}" @@ -1811,11 +1891,11 @@ def subgraph_extract_by_rel_filter( rel_pattern = f"{start_pattern}<-[r:`{rel_type}`]-{end_pattern}" else: rel_pattern = f"{start_pattern}-[r:`{rel_type}`]-{end_pattern}" - + # WHERE on r.* where_parts = [] params = {} - + if rel_conditions: for i, (key, condition) in enumerate(rel_conditions.items()): # *_start / *_end key suffix → real property name @@ -1824,32 +1904,36 @@ def subgraph_extract_by_rel_filter( actual_key = self._sanitize_property_key(actual_key) else: actual_key = self._sanitize_property_key(key) - + if isinstance(condition, (tuple, list)) and len(condition) == 2: operator, value = condition param_name = f"rel_cond_{i}" - + if operator.upper() in ["IN", "CONTAINS"]: - where_parts.append(f"r.`{actual_key}` {operator.upper()} ${param_name}") + where_parts.append( + f"r.`{actual_key}` {operator.upper()} ${param_name}" + ) elif operator.upper() == "STARTS WITH": - where_parts.append(f"r.`{actual_key}` STARTS WITH ${param_name}") + where_parts.append( + f"r.`{actual_key}` STARTS WITH ${param_name}" + ) else: where_parts.append(f"r.`{actual_key}` {operator} ${param_name}") - + params[param_name] = value else: # Equality param_name = f"rel_cond_{i}" where_parts.append(f"r.`{actual_key}` = ${param_name}") params[param_name] = condition - + where_clause = "WHERE " + " AND ".join(where_parts) if where_parts else "" - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" if limit is not None: params["limit"] = limit - + cypher = f""" MATCH {rel_pattern} {where_clause} @@ -1861,22 +1945,22 @@ def subgraph_extract_by_rel_filter( size(allNodes) AS node_count, size(allRels) AS relationship_count """ - + res = self.run(cypher, params) - + if res and res[0]: return { "nodes": res[0].get("nodes", []), "relationships": res[0].get("relationships", []), "node_count": res[0].get("node_count", 0), - "relationship_count": res[0].get("relationship_count", 0) + "relationship_count": res[0].get("relationship_count", 0), } - + return { "nodes": [], "relationships": [], "node_count": 0, - "relationship_count": 0 + "relationship_count": 0, } # ========================================================= @@ -1888,7 +1972,7 @@ def match_path_pattern( pattern: str, where: Optional[str] = None, params: Optional[JsonDict] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Run MATCH p = with optional WHERE (caller-built Cypher). @@ -1911,20 +1995,20 @@ def match_path_pattern( """ # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH p = {pattern} {("WHERE " + where) if where else ""} RETURN p AS path {limit_clause} """ - + p = {} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) # ========================================================= @@ -1938,7 +2022,7 @@ def aggregate_stats( where: Optional[str] = None, params: Optional[JsonDict] = None, metrics: Optional[Sequence[str]] = None, - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Simple aggregation on nodes of a label (optional GROUP BY). @@ -1965,15 +2049,15 @@ def aggregate_stats( """ label = self._sanitize_label(label) metrics = list(metrics) if metrics else ["count(*) AS cnt"] - + # Optional LIMIT limit_clause = "LIMIT $limit" if limit is not None else "" - + if group_by: group_by = self._sanitize_property_key(group_by) group_expr = f"n.`{group_by}` AS {group_by}" return_expr = ", ".join([group_expr] + list(metrics)) - + cypher = f""" MATCH (n:`{label}`) {("WHERE " + where) if where else ""} @@ -1984,21 +2068,22 @@ def aggregate_stats( """ else: return_expr = ", ".join(metrics) - + cypher = f""" MATCH (n:`{label}`) {("WHERE " + where) if where else ""} RETURN {return_expr} {limit_clause} """ - + p = {} if limit is not None: p["limit"] = limit if params: p.update(params) - + return self.run(cypher, p) + # ========================================================= # Paths between two nodes (not necessarily shortest) # ========================================================= @@ -2015,7 +2100,7 @@ def paths_between( return_fields: Optional[List[str]] = None, order_by: Optional[str] = None, order_direction: str = "ASC", - limit: Optional[int] = None + limit: Optional[int] = None, ) -> List[JsonDict]: """ Enumerate paths between A and B with hop bounds, filters, and derived metrics. @@ -2085,25 +2170,25 @@ def paths_between( """ (la, ka, va) = a (lb, kb, vb) = b - + # Validate identifiers and hop bounds la = self._sanitize_label(la) lb = self._sanitize_label(lb) ka = self._sanitize_property_key(ka) kb = self._sanitize_property_key(kb) - + if not (0 <= min_hops <= max_hops <= 10): raise ValueError("Invalid hop bounds: 0 <= min_hops <= max_hops <= 10") - + if direction not in {"out", "in", "both"}: raise ValueError("direction must be 'out', 'in', or 'both'") - + if order_direction not in {"ASC", "DESC"}: raise ValueError("order_direction must be 'ASC' or 'DESC'") - + rt = self._sanitize_rel_type(rel_type) if rel_type else "" rel = f":`{rt}`" if rt else "" - + # Variable-length pattern A … B if direction == "out": pattern = f"(A)-[r{rel}*{min_hops}..{max_hops}]->(B)" @@ -2111,15 +2196,15 @@ def paths_between( pattern = f"(A)<-[r{rel}*{min_hops}..{max_hops}]-(B)" else: pattern = f"(A)-[r{rel}*{min_hops}..{max_hops}]-(B)" - + where_clause = f"WHERE {where}" if where else "" - + if return_fields: # Custom RETURN list return_parts = [] for field in return_fields: field_lower = field.lower() - + if field_lower == "path": return_parts.append("p AS path") elif field_lower == "hops": @@ -2128,37 +2213,45 @@ def paths_between( return_parts.append("pathNodes AS nodes") elif field_lower == "relationships": return_parts.append("pathRels AS relationships") - + elif field_lower == "totalamount": - return_parts.append("REDUCE(s = 0, r IN pathRels | s + r.base_amt) AS totalAmount") - + return_parts.append( + "REDUCE(s = 0, r IN pathRels | s + r.base_amt) AS totalAmount" + ) + elif field_lower == "maxamount": - return_parts.append("REDUCE(m = 0, r IN pathRels | CASE WHEN r.base_amt > m THEN r.base_amt ELSE m END) AS maxAmount") - + return_parts.append( + "REDUCE(m = 0, r IN pathRels | CASE WHEN r.base_amt > m THEN r.base_amt ELSE m END) AS maxAmount" + ) + elif field_lower == "minamount": - return_parts.append("REDUCE(m = 999999, r IN pathRels | CASE WHEN r.base_amt < m THEN r.base_amt ELSE m END) AS minAmount") - + return_parts.append( + "REDUCE(m = 999999, r IN pathRels | CASE WHEN r.base_amt < m THEN r.base_amt ELSE m END) AS minAmount" + ) + elif field_lower == "avgamount": - return_parts.append("REDUCE(s = 0, r IN pathRels | s + r.base_amt) / hops AS avgAmount") - + return_parts.append( + "REDUCE(s = 0, r IN pathRels | s + r.base_amt) / hops AS avgAmount" + ) + else: return_parts.append(field) - + return_clause = "RETURN " + ", ".join(return_parts) else: return_clause = """RETURN p AS path, hops, pathNodes AS nodes, pathRels AS relationships""" - + # ORDER BY (default hops ASC) if order_by: order_clause = f"ORDER BY {order_by} {order_direction}" else: order_clause = "ORDER BY hops ASC" - + limit_clause = "LIMIT $limit" if limit is not None else "" - + cypher = f""" MATCH (A:`{la}` {{`{ka}`: $va}}), (B:`{lb}` {{`{kb}`: $vb}}) MATCH p = {pattern} @@ -2168,13 +2261,12 @@ def paths_between( {order_clause} {limit_clause} """ - + params = {"va": va, "vb": vb} if limit is not None: params["limit"] = limit - - return self.run(cypher, params) + return self.run(cypher, params) # ========================================== @@ -2186,47 +2278,42 @@ def paths_between( Neo4jConfig( uri="bolt://localhost:7687", user="neo4j", - password="password" + password=os.getenv("NEO4J_PASSWORD", ""), ) ) as client: - # Schema introspection schema = client.get_schema() print("=== Schema 信息 ===") print(f"节点类型: {list(schema['node_labels'].keys())}") print(f"关系类型: {list(schema['relationship_types'].keys())}") print(f"模式样例: {schema['patterns'][:3]}\n") - + # 1. Lookup by unique key print("=== 测试1: 唯一键查节点 ===") user = client.get_node_by_unique_key("User", "userId", "u123") print(f"用户: {user}\n") - + # 2. N-hop neighbors print("=== 测试2: N跳邻居 ===") neighbors = client.neighbors_n_hop( - "User", "userId", "u123", + "User", + "userId", + "u123", hops=2, rel_type="FOLLOWS", direction="out", - limit=10 + limit=10, ) print(f"找到 {len(neighbors)} 个邻居\n") - + # 3. Common neighbors print("=== 测试3: 公共邻居 ===") common = client.common_neighbors( - ("User", "userId", "u1"), - ("User", "userId", "u2"), - rel_type="FOLLOWS" + ("User", "userId", "u1"), ("User", "userId", "u2"), rel_type="FOLLOWS" ) print(f"找到 {len(common)} 个公共邻居\n") - + # 4. aggregate_stats print("=== 测试4: 聚合统计 ===") - stats = client.aggregate_stats( - "User", - group_by="country", - limit=5 - ) - print(f"统计结果: {stats}\n") \ No newline at end of file + stats = client.aggregate_stats("User", group_by="country", limit=5) + print(f"统计结果: {stats}\n") diff --git a/web/frontend/app.py b/web/frontend/app.py index ed734cd..3679af4 100644 --- a/web/frontend/app.py +++ b/web/frontend/app.py @@ -1,31 +1,47 @@ +""" +╔══════════════════════════════════════════════════════════════╗ +║ ⚠️ DEPRECATED — 此文件已废弃，请勿使用 ║ +║ ║ +║ app.py 是旧版独立 Flask 应用，与新版 route/ 蓝图架构冲突。 ║ +║ 新架构入口为 run.py，所有路由已迁移至 route/ 子模块。 ║ +║ ║ +║ 本文件保留功能代码以作参考，但独立启动已被禁用。 ║ +║ 如需添加新路由，请在 route/ 目录下创建对应蓝图模块。 ║ +╚══════════════════════════════════════════════════════════════╝ +""" + from datetime import datetime import os import json import logging import time import random -from flask import Flask, render_template, request, Response, stream_with_context, jsonify +from flask import ( + Flask, + render_template, + request, + Response, + stream_with_context, + jsonify, +) from flask_cors import CORS # 配置日志（方便调试API调用过程） logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # 初始化Flask应用 -app = Flask( - __name__, - template_folder="app", - static_folder="static" -) -CORS(app) # 解决跨域问题 +app = Flask(__name__, template_folder="app", static_folder="static") +CORS( + app, origins=os.getenv("CORS_ORIGINS", "http://localhost:5089").split(",") +) # CORS白名单从环境变量读取，默认仅允许localhost MODEL_MAPPING = { - "GPT 4": "qwen3-max", - "Qwen 14B": "qwen3-max", - "Qwen Plus": "qwen3-max" + "GPT 4": "qwen3-max", + "Qwen 14B": "qwen3-max", + "Qwen Plus": "qwen3-max", } # 知识库数据存储 # 知识库数据存储 @@ -58,101 +74,105 @@ { "type": "thinking", "contentType": "text", - "content": "用户想要一个JavaScript函数来在网页上显示随机名言。这需要几个组件：HTML结构来显示名言，JavaScript数组存储名言，以及函数来随机选择和显示名言。" + "content": "用户想要一个JavaScript函数来在网页上显示随机名言。这需要几个组件：HTML结构来显示名言，JavaScript数组存储名言，以及函数来随机选择和显示名言。", }, { "type": "result", "contentType": "dag", - "content": { - "nodes": [ - {"id": "1", "label": "用户问题：分析文档"}, - {"id": "2", "label": "提取关键词"}, - {"id": "3", "label": "检索知识库"}, - {"id": "4", "label": "生成回答"}, - {"id": "5", "label": "JavaScript代码示例"} - ], - "edges": [ - {"from": "1", "to": "2"}, - {"from": "2", "to": "3"}, - {"from": "3", "to": "1"}, - {"from": "3", "to": "4"}, - {"from": "3", "to": "5"} - ] - } - } + "content": { + "nodes": [ + {"id": "1", "label": "用户问题：分析文档"}, + {"id": "2", "label": "提取关键词"}, + {"id": "3", "label": "检索知识库"}, + {"id": "4", "label": "生成回答"}, + {"id": "5", "label": "JavaScript代码示例"}, + ], + "edges": [ + {"from": "1", "to": "2"}, + {"from": "2", "to": "3"}, + {"from": "3", "to": "1"}, + {"from": "3", "to": "4"}, + {"from": "3", "to": "5"}, + ], + }, + }, ], "dag_confirmation": [ { "type": "thinking", "contentType": "text", - "content": "用户确认了DAG结构正确，现在需要基于该DAG生成详细回答。首先我需要回顾DAG中的各个节点和流程。" + "content": "用户确认了DAG结构正确，现在需要基于该DAG生成详细回答。首先我需要回顾DAG中的各个节点和流程。", }, { "type": "thinking", "contentType": "text", - "content": "DAG显示了从用户问题到提取关键词，再到检索知识库，最后生成回答的完整流程。我需要按照这个逻辑展开详细说明。" + "content": "DAG显示了从用户问题到提取关键词，再到检索知识库，最后生成回答的完整流程。我需要按照这个逻辑展开详细说明。", }, { "type": "result", "contentType": "text", - "content": "根据您确认的DAG结构，以下是详细的处理流程说明：" + "content": "根据您确认的DAG结构，以下是详细的处理流程说明：", }, { "type": "result", "contentType": "text", - "content": "1. **用户问题分析（节点A）**：系统首先对用户输入的问题进行语义分析，确定问题类型和核心需求。" + "content": "1. **用户问题分析（节点A）**：系统首先对用户输入的问题进行语义分析，确定问题类型和核心需求。", }, { "type": "result", "contentType": "text", - "content": "2. **关键词提取（节点B）**：从分析后的问题中提取关键信息和术语，为后续知识库检索做准备。" + "content": "2. **关键词提取（节点B）**：从分析后的问题中提取关键信息和术语，为后续知识库检索做准备。", }, { "type": "result", "contentType": "code", "content": { "language": "python", - "code": "def extract_keywords(text):\n # 使用NLP工具提取关键词\n import jieba.analyse\n keywords = jieba.analyse.extract_tags(text, topK=10, withWeight=True)\n return [(word, weight) for word, weight in keywords]" - } + "code": "def extract_keywords(text):\n # 使用NLP工具提取关键词\n import jieba.analyse\n keywords = jieba.analyse.extract_tags(text, topK=10, withWeight=True)\n return [(word, weight) for word, weight in keywords]", + }, }, { "type": "result", "contentType": "text", - "content": "3. **知识库检索（节点C）**：基于提取的关键词在知识库中进行精确匹配和模糊搜索，获取相关文档。" + "content": "3. **知识库检索（节点C）**：基于提取的关键词在知识库中进行精确匹配和模糊搜索，获取相关文档。", }, { "type": "result", "contentType": "text", - "content": "4. **生成回答（节点D）**：结合检索到的知识和AI模型，生成准确、简洁的自然语言回答。" + "content": "4. **生成回答（节点D）**：结合检索到的知识和AI模型，生成准确、简洁的自然语言回答。", }, { "type": "result", "contentType": "code", "content": { "language": "python", - "code": "def categorize_goals(goals):\n categories = {\n 'personal': [],\n 'family': [],\n 'professional': []\n }\n for goal in goals:\n if any(word in goal.lower() for word in ['health', 'fitness', 'learn', 'read']):\n categories['personal'].append(goal)\n elif any(word in goal.lower() for word in ['family', 'spouse', 'children', 'home']):\n categories['family'].append(goal)\n elif any(word in goal.lower() for word in ['career', 'work', 'skill', 'project']):\n categories['professional'].append(goal)\n return categories" - } - } - ] + "code": "def categorize_goals(goals):\n categories = {\n 'personal': [],\n 'family': [],\n 'professional': []\n }\n for goal in goals:\n if any(word in goal.lower() for word in ['health', 'fitness', 'learn', 'read']):\n categories['personal'].append(goal)\n elif any(word in goal.lower() for word in ['family', 'spouse', 'children', 'home']):\n categories['family'].append(goal)\n elif any(word in goal.lower() for word in ['career', 'work', 'skill', 'project']):\n categories['professional'].append(goal)\n return categories", + }, + }, + ], } + ########################这一部分集中渲染页面################################# @app.route("/") def index(): """根路由：返回聊天页面""" - return render_template("template-chatbot-s2-convo.html") + return render_template("template-chatbot-s2-convo.html") + -@app.route('/overview') +@app.route("/overview") def overview(): - return render_template('overview.html') + return render_template("overview.html") + -@app.route('/documents') +@app.route("/documents") def documents(): - return render_template('documents.html') + return render_template("documents.html") -@app.route('/manage_dataset') + +@app.route("/manage_dataset") def manage_dataset_page(): - return render_template('manage_dataset.html') + return render_template("manage_dataset.html") @app.route("/api/chat", methods=["GET"]) @@ -162,13 +182,13 @@ def chat(): try: user_message = request.args.get("message", "").strip() selected_model = request.args.get("model", "") - dag_confirm = request.args.get("dag_confirm", "").strip() + dag_confirm = request.args.get("dag_confirm", "").strip() except Exception as e: logger.error(f"解析请求参数失败：{str(e)}") return Response( json.dumps({"error": "请求格式错误，请检查参数"}), mimetype="application/json", - status=400 + status=400, ) # 2. 验证参数合法性 @@ -176,9 +196,9 @@ def chat(): return Response( json.dumps({"error": "消息内容不能为空"}), mimetype="application/json", - status=400 + status=400, ) - + # 3. 选择测试数据 if dag_confirm == "yes": test_data_key = "dag_confirmation" @@ -188,10 +208,12 @@ def chat(): test_data_key = "dag" else: test_data_key = random.choice(list(TEST_DATA.keys())) - + test_data = TEST_DATA[test_data_key] - - logger.info(f"开始处理请求：模型={selected_model}，消息={user_message[:20]}...，dag_confirm={dag_confirm}，使用测试数据={test_data_key}") + + logger.info( + f"开始处理请求：模型={selected_model}，消息={user_message[:20]}...，dag_confirm={dag_confirm}，使用测试数据={test_data_key}" + ) # 4. 定义流式响应生成函数 @stream_with_context @@ -201,11 +223,11 @@ def generate_stream(): for item in test_data: # 模拟处理时间 time.sleep(0.6) - + # 返回数据 yield f"data: {json.dumps(item)}\n\n" logger.debug(f"返回流式数据：{item}") - + # 发送结束信号 yield "event: end\ndata: Stream completed\n\n" logger.info("流式响应完成") @@ -219,6 +241,7 @@ def generate_stream(): # 5. 返回流式响应（指定SSE格式） return Response(generate_stream(), mimetype="text/event-stream") + ########################API路由################################# @app.route("/api/models", methods=["GET"]) def get_models(): @@ -229,41 +252,33 @@ def get_models(): {"id": 1, "name": "GPT 4", "description": "OpenAI的GPT-4模型"}, {"id": 2, "name": "Qwen 14B", "description": "通义千问14B参数模型"}, {"id": 3, "name": "Qwen Plus", "description": "通义千问增强版模型"}, - {"id": 4, "name": "Llama 3", "description": "Meta的Llama 3开源模型"} + {"id": 4, "name": "Llama 3", "description": "Meta的Llama 3开源模型"}, ] - return jsonify({ - "success": True, - "data": models, - "count": len(models) - }) + return jsonify({"success": True, "data": models, "count": len(models)}) except Exception as e: logger.error(f"获取模型列表失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "获取模型列表失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "获取模型列表失败", "message": str(e)} + ), 500 + @app.route("/api/knowledge_bases", methods=["GET"]) def get_knowledge_bases(): """获取知识库列表 - 简化版本""" try: logger.info("收到知识库查询请求") - + # 直接返回数据，不进行任何复杂处理 - return jsonify({ - "success": True, - "data": knowledge_bases, - "count": len(knowledge_bases) - }) - + return jsonify( + {"success": True, "data": knowledge_bases, "count": len(knowledge_bases)} + ) + except Exception as e: logger.error(f"获取知识库列表失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "获取知识库列表失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "获取知识库列表失败", "message": str(e)} + ), 500 + @app.route("/api/knowledge_bases/", methods=["DELETE"]) def delete_knowledge_base(kb_id): @@ -271,31 +286,23 @@ def delete_knowledge_base(kb_id): try: global knowledge_bases logger.info(f"收到删除知识库请求，ID: {kb_id}") - + # 找到要删除的知识库索引 original_count = len(knowledge_bases) knowledge_bases = [kb for kb in knowledge_bases if kb["id"] != kb_id] - + if len(knowledge_bases) < original_count: logger.info(f"成功删除知识库 ID: {kb_id}") - return jsonify({ - "success": True, - "message": f"成功删除知识库" - }) + return jsonify({"success": True, "message": f"成功删除知识库"}) else: logger.warning(f"未找到指定的知识库 ID: {kb_id}") - return jsonify({ - "success": False, - "error": "未找到指定的知识库" - }), 404 - + return jsonify({"success": False, "error": "未找到指定的知识库"}), 404 + except Exception as e: logger.error(f"删除知识库失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "删除知识库失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "删除知识库失败", "message": str(e)} + ), 500 @app.route("/api/knowledge_bases", methods=["POST"]) @@ -304,23 +311,20 @@ def create_knowledge_base(): try: data = request.get_json() logger.info(f"收到创建知识库请求: {data}") - + name = data.get("name") or data.get("名称") if not data or not name: - return jsonify({ - "success": False, - "error": "知识库名称不能为空" - }), 400 - + return jsonify({"success": False, "error": "知识库名称不能为空"}), 400 + file_type = data.get("file_type") or data.get("文件类型", "text") - + # 验证文件类型是否合法 if file_type not in ["text", "graph"]: file_type = "text" # 默认值 - + # 生成新ID new_id = max([kb["id"] for kb in knowledge_bases]) + 1 if knowledge_bases else 1 - + new_kb = { "id": new_id, "name": name, @@ -329,74 +333,72 @@ def create_knowledge_base(): "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), } knowledge_bases.append(new_kb) - + print(knowledge_bases) logger.info(f"成功创建知识库: {new_kb['name']}, 文件类型: {file_type}") - return jsonify({ - "success": True, - "data": new_kb - }) - + return jsonify({"success": True, "data": new_kb}) + except Exception as e: logger.error(f"创建知识库失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "创建知识库失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "创建知识库失败", "message": str(e)} + ), 500 + def get_dataset_type_function(kb_id): """根据知识库ID获取数据集类型，支持遍历查找""" try: # 将kb_id转换为整数 kb_id_int = int(kb_id) - + # 遍历查找匹配的知识库 for kb in knowledge_bases: if kb["id"] == kb_id_int: return kb.get("file_type") or kb.get("文件类型", "text") - + # 如果没有找到，返回默认值 return "text" - + except (ValueError, TypeError): # 如果kb_id不是有效的数字，返回默认值 return "text" + # Flask 路由示例 -@app.route('/api/dataset_type') +@app.route("/api/dataset_type") def get_dataset_type(): - kb_id = request.args.get('kb_id') - + kb_id = request.args.get("kb_id") + # 检查参数是否存在 if not kb_id: - return jsonify({ - 'success': False, - 'error': '缺少kb_id参数', - 'dataset_type': 'text' # 提供默认值 - }), 400 - + return jsonify( + { + "success": False, + "error": "缺少kb_id参数", + "dataset_type": "text", # 提供默认值 + } + ), 400 + dataset_type = get_dataset_type_function(kb_id) - - return jsonify({ - 'success': True, - 'dataset_type': dataset_type, # "text" 或 "graph" - 'kb_id': kb_id # 返回请求的kb_id用于调试 - }) + + return jsonify( + { + "success": True, + "dataset_type": dataset_type, # "text" 或 "graph" + "kb_id": kb_id, # 返回请求的kb_id用于调试 + } + ) + @app.route("/api/health", methods=["GET"]) def health_check(): """健康检查接口""" - return jsonify({ - "status": "healthy", - "timestamp": datetime.now().isoformat() - }) + return jsonify({"status": "healthy", "timestamp": datetime.now().isoformat()}) if __name__ == "__main__": - app.run( - debug=True, - host="0.0.0.0", - port=5000 - ) \ No newline at end of file + # ⚠️ 独立启动已禁用 — 请使用 run.py 作为入口 + print("[DEPRECATED] app.py 独立启动已禁用。请使用 run.py 启动应用。") + print("[DEPRECATED] 所有新路由请在 route/ 目录下以蓝图模块方式创建。") + # app.run(debug=os.getenv("FLASK_DEBUG", "0") == "1", host="0.0.0.0", port=5000) diff --git a/web/frontend/models.json b/web/frontend/models.json.example similarity index 51% rename from web/frontend/models.json rename to web/frontend/models.json.example index b523e33..55c59f6 100644 --- a/web/frontend/models.json +++ b/web/frontend/models.json.example @@ -1,9 +1,10 @@ [ { + "_comment": "将 api_key 设置为你的实际密钥，或通过环境变量 OPENAI_API_KEY 注入", "id": 1, "name": "gpt-4o-mini", "base_url": "https://gitaigc.com/v1/", - "api_key": "sk-G30rFStBigqXtuyIOkOo7Zh4QNxO8ZAjfZQ5DYPCgMXbPv8q", + "api_key": "YOUR_API_KEY", "created_at": "2025-12-08 16:18:43" } -] \ No newline at end of file +] diff --git a/web/frontend/route/__init__.py b/web/frontend/route/__init__.py index 8994ff1..a0498d3 100644 --- a/web/frontend/route/__init__.py +++ b/web/frontend/route/__init__.py @@ -5,24 +5,29 @@ from flask import Flask from flask_cors import CORS -from flask_socketio import SocketIO +from flask_socketio import SocketIO -PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) +PROJECT_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..") +) if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) -# from aag.api.async_runtime import start_async_runtime, stop_async_runtime +from aag.api.async_runtime import start_async_runtime, stop_async_runtime # Global logging configuration (set once) logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) -socketio = SocketIO(cors_allowed_origins="*") +socketio = SocketIO( + cors_allowed_origins=os.getenv("CORS_ORIGINS", "http://localhost:5089").split(","), + ping_interval=25, + ping_timeout=60, +) def create_app(): @@ -35,10 +40,10 @@ def create_app(): static_folder=os.path.join(base_dir, "..", "static"), ) - CORS(app) + CORS(app, origins=os.getenv("CORS_ORIGINS", "http://localhost:5089").split(",")) # ====== Register blueprints (route modules) ====== - from .routes_pages import bp as pages_bp + from .routes_pages import bp as pages_bp from .routes_health import bp as health_bp from .routes_documents import bp as documents_bp from .routes_manage_dataset import bp as manage_bp @@ -52,14 +57,24 @@ def create_app(): # ====== Initialize SocketIO and async runtime ====== socketio.init_app(app) - # start_async_runtime() - # atexit.register(stop_async_runtime) - + + # 启动后台异步运行时，为路由模块提供共享事件循环 + # 避免每次异步调用创建/销毁事件循环的开销 + try: + start_async_runtime() + logger.info("异步运行时已启动（后台事件循环 + ChatService 初始化完成）") + except Exception as exc: + logger.error("异步运行时启动失败: %s，路由将回退到 asyncio.run()", exc) + + # 确保应用退出时清理后台线程和事件循环（即使 start 失败 stop 也会安全跳过） + atexit.register(stop_async_runtime) + # ====== Register WebSocket event handlers ====== try: - from . import sockets_chat + from . import sockets_chat + logger.info("WebSocket event module sockets_chat loaded") except ImportError as e: logger.warning(f"Failed to import sockets_chat: {e}") - return app \ No newline at end of file + return app diff --git a/web/frontend/route/model_schema.py b/web/frontend/route/model_schema.py index 22f85c7..e303d98 100644 --- a/web/frontend/route/model_schema.py +++ b/web/frontend/route/model_schema.py @@ -1,83 +1,127 @@ import json import os +import base64 import logging from datetime import datetime -DATA_FILE = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "models.json") +DATA_FILE = os.path.join( + os.path.abspath(os.path.join(os.path.dirname(__file__), "..")), "models.json" +) logger = logging.getLogger(__name__) +# 前端模型名称到后端实际模型标识的映射 +# 统一来源：config.py 中的 MODEL_MAPPING 已迁移至此，请从此处导入。 +MODEL_MAPPING = { + "GPT 4": "qwen3-max", + "Qwen 14B": "qwen3-max", + "Qwen Plus": "qwen3-max", +} + + def _init_data_file(): """初始化数据文件""" if not os.path.exists(DATA_FILE): try: - with open(DATA_FILE, 'w', encoding='utf-8') as f: + with open(DATA_FILE, "w", encoding="utf-8") as f: json.dump([], f) except Exception as e: logger.error(f"初始化模型数据文件失败: {e}") + +def _obfuscate_api_key(api_key: str) -> str: + """ + 对 API Key 进行简单混淆（Base64 编码）后存储。 + ⚠️ 注意：这是混淆（obfuscation）而非加密（encryption）， + Base64 可被轻易解码。生产环境应使用密钥管理服务（如 HashiCorp Vault、AWS KMS）。 + """ + if not api_key: + return api_key + return base64.b64encode(api_key.encode("utf-8")).decode("utf-8") + + +def _deobfuscate_api_key(encoded_key: str) -> str: + """反向解码混淆后的 API Key。""" + if not encoded_key: + return encoded_key + try: + return base64.b64decode(encoded_key.encode("utf-8")).decode("utf-8") + except Exception: + # 兼容已存储的未混淆明文 Key + return encoded_key + + def load_models(): - """加载所有模型""" + """加载所有模型（自动解码 API Key）""" _init_data_file() try: - with open(DATA_FILE, 'r', encoding='utf-8') as f: - return json.load(f) + with open(DATA_FILE, "r", encoding="utf-8") as f: + models = json.load(f) + # 解码存储的 API Key（兼容明文和编码两种格式） + for model in models: + if "api_key" in model: + model["api_key"] = _deobfuscate_api_key(model["api_key"]) + return models except Exception as e: logger.error(f"读取模型数据失败: {e}") return [] + def save_models(models): """保存模型列表到文件""" try: - with open(DATA_FILE, 'w', encoding='utf-8') as f: + with open(DATA_FILE, "w", encoding="utf-8") as f: json.dump(models, f, ensure_ascii=False, indent=4) return True except Exception as e: logger.error(f"保存模型数据失败: {e}") return False + def create_model(name, base_url, api_key): """创建新模型""" models = load_models() - + # 生成新 ID (如果列表为空则为1，否则为最大ID+1) new_id = 1 if models: - new_id = max(m.get('id', 0) for m in models) + 1 - + new_id = max(m.get("id", 0) for m in models) + 1 + current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - + new_model = { "id": new_id, "name": name, "base_url": base_url, - "api_key": api_key, # 注意：实际生产环境中 API Key 应当加密存储 - "created_at": current_time + "api_key": _obfuscate_api_key(api_key), # 混淆存储；生产环境应使用 KMS 加密 + "created_at": current_time, } - + models.append(new_model) - + if save_models(models): return new_model return None + def delete_model(model_id): """删除模型""" models = load_models() - + # 过滤掉要删除的 ID - new_models = [m for m in models if m.get('id') != model_id] - + new_models = [m for m in models if m.get("id") != model_id] + if len(new_models) < len(models): save_models(new_models) return True return False + def get_model_by_id(model_id): """根据 ID 获取模型详情""" models = load_models() for m in models: - if m.get('id') == model_id: + if m.get("id") == model_id: return m - return None \ No newline at end of file + return None diff --git a/web/frontend/route/routes_documents.py b/web/frontend/route/routes_documents.py index 117df73..1aa90f1 100644 --- a/web/frontend/route/routes_documents.py +++ b/web/frontend/route/routes_documents.py @@ -3,6 +3,7 @@ import sys import json import asyncio +from pathlib import Path from document_schema import ( load_knowledge_bases, create_knowledge_base, @@ -11,23 +12,58 @@ update_all_knowledge_bases_file_count, ) from DummySocket import DummySocket -sys.path.append("../../") + +# ⚠️ 此文件位于 web/frontend/route/，需要从顶级包 aag/ 导入模块 +# web/ 和 web/frontend/ 已有 __init__.py，但 aag 与 web 是同级顶级包（非父子关系）， +# Python 包相对导入无法跨顶级包边界（from ....aag... 会触发 ValueError）。 +# 折中方案：使用 pathlib.Path 计算项目根目录并临时加入 sys.path， +# 使 aag/ 可作为顶级包被绝对导入（from aag.api.DocumentAPI import server_Test）。 +_PROJECT_ROOT = str( + Path(__file__).resolve().parents[3] +) # route → frontend → web → YiGraph +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + from aag.api.DocumentAPI import server_Test +# 尝试使用后台事件循环运行异步任务，避免每次请求创建/销毁事件循环 +try: + from aag.api.async_runtime import get_background_loop + + _HAS_BACKGROUND_LOOP = True +except ImportError: + _HAS_BACKGROUND_LOOP = False + # 创建蓝图 -bp = Blueprint('documents', __name__, url_prefix='/api') +bp = Blueprint("documents", __name__, url_prefix="/api") # 配置日志 logger = logging.getLogger(__name__) + +def _run_async(coro, timeout: int = 30): + """ + 在后台事件循环中执行异步协程；若运行时未初始化则回退到 asyncio.run()。 + 回退时会记录警告日志，帮助排查生产环境中异步运行时未启动的问题。 + """ + if _HAS_BACKGROUND_LOOP: + try: + loop = get_background_loop() + future = asyncio.run_coroutine_threadsafe(coro, loop) + return future.result(timeout=timeout) + except RuntimeError: + pass + logger.warning( + "异步运行时未初始化，回退到 asyncio.run()（每次调用将创建并销毁事件循环）" + ) + return asyncio.run(coro) + + def load_knowledge_bases(): a23 = DummySocket(json.dumps({"action": "get_datasets"})) - asyncio.run(server_Test.handler(a23)) + _run_async(server_Test.handler(a23)) knowledge_bases_with_count = a23.returnmsg - print(jsonify({ - 'success': True, - 'data': knowledge_bases_with_count - })) + print(jsonify({"success": True, "data": knowledge_bases_with_count})) basescy = json.loads(knowledge_bases_with_count) return basescy["content"]["data"] @@ -48,51 +84,61 @@ def _parse_ws_result(raw_msg: str): return payload return None + def get_knowledge_base_name(kb_id): """根据知识库ID获取知识库名称""" try: - kb_id = int(kb_id) - knowledge_bases = load_knowledge_bases() + kb_id = int(kb_id) + knowledge_bases = load_knowledge_bases() for kb in knowledge_bases: if kb["id"] == kb_id: return kb["name"] - return f"kb_{kb_id}" + return f"kb_{kb_id}" except Exception as e: logger.error(f"获取知识库名称错误: {str(e)}") return f"kb_{kb_id}" + # 获取知识库列表 @bp.route("/knowledge_bases", methods=["GET"]) def get_knowledge_bases(): """获取知识库列表""" try: a1 = DummySocket(json.dumps({"action": "get_datasets"})) - asyncio.run(server_Test.handler(a1)) + _run_async(server_Test.handler(a1)) gkb = a1.returnmsg knowledge_bases = json.loads(gkb) for kb in knowledge_bases["content"]["data"]: if kb["file_type"] == "graph": if kb["file_count"] == 1: - a6 = DummySocket(json.dumps({"action": "get_dataset_schema","ds_name":kb["name"]})) - asyncio.run(server_Test.handler(a6)) + a6 = DummySocket( + json.dumps( + {"action": "get_dataset_schema", "ds_name": kb["name"]} + ) + ) + _run_async(server_Test.handler(a6)) gkb6 = a6.returnmsg gkb6_json = json.loads(gkb6) - if gkb6_json["content"]["data"][0].get("vertex_file",None) is not None: + if ( + gkb6_json["content"]["data"][0].get("vertex_file", None) + is not None + ): kb["file_count"] = 2 - - return jsonify({ - "success": True, - "data": knowledge_bases["content"]["data"], - "count": len(knowledge_bases) - }) - + + return jsonify( + { + "success": True, + "data": knowledge_bases["content"]["data"], + "count": len(knowledge_bases), + } + ) + except Exception as e: logger.error(f"获取知识库列表失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "获取知识库列表失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "获取知识库列表失败", "message": str(e)} + ), 500 + # 创建知识库 @bp.route("/knowledge_bases", methods=["POST"]) @@ -103,10 +149,7 @@ def create_knowledge_base_route(): name = data.get("name") or data.get("名称") if not data or not name: - return jsonify({ - "success": False, - "error": "知识库名称不能为空" - }), 400 + return jsonify({"success": False, "error": "知识库名称不能为空"}), 400 file_type = data.get("file_type") or data.get("文件类型", "text") if file_type not in ["text", "graph"]: @@ -118,79 +161,69 @@ def create_knowledge_base_route(): # file_type=file_type # ) - a2 = DummySocket(json.dumps({"action": "create_dataset","name": name,"type": file_type})) - asyncio.run(server_Test.handler(a2)) + a2 = DummySocket( + json.dumps({"action": "create_dataset", "name": name, "type": file_type}) + ) + _run_async(server_Test.handler(a2)) result = _parse_ws_result(getattr(a2, "returnmsg", "{}")) if result and result.get("success"): # DocumentAPI 返回 content.data 里包含 db_name/message - return jsonify({ - "success": True, - "data": result.get("data", {}) - }) + return jsonify({"success": True, "data": result.get("data", {})}) else: err = None if result: err = result.get("error") or result.get("message") - return jsonify({ - "success": False, - "error": err or "创建知识库失败" - }), 500 - + return jsonify({"success": False, "error": err or "创建知识库失败"}), 500 + except Exception as e: logger.error(f"创建知识库失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "创建知识库失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "创建知识库失败", "message": str(e)} + ), 500 + @bp.route("/knowledge_bases/", methods=["DELETE"]) def delete_knowledge_base_route(kb_id): """删除知识库""" try: - # 获取要删除的知识库名称，用于后续调用 - #knowledge_bases = load_knowledge_bases() + # 获取要删除的知识库名称，用于后续调用 + # knowledge_bases = load_knowledge_bases() kb_to_delete = get_knowledge_base_name(kb_id) # for kb in knowledge_bases: # if kb.get("id") == kb_id: # kb_to_delete = kb # break - + if not kb_to_delete: logger.warning(f"未找到指定的知识库 ID: {kb_id}") - return jsonify({ - "success": False, - "error": "未找到指定的知识库" - }), 404 - - a3 = DummySocket(json.dumps({"action": "delete_dataset","ds_name":kb_to_delete})) - asyncio.run(server_Test.handler(a3)) + return jsonify({"success": False, "error": "未找到指定的知识库"}), 404 + + a3 = DummySocket( + json.dumps({"action": "delete_dataset", "ds_name": kb_to_delete}) + ) + _run_async(server_Test.handler(a3)) result = _parse_ws_result(getattr(a3, "returnmsg", "{}")) if result and result.get("success"): - return jsonify({ - "success": True, - "message": result.get("message") or "成功删除知识库" - }) + return jsonify( + {"success": True, "message": result.get("message") or "成功删除知识库"} + ) else: err = None if result: err = result.get("error") or result.get("message") logger.warning(f"未找到指定的知识库 ID: {kb_id}") - return jsonify({ - "success": False, - "error": err or "未找到指定的知识库" - }), 404 - + return jsonify( + {"success": False, "error": err or "未找到指定的知识库"} + ), 404 + except Exception as e: logger.error(f"删除知识库失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "删除知识库失败", - "message": str(e) - }), 500 - + return jsonify( + {"success": False, "error": "删除知识库失败", "message": str(e)} + ), 500 + # 获取知识库文件个数（实时统计，不更新JSON） @bp.route("/knowledge_bases//file_count", methods=["GET"]) @@ -198,22 +231,17 @@ def get_knowledge_base_file_count(kb_id): """获取指定知识库的实时文件个数（不更新JSON）""" try: file_count = count_files_in_knowledge_base(kb_id) - - return jsonify({ - "success": True, - "data": { - "kb_id": kb_id, - "file_count": file_count - } - }) - + + return jsonify( + {"success": True, "data": {"kb_id": kb_id, "file_count": file_count}} + ) + except Exception as e: logger.error(f"获取知识库文件个数失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "获取文件个数失败", - "message": str(e) - }), 500 + return jsonify( + {"success": False, "error": "获取文件个数失败", "message": str(e)} + ), 500 + # 更新所有知识库文件个数 @bp.route("/knowledge_bases/update_all_file_counts", methods=["POST"]) @@ -221,26 +249,23 @@ def update_all_file_counts_route(): """更新所有知识库的文件个数""" try: success = update_all_knowledge_bases_file_count() - + if success: # 重新加载更新后的知识库数据 knowledge_bases = load_knowledge_bases() - - return jsonify({ - "success": True, - "message": "成功更新所有知识库文件个数", - "data": knowledge_bases - }) + + return jsonify( + { + "success": True, + "message": "成功更新所有知识库文件个数", + "data": knowledge_bases, + } + ) else: - return jsonify({ - "success": False, - "error": "更新文件个数失败" - }), 500 - + return jsonify({"success": False, "error": "更新文件个数失败"}), 500 + except Exception as e: logger.error(f"更新所有文件个数失败：{str(e)}", exc_info=True) - return jsonify({ - "success": False, - "error": "更新所有文件个数失败", - "message": str(e) - }), 500 \ No newline at end of file + return jsonify( + {"success": False, "error": "更新所有文件个数失败", "message": str(e)} + ), 500 diff --git a/web/frontend/route/routes_manage_dataset.py b/web/frontend/route/routes_manage_dataset.py index 3e0d942..5c2d816 100644 --- a/web/frontend/route/routes_manage_dataset.py +++ b/web/frontend/route/routes_manage_dataset.py @@ -6,22 +6,62 @@ import asyncio import base64 import yaml +from pathlib import Path from flask import Blueprint, jsonify, request from flask_socketio import emit, join_room, leave_room from DummySocket import DummySocket -sys.path.append("../../") + +# ⚠️ 此文件位于 web/frontend/route/，需要从顶级包 aag/ 导入模块 +# web/ 和 web/frontend/ 已有 __init__.py，但 aag 与 web 是同级顶级包（非父子关系）， +# Python 包相对导入无法跨顶级包边界（from ....aag... 会触发 ValueError）。 +# 折中方案：使用 pathlib.Path 计算项目根目录并临时加入 sys.path， +# 使 aag/ 可作为顶级包被绝对导入（from aag.api.DocumentAPI import server_Test）。 +_PROJECT_ROOT = str( + Path(__file__).resolve().parents[3] +) # route → frontend → web → YiGraph +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + from aag.utils.path_utils import DATASETS_DATA_DIR from aag.api.DocumentAPI import server_Test +# 尝试使用后台事件循环运行异步任务，避免每次请求创建/销毁事件循环 +try: + from aag.api.async_runtime import get_background_loop + + _HAS_BACKGROUND_LOOP = True +except ImportError: + _HAS_BACKGROUND_LOOP = False + + +def _run_async(coro, timeout: int = 30): + """ + 在后台事件循环中执行异步协程；若运行时未初始化则回退到 asyncio.run()。 + 回退时会记录警告日志，帮助排查生产环境中异步运行时未启动的问题。 + """ + if _HAS_BACKGROUND_LOOP: + try: + loop = get_background_loop() + future = asyncio.run_coroutine_threadsafe(coro, loop) + return future.result(timeout=timeout) + except RuntimeError: + pass + logger.warning( + "异步运行时未初始化，回退到 asyncio.run()（每次调用将创建并销毁事件循环）" + ) + return asyncio.run(coro) + + class MessageCollectingSocket: """包装DummySocket，收集所有消息，优先返回error消息""" + def __init__(self, message): self.socket = DummySocket(message) self.messages = [] self.has_error = False self.error_message = None - + async def send(self, msg): # 调用原始socket的send方法 await self.socket.send(msg) @@ -35,7 +75,7 @@ async def send(self, msg): self.error_message = msg except: pass - + @property def returnmsg(self): # 如果有错误，返回错误消息；否则返回最后一个消息 @@ -44,14 +84,15 @@ def returnmsg(self): elif self.messages: return self.messages[-1] else: - return self.socket.returnmsg if hasattr(self.socket, 'returnmsg') else None - + return self.socket.returnmsg if hasattr(self.socket, "returnmsg") else None + def __getattr__(self, name): # 代理其他属性到原始socket return getattr(self.socket, name) + formatter = logging.Formatter( - '%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s' + "%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s" ) logger = logging.getLogger(__name__) @@ -63,78 +104,79 @@ def __getattr__(self, name): # 添加新的API端点来检查解析状态 -@bp.route('/api/check_parsing_status', methods=['POST']) +@bp.route("/api/check_parsing_status", methods=["POST"]) def api_check_parsing_status(): """检查数据集所有文件的解析状态""" try: data = request.get_json() - kb_id = data.get('kb_id') - files = data.get('files', []) - + kb_id = data.get("kb_id") + files = data.get("files", []) + if not kb_id: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id' - }), 400 - + return jsonify({"success": False, "error": "Missing kb_id"}), 400 + kb_name = get_knowledge_base_name(kb_id) kb_type = get_dataset_type_from_back(kb_id) # 获取解析状态 - a24 = DummySocket(json.dumps({"action": "get_parsing_status","ds_name":kb_name})) - asyncio.run(server_Test.handler(a24)) + a24 = DummySocket( + json.dumps({"action": "get_parsing_status", "ds_name": kb_name}) + ) + _run_async(server_Test.handler(a24)) acps = a24.returnmsg acpss = json.loads(acps) - if kb_type == 'graph': - a25 = DummySocket(json.dumps({"action": "get_dataset_schema","ds_name":kb_name})) - asyncio.run(server_Test.handler(a25)) + if kb_type == "graph": + a25 = DummySocket( + json.dumps({"action": "get_dataset_schema", "ds_name": kb_name}) + ) + _run_async(server_Test.handler(a25)) acps5 = a25.returnmsg acpss5 = json.loads(acps5) - if acpss5["content"]["data"][0].get("vertex_file",None) is not None: - acpss["content"]["data"]["parsing_status"][acpss5["content"]["data"][0]["vertex_file"].split("/")[-1]] = "completed" - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'all_parsed': acpss["content"]["data"]["total_parsing_status"], - 'file_status': acpss["content"]["data"]["parsing_status"], - #'dataset_type': parsing_status['dataset_type'], - #'message': f"解析状态: {len(files)}个文件中{sum(parsing_status['file_status'].values())}个已解析" - }) - + if acpss5["content"]["data"][0].get("vertex_file", None) is not None: + acpss["content"]["data"]["parsing_status"][ + acpss5["content"]["data"][0]["vertex_file"].split("/")[-1] + ] = "completed" + return jsonify( + { + "success": True, + "kb_id": kb_id, + "all_parsed": acpss["content"]["data"]["total_parsing_status"], + "file_status": acpss["content"]["data"]["parsing_status"], + #'dataset_type': parsing_status['dataset_type'], + #'message': f"解析状态: {len(files)}个文件中{sum(parsing_status['file_status'].values())}个已解析" + } + ) + except Exception as e: logger.error(f"检查解析状态失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 + def load_knowledge_bases(): a23 = DummySocket(json.dumps({"action": "get_datasets"})) - asyncio.run(server_Test.handler(a23)) + _run_async(server_Test.handler(a23)) knowledge_bases_with_count = a23.returnmsg - print(jsonify({ - 'success': True, - 'data': knowledge_bases_with_count - })) + logger.debug(f"load_knowledge_bases 原始返回: {knowledge_bases_with_count}") basescy = json.loads(knowledge_bases_with_count) return basescy["content"]["data"] - + + def format_string_to_list(input_str): """ 将逗号分隔的字符串转换为列表格式的字符串 - + 参数: input_str: 逗号分隔的字符串，例如 "name,age" - + 返回: 字符串格式的列表表示，例如 "['name','age']" """ # 分割字符串并去除每个元素两端的空格 - items = [item.strip() for item in input_str.split(',') if item.strip()] - + items = [item.strip() for item in input_str.split(",") if item.strip()] + # 构建格式化的字符串 formatted_items = [f"'{item}'" for item in items] result = f"[{','.join(formatted_items)}]" - + return result @@ -144,102 +186,105 @@ def get_knowledge_base_name(kb_id): # 兼容 "kb_6" 这类前缀 if isinstance(kb_id, str) and kb_id.startswith("kb_"): kb_id = kb_id[3:] - kb_id = int(kb_id) - knowledge_bases = load_knowledge_bases() + kb_id = int(kb_id) + knowledge_bases = load_knowledge_bases() for kb in knowledge_bases: if kb["id"] == kb_id: return kb["name"] - return f"kb_{kb_id}" + return f"kb_{kb_id}" except Exception as e: logger.error(f"获取知识库名称错误: {str(e)}") return f"kb_{kb_id}" + def get_dataset_type_from_back(kb_id): """根据知识库ID获取知识库类型""" try: if isinstance(kb_id, str) and kb_id.startswith("kb_"): kb_id = kb_id[3:] - kb_id = int(kb_id) - knowledge_bases = load_knowledge_bases() + kb_id = int(kb_id) + knowledge_bases = load_knowledge_bases() for kb in knowledge_bases: if kb["id"] == kb_id: return kb["file_type"] - return f"kb_{kb_id}" + return f"kb_{kb_id}" except Exception as e: logger.error(f"获取文件类型错误: {str(e)}") return f"kb_{kb_id}" + def get_files_for_knowledge_base(kb_id): """获取指定知识库的所有文件信息""" try: kb_name = get_knowledge_base_name(kb_id) - a23 = DummySocket(json.dumps({"action": "get_dataset_schema","ds_name": kb_name})) - asyncio.run(server_Test.handler(a23)) + a23 = DummySocket( + json.dumps({"action": "get_dataset_schema", "ds_name": kb_name}) + ) + _run_async(server_Test.handler(a23)) regffkb = json.loads(a23.returnmsg) reregffkb = regffkb["content"]["data"] if reregffkb[0]["type"] == "graph": reregffkb[0]["size"] = reregffkb[0]["edge_size"] - if reregffkb[0].get("vertex_file",None) is not None: - reregffkb.append( { - "id" : 2, - "type" : "graph", - "name" : reregffkb[0]["vertex_file"].split("/")[-1], - "size" : reregffkb[0]["vertex_size"], - "uploadDate" : reregffkb[0]["uploadDate"], - "graph_status" : "completed", - }) + if reregffkb[0].get("vertex_file", None) is not None: + reregffkb.append( + { + "id": 2, + "type": "graph", + "name": reregffkb[0]["vertex_file"].split("/")[-1], + "size": reregffkb[0]["vertex_size"], + "uploadDate": reregffkb[0]["uploadDate"], + "graph_status": "completed", + } + ) else: pass - #return files + # return files return reregffkb except Exception as e: logger.error(f"获取知识库 {kb_id} 文件列表错误: {str(e)}") return [] -@bp.route('/api/knowledge_bases//files', methods=['GET']) + +@bp.route("/api/knowledge_bases//files", methods=["GET"]) def get_knowledge_base_files(kb_id): """获取指定知识库的文件列表""" try: files = get_files_for_knowledge_base(kb_id) - return jsonify({ - 'success': True, - 'data': files - }) + return jsonify({"success": True, "data": files}) except Exception as e: logger.error(f"获取知识库文件列表失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 + -@bp.route('/api/preview_file', methods=['GET']) +@bp.route("/api/preview_file", methods=["GET"]) def api_preview_file(): """获取文件内容用于预览""" try: - kb_id = request.args.get('kb_id') - file_name = request.args.get('file_name') - + kb_id = request.args.get("kb_id") + file_name = request.args.get("file_name") + if not kb_id or not file_name: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id or file_name' - }), 400 - + return jsonify( + {"success": False, "error": "Missing kb_id or file_name"} + ), 400 + # 获取知识库名称 kb_name = get_knowledge_base_name(kb_id) kb_type = get_dataset_type_from_back(kb_id) - + # 获取文件路径 file_path = None - if kb_type == 'text': + if kb_type == "text": # 文本数据集：从schema中获取original_path server_Test.load_each_dataset(kb_name) if file_name in server_Test.each_dataset: - file_path = server_Test.each_dataset[file_name]["schema"].get("original_path") - elif kb_type == 'graph': + file_path = server_Test.each_dataset[file_name]["schema"].get( + "original_path" + ) + elif kb_type == "graph": # 图数据集：从schema中获取original_path server_Test.load_each_dataset(kb_name) - + # 获取dataset_folder的绝对路径 # server_Test.dataset_folder 是相对于 DocumentAPI.py 的路径 # DocumentAPI.py 在 aag/api/ 目录下 @@ -247,41 +292,57 @@ def api_preview_file(): # 获取 DocumentAPI.py 的目录 import inspect from aag.api.DocumentAPI import DocumentAPIServer + api_file = inspect.getfile(DocumentAPIServer) api_dir = os.path.dirname(os.path.abspath(api_file)) - dataset_folder_abs = os.path.normpath(os.path.join(api_dir, server_Test.dataset_folder)) + dataset_folder_abs = os.path.normpath( + os.path.join(api_dir, server_Test.dataset_folder) + ) else: dataset_folder_abs = server_Test.dataset_folder - + # 图数据集可能有多个文件，需要遍历查找 for dataset_name, dataset_info in server_Test.each_dataset.items(): schema = dataset_info.get("schema", {}) - + # 辅助函数：解析路径 def resolve_file_path(original_path, target_file_name): """解析文件路径""" if not original_path: return None - + # 检查文件名是否匹配 path_file_name = os.path.basename(original_path) - if target_file_name != path_file_name and target_file_name not in original_path: + if ( + target_file_name != path_file_name + and target_file_name not in original_path + ): return None - + # 解析路径：从original_path中提取data/之后的部分 if not os.path.isabs(original_path): # original_path格式：../../aag/datasets/data/{kb_name}/graph/{file_name} # 提取 data/ 之后的部分 if "data/" in original_path: data_index = original_path.find("data/") - rel_path = original_path[data_index:] # data/{kb_name}/graph/{file_name} - resolved_path = os.path.normpath(os.path.join(dataset_folder_abs, rel_path)) + rel_path = original_path[ + data_index: + ] # data/{kb_name}/graph/{file_name} + resolved_path = os.path.normpath( + os.path.join(dataset_folder_abs, rel_path) + ) else: # 如果没有data/，直接使用标准路径 - resolved_path = os.path.join(dataset_folder_abs, "data", kb_name, "graph", target_file_name) + resolved_path = os.path.join( + dataset_folder_abs, + "data", + kb_name, + "graph", + target_file_name, + ) else: resolved_path = original_path - + # 检查路径是否存在 if os.path.exists(resolved_path): return resolved_path @@ -291,12 +352,14 @@ def resolve_file_path(original_path, target_file_name): logger.warning(f" - dataset_folder_abs: {dataset_folder_abs}") logger.warning(f" - target_file_name: {target_file_name}") # 尝试列出目录内容以帮助调试 - graph_dir = os.path.join(dataset_folder_abs, "data", kb_name, "graph") + graph_dir = os.path.join( + dataset_folder_abs, "data", kb_name, "graph" + ) if os.path.exists(graph_dir): files_in_dir = os.listdir(graph_dir) logger.warning(f" - graph目录下的文件: {files_in_dir}") return None - + # 先检查edge文件 if "edge" in schema and len(schema["edge"]) > 0: edge_original_path = schema["edge"][0].get("original_path", "") @@ -304,7 +367,7 @@ def resolve_file_path(original_path, target_file_name): if resolved_path: file_path = resolved_path break - + # 再检查vertex文件（即使edge文件匹配但路径不存在，也要继续检查） if "vertex" in schema and len(schema["vertex"]) > 0: vertex_original_path = schema["vertex"][0].get("original_path", "") @@ -312,38 +375,47 @@ def resolve_file_path(original_path, target_file_name): if resolved_path: file_path = resolved_path break - + if not file_path or not os.path.exists(file_path): # 添加详细的错误信息用于调试 - logger.error(f"文件未找到: kb_id={kb_id}, kb_name={kb_name}, file_name={file_name}, file_path={file_path}") - if kb_type == 'graph' and 'dataset_folder_abs' in locals(): + logger.error( + f"文件未找到: kb_id={kb_id}, kb_name={kb_name}, file_name={file_name}, file_path={file_path}" + ) + if kb_type == "graph" and "dataset_folder_abs" in locals(): logger.error(f"dataset_folder_abs: {dataset_folder_abs}") # 列出可能的路径 - possible_path = os.path.join(dataset_folder_abs, "data", kb_name, "graph", file_name) - logger.error(f"尝试的标准路径: {possible_path}, 存在: {os.path.exists(possible_path)}") + possible_path = os.path.join( + dataset_folder_abs, "data", kb_name, "graph", file_name + ) + logger.error( + f"尝试的标准路径: {possible_path}, 存在: {os.path.exists(possible_path)}" + ) # 列出graph目录下的所有文件 graph_dir = os.path.join(dataset_folder_abs, "data", kb_name, "graph") if os.path.exists(graph_dir): files_in_dir = os.listdir(graph_dir) logger.error(f"graph目录下的文件: {files_in_dir}") - return jsonify({ - 'success': False, - 'error': f'File {file_name} not found. Path: {file_path}' - }), 404 - + return jsonify( + { + "success": False, + "error": f"File {file_name} not found. Path: {file_path}", + } + ), 404 + # 获取文件扩展名 - file_extension = file_name.split('.').pop().lower() - + file_extension = file_name.split(".").pop().lower() + # 根据文件类型读取内容 content = None - content_type = 'text' - - if file_extension == 'csv': + content_type = "text" + + if file_extension == "csv": # CSV文件：读取为表格数据 import csv + try: rows = [] - with open(file_path, 'r', encoding='utf-8') as f: + with open(file_path, "r", encoding="utf-8") as f: reader = csv.DictReader(f) # 获取表头 headers = reader.fieldnames @@ -352,19 +424,20 @@ def resolve_file_path(original_path, target_file_name): if i >= 100: break rows.append(row) - + content = { - 'headers': headers, - 'rows': rows, - 'total_rows': sum(1 for _ in open(file_path, 'r', encoding='utf-8')) - 1, # 减去header行 - 'displayed_rows': len(rows) + "headers": headers, + "rows": rows, + "total_rows": sum(1 for _ in open(file_path, "r", encoding="utf-8")) + - 1, # 减去header行 + "displayed_rows": len(rows), } - content_type = 'csv' + content_type = "csv" except UnicodeDecodeError: # 如果UTF-8失败，尝试其他编码 try: rows = [] - with open(file_path, 'r', encoding='gbk') as f: + with open(file_path, "r", encoding="gbk") as f: reader = csv.DictReader(f) headers = reader.fieldnames for i, row in enumerate(reader): @@ -372,58 +445,62 @@ def resolve_file_path(original_path, target_file_name): break rows.append(row) content = { - 'headers': headers, - 'rows': rows, - 'total_rows': sum(1 for _ in open(file_path, 'r', encoding='gbk')) - 1, - 'displayed_rows': len(rows) + "headers": headers, + "rows": rows, + "total_rows": sum( + 1 for _ in open(file_path, "r", encoding="gbk") + ) + - 1, + "displayed_rows": len(rows), } - content_type = 'csv' + content_type = "csv" except Exception as e: logger.error(f"读取CSV文件失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to read CSV file: {str(e)}' - }), 500 + return jsonify( + { + "success": False, + "error": f"Failed to read CSV file: {str(e)}", + } + ), 500 except Exception as e: logger.error(f"读取CSV文件失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to read CSV file: {str(e)}' - }), 500 - elif file_extension == 'txt': + return jsonify( + {"success": False, "error": f"Failed to read CSV file: {str(e)}"} + ), 500 + elif file_extension == "txt": # 直接读取文本文件 try: - with open(file_path, 'r', encoding='utf-8') as f: + with open(file_path, "r", encoding="utf-8") as f: content = f.read() except UnicodeDecodeError: # 如果UTF-8失败，尝试其他编码 try: - with open(file_path, 'r', encoding='gbk') as f: + with open(file_path, "r", encoding="gbk") as f: content = f.read() except: - with open(file_path, 'r', encoding='latin-1') as f: + with open(file_path, "r", encoding="latin-1") as f: content = f.read() - elif file_extension == 'pdf': + elif file_extension == "pdf": # PDF文件：返回base64编码的内容，前端使用PDF.js显示 try: - with open(file_path, 'rb') as f: + with open(file_path, "rb") as f: pdf_content = f.read() - content = base64.b64encode(pdf_content).decode('utf-8') - content_type = 'pdf' + content = base64.b64encode(pdf_content).decode("utf-8") + content_type = "pdf" except Exception as e: logger.error(f"读取PDF文件失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to read PDF file: {str(e)}' - }), 500 - elif file_extension == 'docx': + return jsonify( + {"success": False, "error": f"Failed to read PDF file: {str(e)}"} + ), 500 + elif file_extension == "docx": # DOCX文档：使用mammoth转换为HTML格式，保持原有格式 try: import mammoth - with open(file_path, 'rb') as f: + + with open(file_path, "rb") as f: result = mammoth.convert_to_html(f) content = result.value - content_type = 'html' + content_type = "html" # 如果有警告，记录但不影响显示 if result.messages: logger.warning(f"DOCX转换警告: {result.messages}") @@ -431,71 +508,84 @@ def resolve_file_path(original_path, target_file_name): # 如果mammoth未安装，尝试使用markitdown try: from markitdown import MarkItDown + md = MarkItDown() result = md.convert(file_path) # markitdown可能返回markdown，需要转换为HTML import markdown - if hasattr(result, 'html_content') and result.html_content: + + if hasattr(result, "html_content") and result.html_content: content = result.html_content else: # 将markdown转换为HTML - content = markdown.markdown(result.text_content, extensions=['tables', 'fenced_code']) - content_type = 'html' + content = markdown.markdown( + result.text_content, extensions=["tables", "fenced_code"] + ) + content_type = "html" except Exception as e: logger.error(f"转换DOCX文档失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to convert DOCX document: {str(e)}. Please install mammoth: pip install mammoth' - }), 500 + return jsonify( + { + "success": False, + "error": f"Failed to convert DOCX document: {str(e)}. Please install mammoth: pip install mammoth", + } + ), 500 except Exception as e: logger.error(f"转换DOCX文档失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to convert DOCX document: {str(e)}' - }), 500 - elif file_extension == 'doc': + return jsonify( + { + "success": False, + "error": f"Failed to convert DOCX document: {str(e)}", + } + ), 500 + elif file_extension == "doc": # DOC文档（旧格式）：尝试转换为文本或提示用户 try: # 尝试使用markitdown转换 from markitdown import MarkItDown + md = MarkItDown() result = md.convert(file_path) # 将markdown转换为HTML import markdown - if hasattr(result, 'html_content') and result.html_content: + + if hasattr(result, "html_content") and result.html_content: content = result.html_content else: - content = markdown.markdown(result.text_content, extensions=['tables', 'fenced_code']) - content_type = 'html' + content = markdown.markdown( + result.text_content, extensions=["tables", "fenced_code"] + ) + content_type = "html" except Exception as e: logger.error(f"转换DOC文档失败: {str(e)}") # DOC格式较老，可能无法完美转换，返回文本内容 - return jsonify({ - 'success': False, - 'error': f'DOC格式文件预览受限，建议转换为DOCX格式。转换错误: {str(e)}' - }), 500 + return jsonify( + { + "success": False, + "error": f"DOC格式文件预览受限，建议转换为DOCX格式。转换错误: {str(e)}", + } + ), 500 else: - return jsonify({ - 'success': False, - 'error': f'Unsupported file type: {file_extension}' - }), 400 - - return jsonify({ - 'success': True, - 'content': content, - 'content_type': content_type, - 'file_name': file_name, - 'file_path': file_path - }) - + return jsonify( + {"success": False, "error": f"Unsupported file type: {file_extension}"} + ), 400 + + return jsonify( + { + "success": True, + "content": content, + "content_type": content_type, + "file_name": file_name, + "file_path": file_path, + } + ) + except Exception as e: logger.error(f"预览文件处理错误: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 + -@bp.route('/api/upload_file', methods=['POST']) +@bp.route("/api/upload_file", methods=["POST"]) def api_upload_file(): """HTTP API: 文件上传""" # 用于跟踪已保存的文件，以便在发生错误时清理 @@ -508,126 +598,154 @@ def is_duplicate_error(msg: str) -> bool: return False lower_msg = str(msg).lower() return "already exists" in lower_msg or "已存在" in lower_msg + try: # 检查是否为批量上传（图数据集两个文件的情况） - is_batch_upload = request.form.get('is_batch_upload') == 'true' - kb_id = request.form.get('kb_id') - file_type = request.form.get('file_type') - graph_info_json = request.form.get('graph_info') - + is_batch_upload = request.form.get("is_batch_upload") == "true" + kb_id = request.form.get("kb_id") + file_type = request.form.get("file_type") + graph_info_json = request.form.get("graph_info") + if not kb_id: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id' - }), 400 - + return jsonify({"success": False, "error": "Missing kb_id"}), 400 + # 获取知识库名称 kb_name = get_knowledge_base_name(kb_id) - + # 创建目录结构，统一使用后端 DATASETS_DATA_DIR 绝对路径 file_dir = os.path.join(str(DATASETS_DATA_DIR), kb_name, file_type) os.makedirs(file_dir, exist_ok=True) - - if is_batch_upload and file_type == 'graph': + + if is_batch_upload and file_type == "graph": # 批量上传模式：图数据集上传两个文件 - files = request.files.getlist('files') + files = request.files.getlist("files") if not files or len(files) == 0: - return jsonify({ - 'success': False, - 'error': 'Missing files' - }), 400 - + return jsonify({"success": False, "error": "Missing files"}), 400 + # 先保存所有文件 for file in files: file_path = os.path.join(file_dir, file.filename) file.save(file_path) file_size = os.path.getsize(file_path) - saved_files.append({ - 'filename': file.filename, - 'path': file_path, - 'size': file_size - }) - + saved_files.append( + {"filename": file.filename, "path": file_path, "size": file_size} + ) + # 所有文件保存完成后，再调用DocumentAPI try: graph_info = json.loads(graph_info_json) except (json.JSONDecodeError, TypeError) as e: # JSON解析失败，删除已保存的文件 for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as del_e: - logger.error(f"删除文件失败 {saved_file['path']}: {str(del_e)}") - return jsonify({ - 'success': False, - 'error': f'Invalid graph_info JSON: {str(e)}' - }), 400 + logger.error( + f"删除文件失败 {saved_file['path']}: {str(del_e)}" + ) + return jsonify( + {"success": False, "error": f"Invalid graph_info JSON: {str(e)}"} + ), 400 graph_name = graph_info.get("graphName") # 提取顶点文件相关字段 - vertex_file_name = graph_info["vertexSchema"].get("fileName", "") # 顶点文件名（可选） - vertex_id_field = graph_info["vertexSchema"].get("idField", "") # 顶点ID字段（如果有顶点文件则必填） + vertex_file_name = graph_info["vertexSchema"].get( + "fileName", "" + ) # 顶点文件名（可选） + vertex_id_field = graph_info["vertexSchema"].get( + "idField", "" + ) # 顶点ID字段（如果有顶点文件则必填） vertex_name_field = graph_info["vertexSchema"].get("nameField", "") - vertex_name_field = vertex_id_field if not vertex_name_field else vertex_name_field # 顶点名称字段（如果有顶点文件则必填） - + vertex_name_field = ( + vertex_id_field if not vertex_name_field else vertex_name_field + ) # 顶点名称字段（如果有顶点文件则必填） + # 提取边文件相关字段 edge_file_name = graph_info["edgeSchema"]["fileName"] # 边文件名（必填） - edge_source_field = graph_info["edgeSchema"]["sourceField"] # 源节点字段（必填） - edge_target_field = graph_info["edgeSchema"]["targetField"] # 目标节点字段（必填） - edge_relation_field = graph_info["edgeSchema"].get("relationField", "") # 关系字段（必填） - weight_col = graph_info["edgeSchema"].get("weightField", "") # 权重字段（可选） - is_directed = graph_info["graphProperties"].get("isDirected", True) # 是否为有向图 - vertex_attribute_field = format_string_to_list(graph_info["vertexSchema"].get("propertiesField", "")) # 顶点属性字段（可选），转为后端期望的列表字面量字符串 - vertex_query_field = graph_info["vertexSchema"].get("nameField", "") # 顶点查询字段（可选） + edge_source_field = graph_info["edgeSchema"][ + "sourceField" + ] # 源节点字段（必填） + edge_target_field = graph_info["edgeSchema"][ + "targetField" + ] # 目标节点字段（必填） + edge_relation_field = graph_info["edgeSchema"].get( + "relationField", "" + ) # 关系字段（必填） + weight_col = graph_info["edgeSchema"].get( + "weightField", "" + ) # 权重字段（可选） + is_directed = graph_info["graphProperties"].get( + "isDirected", True + ) # 是否为有向图 + vertex_attribute_field = format_string_to_list( + graph_info["vertexSchema"].get("propertiesField", "") + ) # 顶点属性字段（可选），转为后端期望的列表字面量字符串 + vertex_query_field = graph_info["vertexSchema"].get( + "nameField", "" + ) # 顶点查询字段（可选） # 验证必填字段 if not edge_file_name or not edge_source_field or not edge_target_field: # 删除已保存的文件 for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as del_e: - logger.error(f"删除文件失败 {saved_file['path']}: {str(del_e)}") - return jsonify({ - 'success': False, - 'error': 'Edge file name, source field, target field, and are required for Graph dataset' - }), 400 - + logger.error( + f"删除文件失败 {saved_file['path']}: {str(del_e)}" + ) + return jsonify( + { + "success": False, + "error": "Edge file name, source field, target field, and are required for Graph dataset", + } + ), 400 + # 如果有顶点文件，验证顶点文件相关字段 if vertex_file_name and (not vertex_id_field or not vertex_name_field): # 删除已保存的文件 for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as del_e: - logger.error(f"删除文件失败 {saved_file['path']}: {str(del_e)}") - return jsonify({ - 'success': False, - 'error': 'If vertex file is provided, vertex ID field is required' - }), 400 - + logger.error( + f"删除文件失败 {saved_file['path']}: {str(del_e)}" + ) + return jsonify( + { + "success": False, + "error": "If vertex file is provided, vertex ID field is required", + } + ), 400 + vertex_name_field = format_string_to_list(vertex_name_field) edge_relation_field = format_string_to_list(edge_relation_field) - msg111 = json.dumps({ - "action": "upload_file", - "file_name": edge_file_name, - "ds_name": kb_name, - "vertex_file_name": vertex_file_name if vertex_file_name else None, # 只有提供了文件名才发送 - "vertex_id_field": vertex_id_field if vertex_file_name else None, - "vertex_name_field": vertex_name_field if vertex_file_name else None, - "source_field": edge_source_field, - "target_field": edge_target_field, - "relation_field": edge_relation_field, # 使用 relation_field 而不是 edge_relation_field - "is_directed": is_directed, - #"weight_field": weight_col if weight_col else None - "vertex_attribute_field": vertex_attribute_field, - "vertex_query_field": vertex_query_field - }) + msg111 = json.dumps( + { + "action": "upload_file", + "file_name": edge_file_name, + "ds_name": kb_name, + "vertex_file_name": vertex_file_name + if vertex_file_name + else None, # 只有提供了文件名才发送 + "vertex_id_field": vertex_id_field if vertex_file_name else None, + "vertex_name_field": vertex_name_field + if vertex_file_name + else None, + "source_field": edge_source_field, + "target_field": edge_target_field, + "relation_field": edge_relation_field, # 使用 relation_field 而不是 edge_relation_field + "is_directed": is_directed, + # "weight_field": weight_col if weight_col else None + "vertex_attribute_field": vertex_attribute_field, + "vertex_query_field": vertex_query_field, + } + ) a25 = MessageCollectingSocket(msg111) - asyncio.run(server_Test.handler(a25)) + _run_async(server_Test.handler(a25)) acps = a25.returnmsg - # 如果有错误，优先使用错误消息 + # 如果有错误，优先使用错误消息 if a25.has_error: try: acps = json.loads(a25.error_message) @@ -637,72 +755,75 @@ def is_duplicate_error(msg: str) -> bool: # 非重复错误才清理文件，重复文件保持 if not is_duplicate_error(error_content): for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as e: - logger.error(f"删除文件失败 {saved_file['path']}: {str(e)}") - return jsonify({ - 'success': False, - 'error': error_content - }), 400 + logger.error( + f"删除文件失败 {saved_file['path']}: {str(e)}" + ) + return jsonify({"success": False, "error": error_content}), 400 acps = json.loads(acps) if acps["type"] == "error": error_content = acps.get("content", "上传失败") # 非重复错误才清理文件 if not is_duplicate_error(error_content): for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as e: - logger.error(f"删除文件失败 {saved_file['path']}: {str(e)}") - return jsonify({ - 'success': False, - 'error': error_content - }), 400 - + logger.error( + f"删除文件失败 {saved_file['path']}: {str(e)}" + ) + return jsonify({"success": False, "error": error_content}), 400 + # 返回所有文件的信息 files_info = [] for saved_file in saved_files: - files_info.append({ - 'file_name': saved_file['filename'], - 'file_path': saved_file['path'], - 'file_size': saved_file['size'], - 'upload_time': datetime.now().isoformat() - }) - - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'files': files_info, - 'message': f'Files uploaded successfully' - }) + files_info.append( + { + "file_name": saved_file["filename"], + "file_path": saved_file["path"], + "file_size": saved_file["size"], + "upload_time": datetime.now().isoformat(), + } + ) + + return jsonify( + { + "success": True, + "kb_id": kb_id, + "files": files_info, + "message": f"Files uploaded successfully", + } + ) else: # 单个文件上传模式（原有逻辑） - file = request.files.get('file') + file = request.files.get("file") if not file: - return jsonify({ - 'success': False, - 'error': 'Missing file' - }), 400 - + return jsonify({"success": False, "error": "Missing file"}), 400 + # 保存文件 file_path = os.path.join(file_dir, file.filename) file.save(file_path) saved_file_path = file_path # 记录已保存的文件路径，用于错误清理 - + # 获取文件大小 file_size = os.path.getsize(file_path) - - if file_type == 'text': + + if file_type == "text": # 使用可捕获错误信息的socket - a25 = MessageCollectingSocket(json.dumps({ - "action": "upload_file", - "file_name": file.filename, - "ds_name": kb_name - })) - asyncio.run(server_Test.handler(a25)) + a25 = MessageCollectingSocket( + json.dumps( + { + "action": "upload_file", + "file_name": file.filename, + "ds_name": kb_name, + } + ) + ) + _run_async(server_Test.handler(a25)) acps = a25.returnmsg # 如果有错误，优先使用错误消息 if a25.has_error: @@ -711,17 +832,14 @@ def is_duplicate_error(msg: str) -> bool: error_content = acps.get("content", "上传失败") except: error_content = "上传失败，请检查文件格式和内容" - # 非重复错误才删除文件；重复错误保留（覆盖的文件） + # 非重复错误才删除文件；重复错误保留（覆盖的文件） if not is_duplicate_error(error_content): if os.path.exists(file_path): try: os.remove(file_path) except Exception as e: logger.error(f"删除文件失败 {file_path}: {str(e)}") - return jsonify({ - 'success': False, - 'error': error_content - }), 400 + return jsonify({"success": False, "error": error_content}), 400 # 尝试解析返回结果并检查错误类型 try: acps_json = json.loads(acps) @@ -735,10 +853,7 @@ def is_duplicate_error(msg: str) -> bool: os.remove(file_path) except Exception as e: logger.error(f"删除文件失败 {file_path}: {str(e)}") - return jsonify({ - 'success': False, - 'error': error_content - }), 400 + return jsonify({"success": False, "error": error_content}), 400 else: # 如果图数据上传，还需要传输配置文件 try: @@ -749,59 +864,93 @@ def is_duplicate_error(msg: str) -> bool: try: os.remove(saved_file_path) except Exception as del_e: - logger.error(f"删除文件失败 {saved_file_path}: {str(del_e)}") - return jsonify({ - 'success': False, - 'error': f'Invalid graph_info JSON: {str(e)}' - }), 400 + logger.error( + f"删除文件失败 {saved_file_path}: {str(del_e)}" + ) + return jsonify( + { + "success": False, + "error": f"Invalid graph_info JSON: {str(e)}", + } + ), 400 graph_name = graph_info.get("graphName") # 提取顶点文件相关字段 - vertex_file_name = graph_info["vertexSchema"].get("fileName", "") # 顶点文件名（可选） - vertex_id_field = graph_info["vertexSchema"].get("idField", "") # 顶点ID字段（如果有顶点文件则必填） - vertex_name_field = graph_info["vertexSchema"].get("nameField", "") # 顶点名称字段（如果有顶点文件则必填） - + vertex_file_name = graph_info["vertexSchema"].get( + "fileName", "" + ) # 顶点文件名（可选） + vertex_id_field = graph_info["vertexSchema"].get( + "idField", "" + ) # 顶点ID字段（如果有顶点文件则必填） + vertex_name_field = graph_info["vertexSchema"].get( + "nameField", "" + ) # 顶点名称字段（如果有顶点文件则必填） + # 提取边文件相关字段 - edge_file_name = graph_info["edgeSchema"]["fileName"] # 边文件名（必填） - edge_source_field = graph_info["edgeSchema"]["sourceField"] # 源节点字段（必填） - edge_target_field = graph_info["edgeSchema"]["targetField"] # 目标节点字段（必填） - edge_relation_field = graph_info["edgeSchema"].get("relationField", "") # 关系字段（必填） - weight_col = graph_info["edgeSchema"].get("weightField", "") # 权重字段（可选） - is_directed = graph_info["graphProperties"].get("isDirected", True) # 是否为有向图 + edge_file_name = graph_info["edgeSchema"][ + "fileName" + ] # 边文件名（必填） + edge_source_field = graph_info["edgeSchema"][ + "sourceField" + ] # 源节点字段（必填） + edge_target_field = graph_info["edgeSchema"][ + "targetField" + ] # 目标节点字段（必填） + edge_relation_field = graph_info["edgeSchema"].get( + "relationField", "" + ) # 关系字段（必填） + weight_col = graph_info["edgeSchema"].get( + "weightField", "" + ) # 权重字段（可选） + is_directed = graph_info["graphProperties"].get( + "isDirected", True + ) # 是否为有向图 # 验证必填字段 if not edge_file_name or not edge_source_field or not edge_target_field: if os.path.exists(file_path): os.remove(file_path) - return jsonify({ - 'success': False, - 'error': 'Edge file name, source field, target field, and are required for Graph dataset' - }), 400 - + return jsonify( + { + "success": False, + "error": "Edge file name, source field, target field, and are required for Graph dataset", + } + ), 400 + # 如果有顶点文件，验证顶点文件相关字段 if vertex_file_name and (not vertex_id_field or not vertex_name_field): if os.path.exists(file_path): os.remove(file_path) - return jsonify({ - 'success': False, - 'error': 'If vertex file is provided, vertex ID field and name field are required' - }), 400 + return jsonify( + { + "success": False, + "error": "If vertex file is provided, vertex ID field and name field are required", + } + ), 400 vertex_name_field = format_string_to_list(vertex_name_field) edge_relation_field = format_string_to_list(edge_relation_field) - msg111 = json.dumps({ - "action": "upload_file", - "file_name": edge_file_name, - "ds_name": kb_name, - "vertex_file_name": vertex_file_name if vertex_file_name else None, # 只有提供了文件名才发送 - "vertex_id_field": vertex_id_field if vertex_file_name else None, - "vertex_name_field": vertex_name_field if vertex_file_name else None, - "source_field": edge_source_field, - "target_field": edge_target_field, - "relation_field": edge_relation_field, # 使用 relation_field 而不是 edge_relation_field - "is_directed": is_directed, - #"weight_field": weight_col if weight_col else None - }) + msg111 = json.dumps( + { + "action": "upload_file", + "file_name": edge_file_name, + "ds_name": kb_name, + "vertex_file_name": vertex_file_name + if vertex_file_name + else None, # 只有提供了文件名才发送 + "vertex_id_field": vertex_id_field + if vertex_file_name + else None, + "vertex_name_field": vertex_name_field + if vertex_file_name + else None, + "source_field": edge_source_field, + "target_field": edge_target_field, + "relation_field": edge_relation_field, # 使用 relation_field 而不是 edge_relation_field + "is_directed": is_directed, + # "weight_field": weight_col if weight_col else None + } + ) a25 = MessageCollectingSocket(msg111) - asyncio.run(server_Test.handler(a25)) + _run_async(server_Test.handler(a25)) acps = a25.returnmsg # 如果有错误，优先使用错误消息 if a25.has_error: @@ -815,10 +964,7 @@ def is_duplicate_error(msg: str) -> bool: os.remove(file_path) except Exception as e: logger.error(f"删除文件失败 {file_path}: {str(e)}") - return jsonify({ - 'success': False, - 'error': error_content - }), 400 + return jsonify({"success": False, "error": error_content}), 400 acps = json.loads(acps) if acps["type"] == "error": if os.path.exists(file_path): @@ -826,42 +972,42 @@ def is_duplicate_error(msg: str) -> bool: os.remove(file_path) except Exception as e: logger.error(f"删除文件失败 {file_path}: {str(e)}") - return jsonify({ - 'success': False, - 'error': acps.get("content", "上传失败") - }), 400 - - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'file_name': file.filename, - 'file_path': file_path, - 'file_size': file_size, - 'upload_time': datetime.now().isoformat(), - 'message': f'File {file.filename} uploaded successfully' - }) - + return jsonify( + {"success": False, "error": acps.get("content", "上传失败")} + ), 400 + + return jsonify( + { + "success": True, + "kb_id": kb_id, + "file_name": file.filename, + "file_path": file_path, + "file_size": file_size, + "upload_time": datetime.now().isoformat(), + "message": f"File {file.filename} uploaded successfully", + } + ) + except Exception as e: logger.error(f"文件上传处理错误: {str(e)}") # 清理已保存的文件 # 批量上传的情况 if saved_files: for saved_file in saved_files: - if os.path.exists(saved_file['path']): + if os.path.exists(saved_file["path"]): try: - os.remove(saved_file['path']) + os.remove(saved_file["path"]) except Exception as del_e: - logger.error(f"异常处理：删除文件失败 {saved_file['path']}: {str(del_e)}") + logger.error( + f"异常处理：删除文件失败 {saved_file['path']}: {str(del_e)}" + ) # 单个文件上传的情况 elif saved_file_path and os.path.exists(saved_file_path): try: os.remove(saved_file_path) except Exception as del_e: logger.error(f"异常处理：删除文件失败 {saved_file_path}: {str(del_e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 # def process_graph_file(kb_name, file_name, graph_info): @@ -870,13 +1016,13 @@ def is_duplicate_error(msg: str) -> bool: # # 构建文件路径 # file_dir = os.path.join(UPLOAD_BASE_DIR, kb_name, file_name) # file_path = os.path.join(file_dir, file_name) - + # # 这里调用你的图数据处理逻辑 # # 示例：调用DocumentAPI处理图数据 # # 请根据你的实际需求实现这个函数 # # from your_module import process_graph_data # # process_graph_data(kb_name, file_path, graph_info) - + # except Exception as e: # logger.error(f"处理图数据文件失败: {str(e)}") @@ -887,151 +1033,154 @@ def is_duplicate_error(msg: str) -> bool: # # 构建文件路径 # file_dir = os.path.join(UPLOAD_BASE_DIR, kb_name, file_name) # file_path = os.path.join(file_dir, file_name) - + # # 这里调用你的文本处理逻辑 # # 示例：调用DocumentAPI处理文本数据 # # 请根据你的实际需求实现这个函数 # # from your_module import process_text_data # # process_text_data(kb_name, file_path) - + # except Exception as e: # logger.error(f"处理文本文件失败: {str(e)}") + # API路由保持不变 -@bp.route('/api/knowledge_bases1', methods=['GET']) +@bp.route("/api/knowledge_bases1", methods=["GET"]) def get_knowledge_bases1(): """获取数据集列表（从JSON文件读取文件个数）""" try: a1 = DummySocket(json.dumps({"action": "get_datasets"})) - asyncio.run(server_Test.handler(a1)) + _run_async(server_Test.handler(a1)) gkb = a1.returnmsg knowledge_bases = json.loads(gkb) for kb in knowledge_bases["content"]["data"]: if kb.get("file_type") == "graph" and kb.get("file_count") == 1: try: - a6 = DummySocket(json.dumps({"action": "get_dataset_schema","ds_name":kb["name"]})) - asyncio.run(server_Test.handler(a6)) + a6 = DummySocket( + json.dumps( + {"action": "get_dataset_schema", "ds_name": kb["name"]} + ) + ) + _run_async(server_Test.handler(a6)) gkb6 = a6.returnmsg gkb6_json = json.loads(gkb6) data_list = gkb6_json.get("content", {}).get("data", []) - if isinstance(data_list, list) and data_list and data_list[0].get("vertex_file",None) is not None: + if ( + isinstance(data_list, list) + and data_list + and data_list[0].get("vertex_file", None) is not None + ): kb["file_count"] = 2 except Exception as inner_e: logger.warning(f"统计图文件数时出错，已忽略: {inner_e}") - return jsonify({ - "success": True, - "data": knowledge_bases["content"]["data"], - "count": len(knowledge_bases["content"]["data"]) - }) - + return jsonify( + { + "success": True, + "data": knowledge_bases["content"]["data"], + "count": len(knowledge_bases["content"]["data"]), + } + ) + except Exception as e: logger.error(f"获取知识库列表失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 -@bp.route('/api/dataset_type', methods=['GET']) + +@bp.route("/api/dataset_type", methods=["GET"]) def get_dataset_type(): """获取数据集类型 - 从JSON文件动态获取""" - kb_id = request.args.get('kb_id') - + kb_id = request.args.get("kb_id") + if not kb_id: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id parameter' - }), 400 - + return jsonify({"success": False, "error": "Missing kb_id parameter"}), 400 + try: # kb_id 必须可转换为数字，否则直接返回 400，避免返回伪类型 raw_kb_id = kb_id if isinstance(kb_id, str) and kb_id.startswith("kb_"): kb_id = kb_id[3:] kb_id_int = int(kb_id) - + # 从JSON文件获取真实的数据集类型 dataset_type = get_dataset_type_from_back(kb_id_int) - - return jsonify({ - 'success': True, - 'dataset_type': dataset_type - }) + + return jsonify({"success": True, "dataset_type": dataset_type}) except Exception as e: logger.error(f"获取数据集类型失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Invalid kb_id: {raw_kb_id}' - }), 500 + return jsonify({"success": False, "error": f"Invalid kb_id: {raw_kb_id}"}), 500 -@bp.route('/api/delete_file', methods=['GET']) +@bp.route("/api/delete_file", methods=["GET"]) def api_delete_file(): """HTTP API: 删除文件""" try: - kb_id = request.args.get('kb_id') - file_name = request.args.get('file_name') - + kb_id = request.args.get("kb_id") + file_name = request.args.get("file_name") + if not kb_id or not file_name: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id or file_name' - }), 400 - + return jsonify( + {"success": False, "error": "Missing kb_id or file_name"} + ), 400 + # 获取知识库名称 kb_name = get_knowledge_base_name(kb_id) kb_type = get_dataset_type_from_back(kb_id) - if kb_type == 'graph': - a26 = DummySocket(json.dumps({"action": "get_dataset_schema","ds_name": kb_name})) - asyncio.run(server_Test.handler(a26)) + if kb_type == "graph": + a26 = DummySocket( + json.dumps({"action": "get_dataset_schema", "ds_name": kb_name}) + ) + _run_async(server_Test.handler(a26)) acps = a26.returnmsg acps = json.loads(acps) file_name = acps["content"]["data"][0]["name"] try: - socketserver2 = DummySocket(json.dumps( - {"action":"delete_file", - "file_name":file_name, - "ds_name":kb_name} - )) - asyncio.run(server_Test.handler(socketserver2)) - - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'file_name': file_name, - 'message': f'File {file_name} deleted successfully' - }) - + socketserver2 = DummySocket( + json.dumps( + { + "action": "delete_file", + "file_name": file_name, + "ds_name": kb_name, + } + ) + ) + _run_async(server_Test.handler(socketserver2)) + + return jsonify( + { + "success": True, + "kb_id": kb_id, + "file_name": file_name, + "message": f"File {file_name} deleted successfully", + } + ) + except Exception as e: logger.error(f"删除文件夹失败: {str(e)}") - return jsonify({ - 'success': False, - 'error': f'Failed to delete directory: {str(e)}' - }), 500 - + return jsonify( + {"success": False, "error": f"Failed to delete directory: {str(e)}"} + ), 500 + except Exception as e: logger.error(f"删除文件处理错误: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 -@bp.route('/api/parse_control', methods=['GET']) + +@bp.route("/api/parse_control", methods=["GET"]) def api_parse_control(): """HTTP API: 文件解析控制""" try: - kb_id = request.args.get('kb_id') - file_name = request.args.get('file_name') - action = request.args.get('action') # 'start' 或 'pause' - + kb_id = request.args.get("kb_id") + file_name = request.args.get("file_name") + action = request.args.get("action") # 'start' 或 'pause' + if not kb_id or not file_name or not action: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id, file_name or action' - }), 400 - + return jsonify( + {"success": False, "error": "Missing kb_id, file_name or action"} + ), 400 + # 获取知识库名称 kb_name = get_knowledge_base_name(kb_id) - + # # 构建文件路径 # kb_dir = os.path.join(UPLOAD_BASE_DIR, kb_name) # file_dir = os.path.join(kb_dir, file_name) @@ -1042,21 +1191,23 @@ def api_parse_control(): # 'error': f'File not found: {file_path}' # }), 404 # 使用 MessageCollectingSocket 捕获 DocumentAPI 返回的 error 消息 - socketserver = MessageCollectingSocket(json.dumps( - { - "action": "parsing_file", - "file_name": file_name, - "ds_name": kb_name, - "type": "ollama", - "api_key": "YOUR_API_KEY", - "base_url": "YOUR_BASE_URL", - "llm_name": "llama3.1:70b", - "mode":"single", - "thread_count": 1, - "chunk_size": 3072 - } - )) - asyncio.run(server_Test.handler(socketserver)) + socketserver = MessageCollectingSocket( + json.dumps( + { + "action": "parsing_file", + "file_name": file_name, + "ds_name": kb_name, + "type": "ollama", + "api_key": os.getenv("DEFAULT_API_KEY", "YOUR_API_KEY"), + "base_url": os.getenv("DEFAULT_BASE_URL", "YOUR_BASE_URL"), + "llm_name": "llama3.1:70b", + "mode": "single", + "thread_count": 1, + "chunk_size": 3072, + } + ) + ) + _run_async(server_Test.handler(socketserver)) # 优先使用收集到的 error 消息 raw_msg = socketserver.returnmsg @@ -1066,65 +1217,67 @@ def api_parse_control(): parsed_msg = {} # 如果 DocumentAPI 返回的是 error，则把详细信息透传给前端 if isinstance(parsed_msg, dict) and parsed_msg.get("type") == "error": - a6 = DummySocket(json.dumps({"action": "schema_refine","ds_name": kb_name,"file_name": file_name})) - asyncio.run(server_Test.handler(a6)) - #acps = a6.returnmsg - #acps = json.loads(acps) + a6 = DummySocket( + json.dumps( + { + "action": "schema_refine", + "ds_name": kb_name, + "file_name": file_name, + } + ) + ) + _run_async(server_Test.handler(a6)) + # acps = a6.returnmsg + # acps = json.loads(acps) # content 中已经包含“添加失败 ❌ ... 未找到 JSON 内容，原始响应: ...” error_content = parsed_msg.get("content") or "解析失败，请检查文件内容" logger.error(f"在数据集{kb_name}中解析文件{file_name}失败: {error_content}") - return jsonify({ - "success": False, - "error": error_content - }), 400 + return jsonify({"success": False, "error": error_content}), 400 # ============================================ # 在这里添加您的解析逻辑 # 参数说明: # - kb_id: 知识库ID - # - kb_name: 知识库名称 + # - kb_name: 知识库名称 # - file_name: 文件名 # - file_path: 文件完整路径 # - action: 'start' 或 'pause' # ============================================ - + # 这里可以调用您的 DocumentAPI 或其他处理逻辑 # 例如： # from your_module import your_parse_function # your_parse_function(kb_id, kb_name, file_name, file_path, action) - + # 返回成功响应 - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'file_name': file_name, - 'action': action, - 'message': f'解析请求已接收: {action} {file_name}' - }) - + return jsonify( + { + "success": True, + "kb_id": kb_id, + "file_name": file_name, + "action": action, + "message": f"解析请求已接收: {action} {file_name}", + } + ) + except Exception as e: logger.error(f"解析控制处理错误: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 + -@bp.route('/api/generate_graph', methods=['POST']) +@bp.route("/api/generate_graph", methods=["POST"]) def api_generate_graph(): """HTTP API: 生成知识图谱""" try: data = request.get_json() - kb_id = data.get('kb_id') - kb_name = data.get('kb_name') - files = data.get('files', []) + kb_id = data.get("kb_id") + kb_name = data.get("kb_name") + files = data.get("files", []) dataset_type = get_dataset_type_from_back(kb_id) - graph_name1 = files[0].split('.')[0] + graph_name1 = files[0].split(".")[0] if not kb_id: - return jsonify({ - 'success': False, - 'error': 'Missing kb_id' - }), 400 - + return jsonify({"success": False, "error": "Missing kb_id"}), 400 + # 首先检查所有文件是否已解析 # #parsing_status = get_all_files_parsed_status(kb_id, files) # if not parsing_status['all_parsed']: @@ -1134,66 +1287,73 @@ def api_generate_graph(): # 'error': f'以下文件尚未解析: {", ".join(unparsed_files)}', # 'unparsed_files': unparsed_files # }), 400 - + # 获取知识库名称 kb_real_name = get_knowledge_base_name(kb_id) - + # ============================================ # 在这里调用您的三元组生成逻辑 # ============================================ # 这里是您需要实现的三元组生成逻辑 # 请替换下面的示例代码 # ============================================ - if dataset_type == 'text': + if dataset_type == "text": is_directed = True - msg12 = json.dumps({"action": "get_overall_triplets", - #"graph_name": graph_name1, - "ds_name": kb_name}) + msg12 = json.dumps( + { + "action": "get_overall_triplets", + # "graph_name": graph_name1, + "ds_name": kb_name, + } + ) a33 = DummySocket(msg12) - asyncio.run(server_Test.handler(a33)) + _run_async(server_Test.handler(a33)) agg = a33.returnmsg agg1 = json.loads(agg) triplets = agg1["content"]["data"] if len(triplets) > 500: triplets = triplets[:500] - elif dataset_type == 'graph': - msg12 = json.dumps({"action": "get_file_triplets_from_graph_dataset", + elif dataset_type == "graph": + msg12 = json.dumps( + { + "action": "get_file_triplets_from_graph_dataset", "file_name": files[0], - "ds_name": kb_name}) + "ds_name": kb_name, + } + ) a33 = DummySocket(msg12) - asyncio.run(server_Test.handler(a33)) + _run_async(server_Test.handler(a33)) agg = a33.returnmsg agg1 = json.loads(agg) triplets = agg1["content"]["data"]["edges"] - is_directed = agg1["content"]["data"].get("is_directed",True) + is_directed = agg1["content"]["data"].get("is_directed", True) if len(triplets) > 500: triplets = triplets[:500] - # 示例三元组数据 - 请替换为您的实际逻辑 + # 示例三元组数据 - 请替换为您的实际逻辑 # ============================================ # 三元组生成逻辑结束 # ============================================ - + # 返回成功响应 - return jsonify({ - 'success': True, - 'kb_id': kb_id, - 'kb_name': kb_name, - 'file_count': len(files), - 'triplets': triplets, - 'is_directed': is_directed, - 'message': f'成功生成知识图谱，包含 {len(triplets)} 个三元组' - }) - + return jsonify( + { + "success": True, + "kb_id": kb_id, + "kb_name": kb_name, + "file_count": len(files), + "triplets": triplets, + "is_directed": is_directed, + "message": f"成功生成知识图谱，包含 {len(triplets)} 个三元组", + } + ) + except Exception as e: logger.error(f"生成图谱处理错误: {str(e)}") - return jsonify({ - 'success': False, - 'error': str(e) - }), 500 + return jsonify({"success": False, "error": str(e)}), 500 -#async def main(): - #from aag.api.DocumentAPI import server_Test +# async def main(): +# from aag.api.DocumentAPI import server_Test -#asyncio.run(main()) +# asyncio.run(main()) diff --git a/web/frontend/route/sockets_chat.py b/web/frontend/route/sockets_chat.py index 9207365..f2d2055 100644 --- a/web/frontend/route/sockets_chat.py +++ b/web/frontend/route/sockets_chat.py @@ -3,15 +3,67 @@ import re import logging import asyncio -from flask_socketio import emit +from flask import request +from flask_socketio import emit, disconnect from . import socketio + # Add project path to import AAG services -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))) +sys.path.append( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + ) +) from aag.api.async_runtime import get_background_loop, get_chat_service - +from aag.api.services.chat_service import CHAT_FRIENDLY_ERROR_MSG logger = logging.getLogger(__name__) +# WebSocket 连接跟踪：记录活跃连接的 session ID +_active_connections: set = set() + + +@socketio.on("connect") +def on_connect(): + """WebSocket 连接认证：验证 token 后才允许连接""" + expected_token = os.getenv("YIGRAPH_WS_TOKEN", "") + + # 未设置环境变量则跳过验证（开发环境兼容） + if not expected_token: + logger.warning( + "YIGRAPH_WS_TOKEN 未设置，WebSocket 连接跳过认证（仅建议开发环境使用）" + ) + _active_connections.add(request.sid) + logger.info( + f"WebSocket 客户端已连接（无认证）: {request.sid}, 活跃连接数: {len(_active_connections)}" + ) + return + + # 从请求参数或 Header 中提取 token + token = request.args.get("token", "") + if not token: + token = request.headers.get("X-WS-Token", "") + + if token != expected_token: + logger.warning(f"WebSocket 认证失败，拒绝连接: {request.sid}") + # 认证失败：拒绝连接并返回错误信息 + disconnect() + return False + + _active_connections.add(request.sid) + logger.info( + f"WebSocket 客户端已连接: {request.sid}, 活跃连接数: {len(_active_connections)}" + ) + + +@socketio.on("disconnect") +def on_disconnect(): + """WebSocket 断开连接：清理资源并记录""" + sid = request.sid + _active_connections.discard(sid) + logger.info( + f"WebSocket 客户端已断开: {sid}, 剩余活跃连接数: {len(_active_connections)}" + ) + def split_into_sentences(text: str, max_len: int = 80): if not text: @@ -20,15 +72,17 @@ def split_into_sentences(text: str, max_len: int = 80): # Step 1: split by punctuation (CJK and Western) # (?<=[。！？\n]) after CJK punctuation or newline # (?<=[.!?]\s) after period/question/exclamation + space - sentences = re.split(r'(?<=[。！？\n])|(?<=[.!?]\s)', text) - + sentences = re.split(r"(?<=[。！？\n])|(?<=[.!?]\s)", text) + # Filter empty strings and strip whitespace sentences = [s.strip() for s in sentences if s.strip()] - + # If no sentence boundaries found (e.g. long run-on text), split by paragraphs if len(sentences) <= 1: # Split by blank lines, then by single newlines - paragraphs = [p.strip() for p in text.replace('\n\n', '\n').split('\n') if p.strip()] + paragraphs = [ + p.strip() for p in text.replace("\n\n", "\n").split("\n") if p.strip() + ] if len(paragraphs) > 1: sentences = paragraphs @@ -36,39 +90,39 @@ def split_into_sentences(text: str, max_len: int = 80): def smart_split_long_sentence(sentence: str, max_len: int): if len(sentence) <= max_len: return [sentence] - + result = [] start = 0 text_len = len(sentence) - + while start < text_len: end = start + max_len - + # Already at end, append remainder if end >= text_len: result.append(sentence[start:].strip()) break - + # Look backward from end for first space or CJK punctuation as split point split_pos = -1 for i in range(end, max(start, end - 20), -1): - if sentence[i] in ' \t\n。！？,!?': + if sentence[i] in " \t\n。！？,!?": split_pos = i break - + # No break character found, force split at max_len if split_pos == -1: split_pos = end - - chunk = sentence[start:split_pos + 1].strip() + + chunk = sentence[start : split_pos + 1].strip() if chunk: result.append(chunk) start = split_pos + 1 - + # Guard against infinite loop if start >= text_len: break - + return result final = [] @@ -77,9 +131,10 @@ def smart_split_long_sentence(sentence: str, max_len: int): final.extend(smart_split_long_sentence(sent, max_len)) else: final.append(sent) - + return final if final else [text.strip()] + def smart_split_markdown(text: str, max_len: int = 80): """ Markdown-aware smart split. Preserves code blocks and key syntax; @@ -91,52 +146,54 @@ def smart_split_markdown(text: str, max_len: int = 80): # --- Layer 1: preserve code blocks --- # Split into [plain, code block, plain, code block, ...] # (```[\s\S]*?```) multi-line code, (`[^`\n]+`) inline code - parts = re.split(r'(```[\s\S]*?```|`[^`\n]+`)', text) + parts = re.split(r"(```[\s\S]*?```|`[^`\n]+`)", text) atoms = [] - + for part in parts: if not part: continue - + # 1. Code block (starts with `): treat as single atom - if part.startswith('`'): + if part.startswith("`"): # Very long block: split by newlines for streaming - if len(part) > max_len and '\n' in part: - code_lines = part.split('\n') + if len(part) > max_len and "\n" in part: + code_lines = part.split("\n") for idx, line in enumerate(code_lines): - suffix = '\n' if idx < len(code_lines) - 1 else '' + suffix = "\n" if idx < len(code_lines) - 1 else "" atoms.append(line + suffix) else: atoms.append(part) - + # 2. Plain text: fine-grained split else: # Split by paragraph/newline first (important in Markdown) - lines = part.split('\n') + lines = part.split("\n") for i, line in enumerate(lines): - suffix = '\n' if i < len(lines) - 1 else '' + suffix = "\n" if i < len(lines) - 1 else "" full_line = line + suffix - + if not line.strip(): atoms.append(full_line) continue # Split by sentence punctuation within line - sub_parts = re.split(r'([。！？]|(?<=[.!?])\s)', line) - + sub_parts = re.split(r"([。！？]|(?<=[.!?])\s)", line) + current_sent = "" for sub in sub_parts: current_sent += sub - if sub in ['。', '！', '？'] or (sub.strip() == '' and len(current_sent) > 0): + if sub in ["。", "！", "？"] or ( + sub.strip() == "" and len(current_sent) > 0 + ): atoms.append(current_sent) current_sent = "" - elif len(current_sent) > 0 and current_sent[-1] in '.!?': - pass - + elif len(current_sent) > 0 and current_sent[-1] in ".!?": + pass + if current_sent: atoms.append(current_sent) - + if suffix: if atoms: atoms[-1] += suffix @@ -162,7 +219,7 @@ def smart_split_markdown(text: str, max_len: int = 80): return final_chunks -@socketio.on('chat_request') +@socketio.on("chat_request") def handle_chat_request(data): """WebSocket chat handler: receive user message, push streaming results.""" # 1. Parse parameters @@ -170,26 +227,36 @@ def handle_chat_request(data): user_message = str(data.get("message") or "").strip() selected_model = str(data.get("model") or "") dag_confirm = str(data.get("dag_confirm") or "").strip() - is_dag_modification = str(data.get("is_dag_modification", "false")).lower() == "true" + is_dag_modification = ( + str(data.get("is_dag_modification", "false")).lower() == "true" + ) dag_id = str(data.get("dag_id") or "") modifications = data.get("modifications", "") # may be str or other expert_mode = data.get("expert_mode", False) dataset = str(data.get("dataset") or "").strip() # dataset name from frontend _dtype = data.get("dataset_type") or data.get("file_type") - dataset_type = str(_dtype).strip() if _dtype else None # "text" | "graph" | None + dataset_type = ( + str(_dtype).strip() if _dtype else None + ) # "text" | "graph" | None custom_mode = data.get("mode") except Exception as e: logger.error(f"Failed to parse parameters: {e}") - emit('chat_response', {"error": "Invalid request format. Please check parameters."}) + emit( + "chat_response", + {"error": "Invalid request format. Please check parameters."}, + ) return # 2. Validation if not user_message and not dag_confirm and not is_dag_modification: - emit('chat_response', {"error": "Message content cannot be empty."}) + emit("chat_response", {"error": "Message content cannot be empty."}) return if not dataset: - emit('chat_response', {"error": "Dataset is empty. Please specify a dataset first."}) + emit( + "chat_response", + {"error": "Dataset is empty. Please specify a dataset first."}, + ) return try: @@ -199,15 +266,19 @@ def handle_chat_request(data): # Callback to send streaming data (called from background event loop thread) def send_response(data_chunk): """Send response data to frontend.""" - socketio.emit('chat_response', data_chunk) + socketio.emit("chat_response", data_chunk) # Determine mode if custom_mode == "interact": mode = "interact" else: mode = "expert" if expert_mode else "normal" - logger.info(f"WS request: model={selected_model}, dataset={dataset}, message={user_message[:20]}..., expertMode={expert_mode}, mode={mode}") - print(f"WS request: model={selected_model}, dataset={dataset}, message={user_message[:20]}..., expertMode={expert_mode}, mode={mode}") + logger.info( + f"WS request: model={selected_model}, dataset={dataset}, message={user_message[:20]}..., expertMode={expert_mode}, mode={mode}" + ) + print( + f"WS request: model={selected_model}, dataset={dataset}, message={user_message[:20]}..., expertMode={expert_mode}, mode={mode}" + ) async def process_request(): try: @@ -219,41 +290,53 @@ async def process_request(): if result.get("success"): result_text = result.get("result", "") - paragraphs = [p.strip() for p in result_text.split('\n') if p.strip()] + paragraphs = [ + p.strip() for p in result_text.split("\n") if p.strip() + ] for i, para in enumerate(paragraphs): - send_response({ - 'type': 'result', - 'contentType': 'text', - 'content': para - }) - await asyncio.sleep(0.5 if i < len(paragraphs)-1 else 0.3) + send_response( + { + "type": "result", + "contentType": "text", + "content": para, + } + ) + await asyncio.sleep(0.5 if i < len(paragraphs) - 1 else 0.3) else: - send_response({ - 'error': result.get("error", "Analysis execution failed.") - }) - send_response({'type': 'stream_end'}) + send_response( + {"error": result.get("error", "Analysis execution failed.")} + ) + send_response({"type": "stream_end"}) return if is_dag_modification or (dag_confirm == "no" and modifications): engine = chat_service.engine_service.get_engine() engine.specific_dataset(dataset, dataset_type) modification_request = modifications or user_message - logger.info(f"DAG modification request received: {modification_request}") + logger.info( + f"DAG modification request received: {modification_request}" + ) - result = await chat_service.process_dag_modification(modification_request) + result = await chat_service.process_dag_modification( + modification_request + ) if result.get("success"): - dag_content = chat_service._convert_dag_to_frontend_format(result) - send_response({ - 'type': 'result', - 'contentType': 'dag', - 'content': dag_content - }) + dag_content = chat_service._convert_dag_to_frontend_format( + result + ) + send_response( + { + "type": "result", + "contentType": "dag", + "content": dag_content, + } + ) else: - send_response({ - 'error': result.get("error", "DAG modification failed.") - }) - send_response({'type': 'stream_end'}) + send_response( + {"error": result.get("error", "DAG modification failed.")} + ) + send_response({"type": "stream_end"}) return # Normal chat request — streaming (stream_end is sent internally) @@ -265,12 +348,18 @@ async def process_request(): dataset_type=dataset_type, mode=mode, expert_mode=expert_mode, - callback=send_response + callback=send_response, ) except Exception as exc: logger.error(f"Background processing failed: {exc}", exc_info=True) - send_response({'type': 'result', 'contentType': 'text', 'content': CHAT_FRIENDLY_ERROR_MSG}) - send_response({'type': 'stream_end'}) + send_response( + { + "type": "result", + "contentType": "text", + "content": CHAT_FRIENDLY_ERROR_MSG, + } + ) + send_response({"type": "stream_end"}) future = asyncio.run_coroutine_threadsafe(process_request(), loop) future.result() @@ -280,5 +369,12 @@ async def process_request(): except Exception as e: error_msg = f"Processing failed: {str(e)}" logger.error(error_msg, exc_info=True) - emit('chat_response', {'type': 'result', 'contentType': 'text', 'content': CHAT_FRIENDLY_ERROR_MSG}) - emit('chat_response', {'type': 'stream_end'}) + emit( + "chat_response", + { + "type": "result", + "contentType": "text", + "content": CHAT_FRIENDLY_ERROR_MSG, + }, + ) + emit("chat_response", {"type": "stream_end"}) diff --git a/web/frontend/run.py b/web/frontend/run.py index eb66711..9ec1669 100644 --- a/web/frontend/run.py +++ b/web/frontend/run.py @@ -1,8 +1,10 @@ +import os import socket from route import create_app, socketio app = create_app() + def find_free_port(start_port=5089, max_tries=50): """Find an available port starting from start_port.""" port = start_port @@ -15,13 +17,11 @@ def find_free_port(start_port=5089, max_tries=50): port += 1 raise RuntimeError("No available ports found.") + if __name__ == "__main__": port = find_free_port() print(f"Using available port: {port}") socketio.run( - app, - debug=True, - host="0.0.0.0", - port=port - ) \ No newline at end of file + app, debug=os.getenv("FLASK_DEBUG", "0") == "1", host="0.0.0.0", port=port + )