884 lines
40 KiB
Python
884 lines
40 KiB
Python
"""
|
||
===========================================
|
||
图谱检索模块 - GraphRetrieval
|
||
===========================================
|
||
"""
|
||
import logging
|
||
import time
|
||
import os
|
||
from typing import Optional, List, Dict, Tuple, Any, Literal
|
||
# 导入拆分后的模块
|
||
from graph_search.query_intent import is_aggregate_query
|
||
from graph_search.cypher_validator import CypherValidator
|
||
from graph_search.cypher_fixer import parse_cypher_error,force_undirected_relationships,fix_cypher_based_on_error,fix_common_cypher_errors
|
||
from graph_search.entry_node_handler import filter_and_sort_entry_nodes,format_entry_nodes
|
||
from graph_search.result_formatter import format_entry_nodes_as_results
|
||
from graph_search.aggregate_handler import extract_aggregate_functions, is_aggregate_result_zero
|
||
from graph_search.init_prompts import *
|
||
# 加载 .env 文件中的环境变量
|
||
try:
|
||
from dotenv import load_dotenv
|
||
load_dotenv()
|
||
except ImportError:
|
||
# 如果 python-dotenv 未安装，跳过加载（环境变量可能已通过其他方式设置）
|
||
pass
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ========== ========== 导入依赖 ========== ==========
|
||
# 注意：需要安装以下包
|
||
# pip install neo4j langchain langchain-community neo4j-graphrag
|
||
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
|
||
from neo4j_graphrag.retrievers.text2cypher import extract_cypher
|
||
|
||
|
||
class GraphRetrieval:
|
||
"""
|
||
知识图谱检索模块（带完整验证）
|
||
|
||
功能说明：
|
||
1. 生成Cypher查询语句（支持错误反馈自动修正）
|
||
2. 多层验证：语法、Schema、Dry Run、LLM语义验证
|
||
3. 多轮生成和修正机制（错误反馈给生成函数）
|
||
4. 降级策略：模板匹配 -> LLM生成 -> 简单查询
|
||
"""
|
||
def __init__(self,
|
||
driver, # Neo4j driver（GraphDatabase.driver实例）
|
||
neo4j_schema: str, # Neo4j schema字符串（从graph_service.py传递，由Neo4jGraph获取）
|
||
llm = None, # LLM实例（用于生成和验证Cypher，如ChatTongyi）
|
||
allowed_node_types: List[str] = None, # 允许的节点类型列表（从graph_service.py传递）
|
||
allowed_relationship_types: List[str] = None, # 允许的关系类型列表（从graph_service.py传递）
|
||
):
|
||
"""
|
||
初始化图谱检索器
|
||
"""
|
||
self.driver = driver
|
||
self.neo4j_schema = neo4j_schema # 完整schema（字符串格式，从graph_service.py传递）
|
||
self.llm = llm
|
||
# Dry Run验证
|
||
self.skip_dry_run_validation = False
|
||
# ========== 配置选项：Cypher生成最大尝试次数 ==========
|
||
# 配置方式：在 .env 文件中添加 CYPHER_MAX_ATTEMPTS=2 或 CYPHER_MAX_ATTEMPTS=3
|
||
self.default_max_attempts = int(os.getenv("CYPHER_MAX_ATTEMPTS", "2"))
|
||
logger.info(f"GraphRetrieval: Cypher生成默认最大尝试次数: {self.default_max_attempts}（通过 .env 文件或环境变量 CYPHER_MAX_ATTEMPTS 配置）")
|
||
|
||
# ========== 保存最后使用的Cypher（用于调试） ==========
|
||
self.last_cypher = None
|
||
|
||
# ========== 初始化验证器 ==========
|
||
self.validator = CypherValidator(
|
||
driver=self.driver,
|
||
neo4j_schema=self.neo4j_schema,
|
||
allowed_node_types=allowed_node_types,
|
||
allowed_relationship_types=allowed_relationship_types
|
||
)
|
||
|
||
def _force_undirected_relationships(self, cypher: str) -> str:
|
||
"""包装函数，调用模块中的函数"""
|
||
return force_undirected_relationships(cypher)
|
||
|
||
def _parse_cypher_error(self, error_message: str, cypher: str) -> Dict[str, Any]:
|
||
"""包装函数，调用模块中的函数"""
|
||
return parse_cypher_error(error_message, cypher)
|
||
|
||
def _fix_cypher_based_on_error(self, cypher: str, error_info: Dict[str, Any]) -> Optional[str]:
|
||
"""包装函数，调用模块中的函数"""
|
||
return fix_cypher_based_on_error(cypher, error_info)
|
||
|
||
def _fix_common_cypher_errors(self, cypher: str) -> str:
|
||
"""包装函数，调用模块中的函数"""
|
||
return fix_common_cypher_errors(cypher)
|
||
|
||
# ========== ========== 规则验证层 ========== ==========
|
||
|
||
def validate_cypher_syntax(self, cypher: str) -> Tuple[bool, Optional[str]]:
|
||
"""包装函数，调用验证器"""
|
||
return self.validator.validate_syntax(cypher)
|
||
|
||
def validate_cypher_schema(self, cypher: str) -> Tuple[bool, Optional[str]]:
|
||
"""包装函数，调用验证器的schema验证"""
|
||
return self.validator.validate_schema(cypher)
|
||
|
||
def dry_run_cypher(self, cypher: str) -> Tuple[bool, Optional[str], Optional[Dict[str, Any]]]:
|
||
"""包装函数，调用验证器"""
|
||
return self.validator.dry_run(cypher, parse_error_func=self._parse_cypher_error)
|
||
|
||
# ========== ========== Cypher生成（支持错误反馈） ========== ==========
|
||
|
||
async def generate_cypher(
|
||
self,
|
||
query: str,
|
||
entry_nodes: dict,
|
||
previous_cypher: Optional[str] = None,
|
||
validation_errors: Optional[List[str]] = None,
|
||
is_aggregate: bool = False,
|
||
error_type: Optional[str] = None, # 'syntax_error' | 'schema_error' | 'empty_result'
|
||
query_classification: Optional[Dict[str, Any]] = None, # 从重写中获得的分类信息（可选）
|
||
) -> Optional[str]:
|
||
"""
|
||
生成Cypher语句（支持错误反馈，自动修正）
|
||
"""
|
||
from graph_search.query_intent import (
|
||
has_time_in_query,
|
||
get_current_date_info,
|
||
classify_time_aggregate_type
|
||
)
|
||
|
||
is_correction = validation_errors and previous_cypher
|
||
|
||
# ========== 判断是否包含时间信息（优先使用分类信息） ==========
|
||
if query_classification:
|
||
has_time = query_classification.get("has_time", False)
|
||
# 如果分类信息中有 time_type，直接使用；否则根据 has_time 和 is_aggregate 判断
|
||
if is_aggregate and has_time:
|
||
time_aggregate_type = query_classification.get("time_type")
|
||
# 如果 time_type 为 None，尝试自动识别
|
||
if time_aggregate_type is None:
|
||
time_aggregate_type = classify_time_aggregate_type(query)
|
||
else:
|
||
time_aggregate_type = None
|
||
logger.info(f"使用分类信息: has_time={has_time}, time_aggregate_type={time_aggregate_type}")
|
||
else:
|
||
# 如果没有分类信息，使用原有逻辑
|
||
has_time = has_time_in_query(query)
|
||
time_aggregate_type = None
|
||
if is_aggregate and has_time:
|
||
time_aggregate_type = classify_time_aggregate_type(query)
|
||
logger.info(f"使用自动识别: has_time={has_time}, time_aggregate_type={time_aggregate_type}")
|
||
|
||
# ========== 获取系统时间信息（如果需要） ==========
|
||
current_date_info = ""
|
||
current_date_info_section = ""
|
||
date_info = {}
|
||
current_month_int = None # 整数类型的月份（用于 prompt 中的 {current_month:02d} 格式化）
|
||
current_month_formatted = "" # 格式化的月份字符串（如 "01", "12"）
|
||
if has_time:
|
||
date_info = get_current_date_info()
|
||
current_date_info = "\n".join([f"- {k}: {v}" for k, v in date_info.items()])
|
||
current_date_info_section = f"系统时间信息:\n{current_date_info}\n\n"
|
||
# 获取整数类型的月份（用于 prompt 模板中的 {current_month:02d} 格式化）
|
||
current_month_int = int(date_info.get('current_month', 1))
|
||
# 同时提供格式化好的字符串版本（备用）
|
||
current_month_formatted = f"{current_month_int:02d}"
|
||
|
||
# ========== 过滤和排序入口节点（按标签内排序和限制） ==========
|
||
# 注意：如果 entry_nodes 已经在 search 方法中过滤过，这里再次过滤是幂等的（结果相同）
|
||
entry_nodes = filter_and_sort_entry_nodes(
|
||
entry_nodes,
|
||
max_nodes_per_label=5
|
||
)
|
||
|
||
# ========== 格式化入口节点信息 ==========
|
||
|
||
entry_nodes_str = format_entry_nodes(entry_nodes)
|
||
|
||
# ========== 路径模板占位说明 ==========
|
||
# 注意：路径模板生成功能已移除，LLM 需要根据 schema 自行选择合理的路径
|
||
path_templates_str = "（请根据 schema 和入口节点信息，自行选择合理的查询路径。）"
|
||
|
||
# ========== 根据统计 / 普通查询，构造查询意图提示 ==========
|
||
if is_aggregate:
|
||
# 统计类：强制要求使用聚合函数，返回简洁聚合结果
|
||
query_intent_hint = (
|
||
"【查询类型】统计类问题，请根据用户问题和入口节点，生成包含聚合函数的 Cypher 查询，"
|
||
"例如使用 COUNT、SUM、AVG、MAX、MIN 等，返回精简的聚合结果（如总数、平均值、排名等），"
|
||
"不要返回大量明细记录。\n"
|
||
)
|
||
else:
|
||
# 明细类：正常检索，用于后续 RAG
|
||
query_intent_hint = (
|
||
"【查询类型】检索类问题，请根据用户问题和入口节点，生成返回明细节点/关系的 Cypher 查询，"
|
||
"用于后续的图谱问答检索。\n"
|
||
)
|
||
|
||
full_query = query_intent_hint + query
|
||
|
||
# ========== 根据是否有错误，选择不同的prompt ==========
|
||
if validation_errors and previous_cypher:
|
||
# ========== 修正模式：包含错误信息和之前的Cypher ==========
|
||
errors_str = "\n".join([f" - {err}" for err in validation_errors])
|
||
|
||
# ========== 根据错误类型选择对应的修正 prompt（统计类和明细类都使用相同的分类） ==========
|
||
if error_type == 'syntax_error':
|
||
# 语法错误：使用专门的语法修正 prompt
|
||
prompt = correct_syntax_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
errors=errors_str,
|
||
cypher=previous_cypher
|
||
)
|
||
elif error_type == 'schema_error':
|
||
# Schema错误：使用专门的 Schema 修正 prompt
|
||
prompt = correct_schema_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
errors=errors_str,
|
||
cypher=previous_cypher
|
||
)
|
||
elif error_type == 'empty_result':
|
||
# 空结果错误：使用专门的空结果修正 prompt
|
||
prompt = correct_empty_result_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
errors=errors_str,
|
||
cypher=previous_cypher
|
||
)
|
||
else:
|
||
# 未分类错误：使用通用修正 prompt（向后兼容）
|
||
# 统计类使用 correct_aggregate_cypher_prompt，明细类使用 correct_cypher_prompt
|
||
if is_aggregate:
|
||
prompt = correct_aggregate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
errors=errors_str,
|
||
cypher=previous_cypher
|
||
)
|
||
else:
|
||
prompt = correct_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
errors=errors_str,
|
||
cypher=previous_cypher
|
||
)
|
||
else:
|
||
# ========== 初始生成模式 ==========
|
||
if is_aggregate:
|
||
# 统计类：根据是否带时间选择不同的prompt
|
||
if has_time:
|
||
# 时间统计查询：根据时间段类型选择对应的prompt
|
||
if time_aggregate_type == "single":
|
||
# 单时间段统计
|
||
prompt = generate_single_time_aggregate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
current_date_info=current_date_info,
|
||
current_year=date_info.get("current_year", ""),
|
||
current_month=current_month_int, # 传递整数，支持 :02d 格式化
|
||
last_year=date_info.get("last_year", ""),
|
||
next_year=date_info.get("next_year", ""),
|
||
next_month_start=date_info.get("next_month_start", ""),
|
||
)
|
||
elif time_aggregate_type == "dual":
|
||
# 双时间段统计（同比/环比）
|
||
# 计算上月（last_month）和前年（last_last_year，用于多时间段）
|
||
last_month = ""
|
||
if current_month_formatted:
|
||
month_int = int(current_month_formatted)
|
||
year_int = int(date_info.get("current_year", "2025"))
|
||
if month_int == 1:
|
||
last_month = f"{year_int - 1}-12"
|
||
else:
|
||
last_month = f"{year_int}-{month_int - 1:02d}"
|
||
|
||
prompt = generate_dual_time_aggregate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
current_date_info=current_date_info,
|
||
current_year=date_info.get("current_year", ""),
|
||
current_month=current_month_int, # 传递整数，支持 :02d 格式化
|
||
last_year=date_info.get("last_year", ""),
|
||
last_month=last_month,
|
||
)
|
||
else: # multi
|
||
# 多时间段统计
|
||
# 计算前年（last_last_year）
|
||
last_last_year = str(int(date_info.get("current_year", "2025")) - 2) if date_info else ""
|
||
prompt = generate_multi_time_aggregate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema,
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
current_date_info=current_date_info,
|
||
current_year=date_info.get("current_year", ""),
|
||
last_year=date_info.get("last_year", ""),
|
||
last_last_year=last_last_year,
|
||
)
|
||
else:
|
||
# 普通统计查询：使用原有prompt
|
||
prompt = generate_aggregate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema, # 使用完整schema
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
path_templates=path_templates_str,
|
||
)
|
||
else:
|
||
# 普通查询：支持时间过滤（在prompt中已包含时间处理说明）
|
||
prompt = generate_cypher_prompt.format_prompt(
|
||
schema=self.neo4j_schema, # 使用完整schema
|
||
entry_nodes=entry_nodes_str,
|
||
query=full_query,
|
||
path_templates=path_templates_str,
|
||
current_date_info=current_date_info if has_time else "",
|
||
current_year=date_info.get("current_year", "") if has_time else "",
|
||
current_month=current_month_int if has_time else 1, # 传递整数，支持 :02d 格式化
|
||
next_year=date_info.get("next_year", "") if has_time else "",
|
||
next_month_start=date_info.get("next_month_start", "") if has_time else "",
|
||
current_date_info_section=current_date_info_section,
|
||
)
|
||
|
||
|
||
try:
|
||
# ========== 调用LLM生成Cypher ==========
|
||
# prompt可能是字符串或ChatPromptTemplate对象
|
||
if isinstance(prompt, str):
|
||
query_text = prompt
|
||
elif hasattr(prompt, 'messages'):
|
||
# 如果是ChatPromptTemplate格式，转换为字符串
|
||
messages = prompt.messages
|
||
query_text = "\n".join([
|
||
msg.content for msg in messages
|
||
if hasattr(msg, 'content') and msg.content
|
||
])
|
||
else:
|
||
query_text = str(prompt)
|
||
|
||
llm_output = await self.llm.ainvoke(query_text)
|
||
|
||
# ========== 提取Cypher语句 ==========
|
||
content = llm_output.content if hasattr(llm_output, 'content') else str(llm_output)
|
||
|
||
# ========== 【新增】过滤 Qwen 的 <think> 标签 ==========
|
||
import re
|
||
content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL).strip()
|
||
|
||
# 优先尝试从标记中提取（如果 LLM 按照格式输出）
|
||
cypher = None
|
||
if "【Cypher查询语句】" in content or "Cypher查询语句" in content:
|
||
# 查找标记后的内容
|
||
markers = ["【Cypher查询语句】", "Cypher查询语句", "Cypher:", "Cypher语句："]
|
||
for marker in markers:
|
||
marker_pos = content.find(marker)
|
||
if marker_pos != -1:
|
||
# 提取标记后的内容
|
||
cypher_text = content[marker_pos + len(marker):].strip()
|
||
# 使用 extract_cypher 提取
|
||
cypher = extract_cypher(cypher_text)
|
||
if cypher:
|
||
break
|
||
|
||
# 如果标记提取失败，使用原有的 extract_cypher（会从整个内容中提取）
|
||
if not cypher:
|
||
cypher = extract_cypher(content)
|
||
|
||
if cypher:
|
||
# ========== 后处理：修复常见语法错误 ==========
|
||
cypher = self._fix_common_cypher_errors(cypher)
|
||
|
||
# ========== 强制转换为无方向关系（确保所有关系都是无方向的） ==========
|
||
# 即使prompt中已要求无方向关系，也要强制转换作为安全措施
|
||
cypher = self._force_undirected_relationships(cypher)
|
||
|
||
# 打印更全面的 Cypher 日志，便于调试（长度做一个上限，避免日志爆炸）
|
||
preview = cypher if len(cypher) <= 1000 else cypher[:1000] + "..."
|
||
logger.info(f"Cypher生成完整语句:\n{preview}")
|
||
else:
|
||
logger.warning(f"Cypher生成结果为空，LLM原始输出: {str(content)[:500]}...")
|
||
return cypher
|
||
except Exception as e:
|
||
logger.error(f"LLM生成Cypher失败: {e}")
|
||
return None
|
||
|
||
# ========== ========== 多轮生成和验证 ========== ==========
|
||
|
||
async def generate_and_validate_cypher(
|
||
self,
|
||
query: str,
|
||
entry_nodes: dict,
|
||
max_attempts: int = 3,
|
||
is_aggregate: bool = False,
|
||
query_classification: Optional[Dict[str, Any]] = None, # 从重写中获得的分类信息（可选）
|
||
) -> Tuple[Optional[str], List[str]]:
|
||
"""
|
||
生成和验证Cypher（完整流程：多轮生成+验证+自动修正）
|
||
|
||
工作流程：
|
||
1. 生成Cypher语句
|
||
2. 验证Cypher语句（语法+LLM语义验证）
|
||
3. 如果有错误，将错误反馈给生成函数，重新生成
|
||
4. 关系方向校正
|
||
5. 返回最终Cypher
|
||
"""
|
||
errors_log = []
|
||
previous_cypher = None
|
||
error_type = None # 当前错误类型：'syntax_error' | 'schema_error' | 'empty_result'
|
||
|
||
for attempt in range(1, max_attempts + 1):
|
||
# ========== 1. 生成/修正Cypher ==========
|
||
cypher = await self.generate_cypher(
|
||
query=query,
|
||
entry_nodes=entry_nodes,
|
||
previous_cypher=previous_cypher,
|
||
validation_errors=errors_log if attempt > 1 else None,
|
||
is_aggregate=is_aggregate,
|
||
error_type=error_type, # 传递错误类型，用于选择对应的修正 prompt
|
||
query_classification=query_classification, # 传递分类信息
|
||
)
|
||
|
||
if not cypher:
|
||
errors_log.append(f"第{attempt}轮: Cypher生成失败")
|
||
continue
|
||
|
||
# ========== 2. 快速规则验证（必须通过） ==========
|
||
syntax_valid, syntax_error = self.validate_cypher_syntax(cypher)
|
||
|
||
if not syntax_valid:
|
||
errors_log.append(f"语法错误: {syntax_error}")
|
||
error_type = 'syntax_error' # 标记为语法错误
|
||
previous_cypher = cypher
|
||
continue
|
||
|
||
# ========== 3. Schema验证（在Dry Run之前执行） ==========
|
||
schema_valid, schema_error = self.validate_cypher_schema(cypher)
|
||
|
||
if not schema_valid:
|
||
errors_log.append(f"Schema错误: {schema_error}")
|
||
error_type = 'schema_error'
|
||
previous_cypher = cypher
|
||
continue
|
||
|
||
# ========== 4. Dry Run验证（可选，可通过环境变量跳过） ==========
|
||
if self.skip_dry_run_validation:
|
||
dry_run_valid, dry_run_error, error_info = True, None, None # 跳过验证
|
||
else:
|
||
dry_run_valid, dry_run_error, error_info = self.dry_run_cypher(cypher)
|
||
|
||
if not dry_run_valid:
|
||
# ========== 从 error_info 中提取错误类型 ==========
|
||
if error_info and error_info.get('error_type'):
|
||
error_type = error_info['error_type'] # 'syntax_error' 或 'schema_error'
|
||
else:
|
||
# 如果 error_info 中没有 error_type，根据错误信息推断
|
||
if 'Schema错误' in dry_run_error or 'schema' in dry_run_error.lower():
|
||
error_type = 'schema_error'
|
||
else:
|
||
error_type = 'syntax_error'
|
||
|
||
# ========== 对于其他错误，尝试基于实际错误信息进行智能修复 ==========
|
||
if error_info:
|
||
fixed_cypher = self._fix_cypher_based_on_error(cypher, error_info)
|
||
if fixed_cypher and fixed_cypher != cypher:
|
||
logger.info(f"[智能修复] 基于Dry Run错误信息自动修复Cypher")
|
||
logger.debug(f"[智能修复] 原始: {cypher[:200]}...")
|
||
logger.debug(f"[智能修复] 修复后: {fixed_cypher[:200]}...")
|
||
# 重新验证修复后的Cypher
|
||
cypher = fixed_cypher
|
||
error_type = None # 重置错误类型，重新验证
|
||
# 重新进行验证（从语法验证开始）
|
||
continue
|
||
|
||
# 如果无法自动修复，记录错误并继续
|
||
error_detail = f"执行计划错误: {dry_run_error}"
|
||
if error_info and error_info.get("error_position"):
|
||
pos = error_info["error_position"]
|
||
error_detail += f" (位置: 第{pos.get('line', '?')}行, 第{pos.get('column', '?')}列)"
|
||
errors_log.append(error_detail)
|
||
previous_cypher = cypher
|
||
continue
|
||
|
||
# ========== 所有验证通过 ==========
|
||
# 注意：关系方向校正已移除，因为prompt中已明确要求使用无方向关系
|
||
self.last_cypher = cypher # 保存最后使用的Cypher
|
||
return cypher, errors_log
|
||
|
||
# ========== 所有尝试都失败 ==========
|
||
logger.error(f"GraphRetrieval: Cypher生成失败，所有{max_attempts}轮尝试均失败")
|
||
return None, errors_log
|
||
|
||
async def search(
|
||
self,
|
||
query: str,
|
||
entry_nodes: dict,
|
||
max_attempts: int = None,
|
||
query_type: Literal["auto", "aggregate", "detail"] = "auto",
|
||
query_classification: Optional[Dict[str, Any]] = None,
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
执行图谱检索（完整流程）
|
||
"""
|
||
search_start = time.time()
|
||
logger.info("GraphRetrieval.search: 开始图谱检索")
|
||
|
||
# ========== 0. 过滤入口节点 ==========
|
||
entry_nodes = filter_and_sort_entry_nodes(entry_nodes, max_nodes_per_label=5)
|
||
logger.info("输出过滤后的入口节点")
|
||
# logger.info(entry_nodes[0].get("名称"))
|
||
# logger.info(entry_nodes)
|
||
|
||
|
||
# ========== 1. 判断是否为聚合查询 ==========
|
||
if query_classification:
|
||
classification_query_type = query_classification.get("query_type", "detail")
|
||
is_aggregate = (classification_query_type == "aggregate")
|
||
else:
|
||
if query_type == "aggregate":
|
||
is_aggregate = True
|
||
elif query_type == "detail":
|
||
is_aggregate = False
|
||
else:
|
||
is_aggregate = is_aggregate_query(query)
|
||
|
||
# ========== 2. 生成 Cypher ==========
|
||
if max_attempts is None:
|
||
max_attempts = self.default_max_attempts
|
||
|
||
cypher, errors = await self.generate_and_validate_cypher(
|
||
query=query,
|
||
entry_nodes=entry_nodes,
|
||
max_attempts=max_attempts,
|
||
is_aggregate=is_aggregate,
|
||
query_classification=query_classification,
|
||
)
|
||
|
||
if not cypher:
|
||
return await self._fallback_to_simple_or_entry_nodes(query, entry_nodes, is_aggregate)
|
||
|
||
# ========== 3. 执行查询（支持拆分）==========
|
||
records, path_records = self._execute_cypher_with_split(cypher, is_aggregate)
|
||
|
||
# ========== 4. 处理空结果或 0 值（仅聚合）==========
|
||
if not records:
|
||
logger.warning("Cypher查询结果为空，尝试重试")
|
||
new_cypher, _ = await self.generate_and_validate_cypher(
|
||
query=query,
|
||
entry_nodes=entry_nodes,
|
||
max_attempts=max_attempts,
|
||
is_aggregate=is_aggregate,
|
||
query_classification=query_classification,
|
||
)
|
||
if new_cypher and new_cypher != cypher:
|
||
new_records, new_path_records = self._execute_cypher_with_split(new_cypher, is_aggregate)
|
||
if new_records:
|
||
self.last_cypher = new_cypher
|
||
return self._process_results(new_records, new_path_records, query, is_aggregate)
|
||
|
||
elif is_aggregate and is_aggregate_result_zero([dict(r) for r in records], cypher):
|
||
logger.warning("统计类查询返回0，尝试重试")
|
||
new_cypher, _ = await self.generate_and_validate_cypher(
|
||
query=query,
|
||
entry_nodes=entry_nodes,
|
||
max_attempts=max_attempts,
|
||
is_aggregate=is_aggregate,
|
||
query_classification=query_classification,
|
||
)
|
||
if new_cypher and new_cypher != cypher:
|
||
new_records, new_path_records = self._execute_cypher_with_split(new_cypher, is_aggregate)
|
||
if new_records:
|
||
self.last_cypher = new_cypher
|
||
return self._process_results(new_records, new_path_records, query, is_aggregate)
|
||
|
||
# ========== 5. 正常返回 ==========
|
||
self.last_cypher = cypher
|
||
result = self._process_results(records, path_records, query, is_aggregate)
|
||
|
||
search_elapsed = (time.time() - search_start) * 1000
|
||
logger.info(f"[时间统计] GraphRetrieval.search: 完成，总耗时: {search_elapsed:.2f}ms, 结果数: {len(result)}")
|
||
return result
|
||
|
||
def _execute_cypher_with_split(self, cypher: str, is_aggregate: bool):
|
||
"""执行 Cypher，支持拆分为 path + aggregate"""
|
||
path_cypher, aggregate_cypher = self._split_aggregate_cypher(cypher)
|
||
|
||
if path_cypher and aggregate_cypher:
|
||
logger.info("[Cypher拆分] 执行 path 和 aggregate 查询")
|
||
try:
|
||
path_result = self.driver.execute_query(path_cypher)
|
||
path_records = path_result.records
|
||
except Exception as e:
|
||
logger.warning(f"[Cypher拆分] Path 查询失败: {e}")
|
||
path_records = None
|
||
|
||
try:
|
||
aggregate_result = self.driver.execute_query(aggregate_cypher)
|
||
records = aggregate_result.records
|
||
except Exception as e:
|
||
logger.error(f"[Cypher拆分] Aggregate 查询失败: {e}，回退到原始 Cypher")
|
||
result = self.driver.execute_query(cypher)
|
||
return result.records, None
|
||
return records, path_records
|
||
else:
|
||
result = self.driver.execute_query(cypher)
|
||
return result.records, None
|
||
|
||
def _process_results(
|
||
self,
|
||
records,
|
||
path_records,
|
||
query: str,
|
||
is_aggregate: bool
|
||
) -> List[Dict[str, Any]]:
|
||
"""
|
||
统一处理查询结果：
|
||
- 聚合查询：注入 result 到 path，rerank path
|
||
- 非聚合：确保每条记录有可序列化的 'result' 字段，并 rerank
|
||
"""
|
||
|
||
if is_aggregate:
|
||
aggregate_results = [dict(r) for r in records]
|
||
aggregate_value = aggregate_results[0].get('result', aggregate_results[0])
|
||
|
||
if path_records:
|
||
raw_results = [dict(r) for r in path_records]
|
||
# 注入聚合值
|
||
for res in raw_results:
|
||
res['result'] = aggregate_value
|
||
if raw_results:
|
||
raw_results[0]['_aggregate_value'] = aggregate_value
|
||
|
||
raw_results = self._rerank_results(raw_results, query, is_aggregate=False)
|
||
return raw_results
|
||
else:
|
||
return aggregate_results
|
||
|
||
|
||
else:
|
||
value_results = [dict(r) for r in records]
|
||
if path_records:
|
||
raw_results = [dict(r) for r in path_records]
|
||
# 注入明细查询结果，防止 value_results 索引越界
|
||
for i in range(len(raw_results)):
|
||
if i < len(value_results):
|
||
raw_results[i]['result'] = value_results[i]
|
||
else:
|
||
raw_results[i]['result'] = None # 或 {}、""、[] 等默认值
|
||
|
||
raw_results = self._rerank_results(raw_results, query, is_aggregate=False)
|
||
return raw_results
|
||
else:
|
||
return value_results
|
||
|
||
async def _fallback_to_simple_or_entry_nodes(
|
||
self,
|
||
query: str,
|
||
entry_nodes: dict,
|
||
is_aggregate: bool,
|
||
) -> List[Dict[str, Any]]:
|
||
"""降级：先试简单模板，再试入口节点"""
|
||
# 尝试简单模板
|
||
simple_cypher = self._generate_simple_template_cypher(entry_nodes, query, is_aggregate)
|
||
if simple_cypher:
|
||
try:
|
||
result = self.driver.execute_query(simple_cypher)
|
||
if result.records:
|
||
raw_results = [dict(r) for r in result.records]
|
||
raw_results = self._rerank_results(raw_results, query, is_aggregate)
|
||
logger.info(f"[降级策略] 简单模板查询成功，返回 {len(raw_results)} 条结果")
|
||
return raw_results
|
||
except Exception as e:
|
||
logger.warning(f"[降级策略] 简单模板查询失败: {e}")
|
||
|
||
# 兜底：入口节点
|
||
fallback_results = format_entry_nodes_as_results(entry_nodes)
|
||
if fallback_results:
|
||
logger.info(f"[降级策略] 返回 {len(fallback_results)} 条入口节点信息（兜底）")
|
||
return fallback_results
|
||
return []
|
||
|
||
|
||
def _split_aggregate_cypher(self, cypher: str) -> Tuple[Optional[str], Optional[str]]:
|
||
if not cypher or not cypher.strip():
|
||
return None, None
|
||
|
||
try:
|
||
import re
|
||
|
||
# 提取 MATCH 模式部分（支持前导任意内容，如 "cypher\n"）
|
||
match_match = re.search(
|
||
r'(?i)\bMATCH\s+(.+?)(?=\s+\b(?:WHERE|WITH|RETURN)\b)',
|
||
cypher,
|
||
re.DOTALL
|
||
)
|
||
if not match_match:
|
||
logger.warning(f"[Cypher处理] 未找到合法的 MATCH 子句。输入片段: {cypher[:100]}...")
|
||
return None, None
|
||
|
||
full_match_clause = match_match.group(1).strip()
|
||
|
||
# 提取 WHERE 子句（如果有）
|
||
where_match = re.search(
|
||
r'(?i)\bWHERE\s+(.+?)(?=\s+\b(?:WITH|RETURN)\b)',
|
||
cypher,
|
||
re.DOTALL
|
||
)
|
||
where_clause = " WHERE " + where_match.group(1).strip() if where_match else ""
|
||
|
||
# 取第一个连通模式（逗号分隔时只取第一个）
|
||
patterns = [p.strip() for p in full_match_clause.split(',') if p.strip()]
|
||
if not patterns:
|
||
logger.warning("[Cypher处理] MATCH 子句中无有效模式")
|
||
return None, None
|
||
|
||
first_pattern = patterns[0]
|
||
|
||
# 构造 path 查询（自动命名 path 变量）
|
||
path_var = "path"
|
||
path_pattern = f"{path_var} = {first_pattern}"
|
||
|
||
# 强制关系无向化（防御性）
|
||
path_pattern = re.sub(r'<?-\[', '-[', path_pattern)
|
||
path_pattern = re.sub(r'\]->?', ']-', path_pattern)
|
||
|
||
path_cypher = f"MATCH {path_pattern}{where_clause} RETURN DISTINCT {path_var}"
|
||
aggregate_cypher = cypher
|
||
|
||
logger.debug(f"[Cypher处理] Path查询: {path_cypher}")
|
||
return path_cypher, aggregate_cypher
|
||
|
||
except Exception as e:
|
||
logger.exception(f"[Cypher处理] 生成 path 查询失败: {e}")
|
||
return None, None
|
||
|
||
def _generate_simple_template_cypher(
|
||
self,
|
||
entry_nodes: dict,
|
||
query: str,
|
||
is_aggregate: bool
|
||
) -> Optional[str]:
|
||
"""
|
||
生成简单的模板Cypher查询（降级策略）
|
||
"""
|
||
if not entry_nodes:
|
||
return None
|
||
|
||
try:
|
||
# 选择第一个标签和第一个节点作为起点
|
||
first_label = list(entry_nodes.keys())[0]
|
||
first_nodes = entry_nodes[first_label]
|
||
if not first_nodes:
|
||
return None
|
||
|
||
first_node = first_nodes[0]
|
||
node_name = first_node.get("name") or first_node.get("_name") or (first_node.get("node", {}).get("名称"))
|
||
if not node_name:
|
||
return None
|
||
|
||
# 生成简单的查询：只查询入口节点及其直接邻居
|
||
if is_aggregate:
|
||
# 统计类：统计邻居数量
|
||
cypher = (
|
||
f"MATCH path=(n:{first_label} {{名称: '{node_name}'}})-[]-(m)\n"
|
||
f"RETURN path, m"
|
||
)
|
||
else:
|
||
# 明细类：返回邻居节点
|
||
cypher = (
|
||
f"MATCH path=(n:{first_label} {{名称: '{node_name}'}})-[]-(m)\n"
|
||
f"RETURN path, m\n"
|
||
f"LIMIT 50"
|
||
)
|
||
|
||
logger.debug(f"[降级策略] 生成简单模板查询: {cypher}")
|
||
return cypher
|
||
except Exception as e:
|
||
logger.warning(f"[降级策略] 生成简单模板查询失败: {e}")
|
||
return None
|
||
|
||
|
||
def _extract_path_info(self, record: Dict[str, Any]) -> List[str]:
|
||
"""
|
||
从单个记录中提取路径相关信息，用于rerank
|
||
"""
|
||
path_texts = []
|
||
|
||
# 遍历记录中的所有值，查找路径信息
|
||
for key, value in record.items():
|
||
if value is None:
|
||
continue
|
||
|
||
# 如果值是路径对象
|
||
if hasattr(value, 'nodes') and hasattr(value, 'relationships'):
|
||
# 提取路径中的节点和关系信息
|
||
try:
|
||
# 处理路径中的节点
|
||
for node in value.nodes:
|
||
node_info = f"节点: {dict(node).get('name', dict(node).get('名称', '未知节点'))} 类型: {list(node.labels)}"
|
||
path_texts.append(node_info)
|
||
|
||
# 处理路径中的关系
|
||
for rel in value.relationships:
|
||
rel_info = f"关系: {rel.type}"
|
||
path_texts.append(rel_info)
|
||
except Exception as e:
|
||
logger.warning(f"提取路径信息时出错: {e}")
|
||
|
||
# 如果值是字典，可能是节点或关系对象
|
||
elif isinstance(value, dict):
|
||
if 'labels' in value or 'type' in value:
|
||
# 这可能是节点或关系
|
||
if 'labels' in value: # 节点
|
||
node_info = f"节点: {value.get('name', value.get('名称', '未知节点'))} 类型: {value.get('labels', [])}"
|
||
path_texts.append(node_info)
|
||
elif 'type' in value: # 关系
|
||
rel_info = f"关系: {value.get('type')}"
|
||
path_texts.append(rel_info)
|
||
|
||
# 如果值是列表，可能包含多个节点或关系
|
||
elif isinstance(value, list):
|
||
for item in value:
|
||
if isinstance(item, dict):
|
||
if 'labels' in item: # 节点
|
||
node_info = f"节点: {item.get('name', item.get('名称', '未知节点'))} 类型: {item.get('labels', [])}"
|
||
path_texts.append(node_info)
|
||
elif 'type' in item: # 关系
|
||
rel_info = f"关系: {item.get('type')}"
|
||
path_texts.append(rel_info)
|
||
|
||
return path_texts
|
||
|
||
def _rerank_results(self, results: List[Dict[str, Any]], query: str, is_aggregate: bool = False) -> List[
|
||
Dict[str, Any]]:
|
||
"""
|
||
使用 model_api 中的 rerank 对结果进行重排序。
|
||
"""
|
||
if not results:
|
||
return results
|
||
|
||
# ✅ 统计类查询：不 rerank，直接赋分 1.0
|
||
if is_aggregate:
|
||
logger.info("检测到统计类查询，跳过 rerank，统一设置 rerank_score=1.0")
|
||
reranked_results = []
|
||
for result in results:
|
||
result_copy = result.copy()
|
||
result_copy['rerank_score'] = 1.0
|
||
reranked_results.append(result_copy)
|
||
return reranked_results
|
||
|
||
# === 以下为原语义 rerank 逻辑 ===
|
||
try:
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
|
||
documents = []
|
||
for result in results:
|
||
path_texts = self._extract_path_info(result)
|
||
doc_text = " ".join(path_texts) if path_texts else str(result)
|
||
documents.append(doc_text)
|
||
|
||
rerank_results = OpenaiAPI.rerank_query(query, documents, top_n=len(documents))
|
||
|
||
reranked_results = []
|
||
sorted_rerank_results = sorted(rerank_results, key=lambda x: x.get('scores', 0), reverse=True)
|
||
for rerank_result in sorted_rerank_results:
|
||
original_index = rerank_result['index']
|
||
if original_index < len(results):
|
||
result_copy = results[original_index].copy()
|
||
result_copy['rerank_score'] = rerank_result['scores']
|
||
reranked_results.append(result_copy)
|
||
|
||
return reranked_results
|
||
|
||
except Exception as e:
|
||
logger.error(f"重排序过程中出错: {e}")
|
||
# 失败时也统一加默认分数（保持接口一致）
|
||
fallback_results = []
|
||
for result in results:
|
||
result_copy = result.copy()
|
||
result_copy['rerank_score'] = 0.0
|
||
fallback_results.append(result_copy)
|
||
return fallback_results
|
||
|