""" =========================================== 图谱检索服务模块 - GraphService =========================================== 功能：整合路由、节点检索和图谱检索的完整服务 提供统一的接口供app.py调用 =========================================== """ import os import logging import time from typing import Dict, Any, Optional, List, Tuple # 加载 .env 文件中的环境变量 try: from dotenv import load_dotenv load_dotenv() except ImportError: # 如果 python-dotenv 未安装，跳过加载（环境变量可能已通过其他方式设置） pass logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # 导入固定的schema try: from graph_search.schema_sample import SCHEMA_SAMPLE except ImportError: logger.warning("无法导入SCHEMA_SAMPLE，将使用空schema") SCHEMA_SAMPLE = "" def remove_knowledge_source(data: Any) -> Any: """ 递归删除数据中的 knowledge_source 字段（用于 results 字段） Args: data: 要处理的数据（可以是字典、列表或基本类型） Returns: 清理后的数据（已删除所有 knowledge_source 字段） """ if isinstance(data, dict): # 创建新字典，排除 knowledge_source result = {} for key, value in data.items(): if key != 'knowledge_source': # 递归处理嵌套的字典、列表 result[key] = remove_knowledge_source(value) return result elif isinstance(data, list): # 递归处理列表中的每个元素 return [remove_knowledge_source(item) for item in data] else: # 基本类型直接返回 return data # ========== ========== 导入依赖 ========== ========== try: from neo4j import GraphDatabase from langchain_community.graphs.neo4j_graph import Neo4jGraph from graph_search.graph_retrieval import GraphRetrieval from graph_search.route_label import RouteLabel from graph_search.node_retrieval import NodeRetrieval from graph_search.result_formatter import format_results, format_entry_nodes_as_results, remove_sensitive_fields, serialize_graph_for_llm from graph_search.init_prompts import get_all_ontology_labels, get_all_relationship_types from modelsAPI.model_api import OpenaiAPI GRAPH_SERVICE_AVAILABLE = True except ImportError as e: GRAPH_SERVICE_AVAILABLE = False logger.warning(f"图谱服务依赖未安装: {e}") class GraphService: """ 图谱检索服务类 功能：整合路由、节点检索和图谱检索的完整流程 """ def __init__(self,neo4j_url: str = None,neo4j_user: str = None,neo4j_password: str = None,driver = None ): """ 初始化图谱服务 Args: neo4j_url: Neo4j连接URL neo4j_user: Neo4j用户名 neo4j_password: Neo4j密码 driver: 可选的Neo4j driver实例，如果提供则复用外部driver（如app.py中的driver） """ if not GRAPH_SERVICE_AVAILABLE: raise ImportError("图谱服务依赖未安装") # ========== 从环境变量获取配置 ========== self.neo4j_url = neo4j_url or os.getenv("NEO4J_URI") or os.getenv("NEO4J_URL", "bolt://localhost:7687") self.neo4j_user = neo4j_user or os.getenv("NEO4J_USERNAME") or os.getenv("NEO4J_USER", "neo4j") self.neo4j_password = neo4j_password or os.getenv("NEO4J_PASSWORD", "neo4j") # ========== 初始化LLM API ========== self.llm_api = OpenaiAPI() # ========== 延迟初始化Neo4j连接 ========== self._graph_retrieval = None self._driver = driver self._get_driver(driver) # 如果提供了外部driver，直接使用；否则延迟初始化 self._neo4j_graph = None # Neo4jGraph 实例，用于获取 schema self._external_driver = driver is not None # 标记是否使用了外部driver # ========== 定义可选节点标签和关系类型 ========== self.node_types = get_all_ontology_labels(self._driver) self.relationship_types = get_all_relationship_types(self._driver) def _get_driver(self, driver): """获取Neo4j driver实例（延迟初始化，可复用）""" if driver is None: try: self._driver = GraphDatabase.driver( self.neo4j_url, auth=(self.neo4j_user, self.neo4j_password) ) self._external_driver = False # 标记为内部创建的driver except Exception as e: logger.error(f"Neo4j Driver实例初始化失败: {e}", exc_info=True) raise def _get_neo4j_graph(self): """获取Neo4jGraph实例（延迟初始化，用于获取schema）""" if self._neo4j_graph is None: try: self._neo4j_graph = Neo4jGraph( url=self.neo4j_url, username=self.neo4j_user, password=self.neo4j_password, enhanced_schema=True # 使用增强的schema信息，提供更详细的节点标签、关系类型和属性等 ) except Exception as e: logger.error(f"Neo4jGraph实例初始化失败: {e}", exc_info=True) raise return self._neo4j_graph def _fetch_schema_from_db_original(self): """ 原来的方式：通过 Neo4jGraph 自动获取 Neo4j schema 返回: neo4j_schema: str - 供 LLM 使用的 schema 文本（Neo4jGraph 返回的完整 schema） structured_schema: dict - 空字典（不进行解析，Schema验证会自动跳过） """ try: neo4j_graph = self._get_neo4j_graph() schema_text = neo4j_graph.schema if not schema_text: logger.warning("Neo4jGraph schema 返回为空") return "", {} return schema_text, {} except Exception as e: logger.warning(f"通过 Neo4jGraph 获取 schema 失败: {e}", exc_info=True) return "", {} def _fetch_schema_from_db_fixed(self): """ 固定方式：使用 schema_sample.py 中定义的固定 SCHEMA_SAMPLE 返回: neo4j_schema: str - 供 LLM 使用的 schema 文本 structured_schema: dict - 空字典（不进行解析，Schema验证会自动跳过） """ try: schema_text = SCHEMA_SAMPLE if not schema_text: logger.warning("SCHEMA_SAMPLE 为空") return "", {} return schema_text, {} except Exception as e: logger.warning(f"获取固定schema失败: {e}", exc_info=True) return "", {} def _fetch_schema_from_db(self): """ 获取schema（当前使用固定方式，可通过注释切换到原来的方式） 返回: neo4j_schema: str - 供 LLM 使用的 schema 文本 structured_schema: dict - 空字典（不进行解析，Schema验证会自动跳过） """ # ========== 当前使用的方式（可通过注释切换） ========== return self._fetch_schema_from_db_fixed() # ========== 原来的方式（切换时取消注释，注释上面一行） ========== # return self._fetch_schema_from_db_original() def _get_graph_retrieval(self) -> GraphRetrieval: """获取GraphRetrieval实例（延迟初始化）""" if self._graph_retrieval is None: # ========== 创建LLM包装器（适配GraphRetrieval的接口） ========== class LLMWrapper: """LLM包装器，将OpenaiAPI包装成GraphRetrieval需要的格式""" def __init__(self, llm_api): self.llm_api = llm_api async def ainvoke(self, prompt): """异步调用LLM""" if isinstance(prompt, str): query = prompt elif hasattr(prompt, 'messages'): # 如果是ChatPromptTemplate格式，转换为字符串 messages = prompt.messages query = "\n".join([msg.content for msg in messages if hasattr(msg, 'content')]) else: query = str(prompt) model = os.getenv("OPENAI_MODEL") content = await self.llm_api.open_api_chat_async(query, model, temperature=0) # 返回类似LLM输出的对象 class LLMOutput: def __init__(self, content): self.content = content return LLMOutput(content) llm_wrapper = LLMWrapper(self.llm_api) # ========== 获取 schema 信息 ========== # 使用 Neo4jGraph 获取 schema neo4j_schema, structured_schema = self._fetch_schema_from_db() if not neo4j_schema: logger.warning("GraphService: Neo4jGraph schema 为空，图谱检索可能无法正常工作") self._graph_retrieval = GraphRetrieval( driver=self._driver, neo4j_schema=neo4j_schema, llm=llm_wrapper, allowed_node_types=self.node_types, allowed_relationship_types=self.relationship_types ) return self._graph_retrieval async def extract_entry_nodes(self, query: str, top_n: int = 10) -> Tuple[Dict[str, List[Dict[str, Any]]], str, Dict[str, Any]]: """ 提取入口节点 功能： 1. 路由标签识别 2. 节点检索 Args: query: 用户查询 top_n: 每个标签检索的节点数量 Returns: tuple: (入口节点字典, 重写后的查询, 分类信息) 如果未重写，则返回原始query 分类信息包含 query_type, has_time, time_type """ extract_start = time.time() # ========== 1. 路由标签识别 ========== route_start = time.time() # 获取schema信息，用于帮助模型识别节点类型 # 使用 Neo4jGraph 获取 schema schema = None try: schema, _ = self._fetch_schema_from_db() if not schema: logger.warning("获取schema失败，将不使用schema信息") except Exception as e: logger.warning(f"获取schema失败，将不使用schema信息: {e}") schema = None route_label = RouteLabel(self.llm_api, self.node_types, schema=schema) route_res, rewritten_query, query_classification = await route_label.route(query) logger.info("输出route_res") logger.info(route_res) route_elapsed = (time.time() - route_start) * 1000 logger.info(f"[时间统计] 路由标签识别完成，耗时: {route_elapsed:.2f}ms") # 使用重写后的查询覆盖原始查询 if rewritten_query and rewritten_query != query: query = rewritten_query if not route_res: logger.warning("未能识别到入口节点") return {}, query, query_classification # ========== 2. 节点检索 ========== node_retrieval_start = time.time() try: # 创建嵌入模型包装器 class EmbeddingWrapper: """嵌入模型包装器""" def embed_documents(self, texts: List[str]) -> List[List[float]]: """批量获取嵌入向量（一次性处理所有文本）""" embed_start = time.time() if not texts: return [] try: # 一次性批量处理所有文本 result = OpenaiAPI.batch_embeddings(texts) # 调用服务器嵌入模型 embed_elapsed = (time.time() - embed_start) * 1000 # 从结果中提取embedding向量 if result and isinstance(result, list): embedding_list = [item.get('embedding') if isinstance(item, dict) else item for item in result] # 验证返回结果 if embedding_list and len(embedding_list) == len(texts): return embedding_list else: logger.error(f"EmbeddingWrapper: 批量向量化返回数量不匹配，期望: {len(texts)}, 实际: {len(embedding_list) if embedding_list else 0}") return [None] * len(texts) else: logger.error(f"EmbeddingWrapper: 批量向量化返回格式异常，期望列表，实际: {type(result)}") return [None] * len(texts) except Exception as e: logger.error(f"EmbeddingWrapper: 批量向量化失败: {e}", exc_info=True) return [None] * len(texts) embeddings = EmbeddingWrapper() node_retrieval = NodeRetrieval(self._driver, embeddings) entry_nodes = await node_retrieval.retrieve_nodes(route_res, top_n) node_retrieval_elapsed = (time.time() - node_retrieval_start) * 1000 logger.info(f"[时间统计] 节点检索完成，耗时: {node_retrieval_elapsed:.2f}ms") extract_elapsed = (time.time() - extract_start) * 1000 logger.info(f"[时间统计] 入口节点提取完成，总耗时: {extract_elapsed:.2f}ms (路由: {route_elapsed:.2f}ms, 检索: {node_retrieval_elapsed:.2f}ms)") logger.info(f"{entry_nodes}") logger.info(f"{query}") logger.info(f"{query_classification}") return entry_nodes, query, query_classification except Exception as e: logger.error(f"入口节点提取失败: {e}", exc_info=True) return {}, query, query_classification # 注意：不再关闭driver，因为它是可复用的 # ========== ========== 图谱检索：简单版本（只判断统计/非统计） ========== ========== async def search( self, query: str, entry_nodes: Optional[Dict[str, List[Dict[str, Any]]]] = None, top_k: int = 10, max_attempts: int = 3, query_type: str = "auto", # "auto" | "aggregate" | "detail" ) -> Dict[str, Any]: """ 执行图谱检索（完整流程） Args: query: 用户查询 entry_nodes: 入口节点（可选，不提供则自动提取） top_k: 节点检索数量 max_attempts: Cypher生成最大尝试次数 Returns: Dict: 检索结果，包含code, message, data, meta信息，符合app.py要求 """ total_start = time.time() logger.info(f"[时间统计] ========== 开始图谱检索，查询: '{query}' ==========") try: # ========== 1. 入口节点处理（如果已提供，直接使用；否则提取） ========== entry_nodes_start = time.time() query_classification = None # 初始化分类信息 if not entry_nodes: logger.info(f"[时间统计] 入口节点未提供，开始自动提取...") entry_nodes, rewritten_query, query_classification = await self.extract_entry_nodes(query = query, top_n = 10) entry_nodes_elapsed = (time.time() - entry_nodes_start) * 1000 # 使用重写后的查询覆盖原始查询 if rewritten_query and rewritten_query != query: query = rewritten_query if not entry_nodes: total_elapsed = (time.time() - total_start) * 1000 logger.warning(f"[时间统计] 未能识别到入口节点，总耗时: {total_elapsed:.2f}ms") return { "code": 200, "message": "success", "data": { "nodes": [], "links": [], "results": [] }, "meta": { "elapsed_ms": round(total_elapsed, 2), "error": "未能识别到入口节点", "query_type": "no_entry_nodes", "query_classification": query_classification } } logger.info(f"[时间统计] 入口节点提取完成，耗时: {entry_nodes_elapsed:.2f}ms") else: entry_nodes_elapsed = (time.time() - entry_nodes_start) * 1000 # ========== 2. 执行图谱检索 ========== graph_search_start = time.time() graph_retrieval = self._get_graph_retrieval() # 如果从重写中获得了分类信息，优先使用分类信息确定 query_type if query_classification: classification_query_type = query_classification.get("query_type", "detail") effective_query_type = classification_query_type else: # 如果没有分类信息，使用传入的 query_type 参数 effective_query_type = query_type if query_type in ("auto", "aggregate", "detail") else "auto" raw_results = await graph_retrieval.search( query=query, entry_nodes=entry_nodes, max_attempts=max_attempts, query_type=effective_query_type, query_classification=query_classification # 传递分类信息 ) graph_search_elapsed = (time.time() - graph_search_start) * 1000 logger.info(f"[时间统计] 图谱检索完成，耗时: {graph_search_elapsed:.2f}ms") # ========== 3. 格式化结果 ========== # 判断是否为统计类查询（优先使用分类信息） if query_classification: is_aggregate = query_classification.get("query_type") == "aggregate" else: is_aggregate = effective_query_type == "aggregate" or (effective_query_type == "auto" and raw_results and raw_results[0].get('_aggregate_value') is not None) # 从 rerank 后的结果中提取路径信息（节点和关系） path_data = {"nodes": [], "links": []} try: # raw_results 已经是 rerank 后的结果（字典列表，包含 rerank_score） if len(raw_results) > top_k: path_results_for_format = raw_results[:top_k] else: path_results_for_format = raw_results # 使用限制后的结果来格式化路径信息 path_data = format_results(path_results_for_format) except Exception as e: logger.warning(f"格式化结果失败: {e}", exc_info=True) # 过滤原始结果，递归删除 embedding、fulltext 和 path，并展开嵌套的 result 字段 # raw_results 已经是 rerank 后的结果，需要去除 path 字段，保留其他信息作为 results filtered_results = [] if is_aggregate: # 统计类查询：results 只包含统计值（一个值） if raw_results and raw_results[0].get('_aggregate_value') is not None: # 从元数据中提取统计值 aggregate_value = raw_results[0]['_aggregate_value'] # 删除 knowledge_source filtered_results = [remove_knowledge_source(aggregate_value)] else: # 如果没有 _aggregate_value，尝试从第一个结果的 result 字段提取 if raw_results: first_result = raw_results[0] if isinstance(first_result, dict) and 'result' in first_result: aggregate_value = first_result['result'] # 删除 knowledge_source filtered_results = [remove_knowledge_source(aggregate_value)] else: # 降级：使用第一个结果，删除 knowledge_source filtered_results = [remove_knowledge_source(first_result)] else: # 非统计类查询：保持原有逻辑 for result in raw_results: # 创建结果副本，去除 path 相关字段 result_copy = result.copy() # 删除 path 和 paths 字段（这些已经提取到 nodes 和 links 中了） result_copy.pop('path', None) result_copy.pop('paths', None) result_copy.pop('_aggregate_value', None) # 删除统计类查询的元数据 # 先递归删除敏感字段（embedding、fulltext） cleaned_result = remove_sensitive_fields(result_copy) # 如果结果中有 result 字段，则展开它 if isinstance(cleaned_result, dict) and 'result' in cleaned_result: result_value = cleaned_result['result'] # 如果 result 的值是列表，则展开列表中的所有元素 if isinstance(result_value, list): # 删除敏感字段和 knowledge_source filtered_results.extend([remove_knowledge_source(remove_sensitive_fields(item)) for item in result_value]) else: # 如果 result 的值是单个对象，直接使用，删除 knowledge_source filtered_results.append(remove_knowledge_source(remove_sensitive_fields(result_value))) else: # 没有 result 字段，直接使用清理后的结果，删除 knowledge_source filtered_results.append(remove_knowledge_source(cleaned_result)) # ========== 4. 格式化结果用于 LLM（在返回前应用格式化字符串到 results）========== # 构建临时的 graph_response 字典用于序列化 temp_graph_response = { "data": { "nodes": path_data.get("nodes", []), "results": filtered_results } } # 使用 serialize_graph_for_llm 格式化结果 formatted_results_str = serialize_graph_for_llm(temp_graph_response) # 将格式化后的字符串应用到 results（替换原有的 results） filtered_results = formatted_results_str # 构建响应 total_elapsed = (time.time() - total_start) * 1000 response = { "code": 200, "message": "success", "data": { "nodes": path_data.get("nodes", []), "links": path_data.get("links", []), "results": filtered_results }, "meta": { "result_count": len(filtered_results), "query_type": "graph_search", "elapsed_ms": round(total_elapsed, 2), "cypher": graph_retrieval.last_cypher if hasattr(graph_retrieval, 'last_cypher') else None, "query_classification": query_classification # 添加分类信息到返回结果中 } } logger.info(f"[时间统计] ========== 图谱检索完成 ==========") logger.info(f"[时间统计] 总耗时: {total_elapsed:.2f}ms") logger.info(f"[时间统计] - 入口节点提取: {entry_nodes_elapsed:.2f}ms") logger.info(f"[时间统计] - 图谱检索: {graph_search_elapsed:.2f}ms") return response except Exception as e: total_elapsed = (time.time() - total_start) * 1000 logger.error(f"[时间统计] 图谱检索异常，总耗时: {total_elapsed:.2f}ms, 错误: {e}", exc_info=True) return { "code": 500, "message": "图谱检索失败", "data": { "nodes": [], "links": [], "results": [] }, "meta": { "elapsed_ms": round(total_elapsed, 2), "error": str(e), "query_type": "error" } } def close(self): """ 关闭连接，释放资源 注意：在应用关闭时调用此方法 注意：如果使用的是外部driver（如app.py传入的），则不会关闭driver，由外部管理 """ try: if self._driver and not self._external_driver: # 只关闭内部创建的driver，不关闭外部传入的driver self._driver.close() self._driver = None logger.info("Neo4j Driver连接已关闭") elif self._external_driver: logger.info("使用的是外部driver，不关闭连接（由外部管理）") if self._neo4j_graph: # Neo4jGraph 内部使用 driver，如果 driver 已关闭，这里不需要额外操作 # 但为了清晰，我们重置引用 self._neo4j_graph = None logger.info("Neo4jGraph 实例已释放") except Exception as e: logger.warning(f"关闭连接时出错: {e}") def __del__(self): """析构函数，确保资源被释放""" self.close()