285 lines
11 KiB
Python
285 lines
11 KiB
Python
"""
|
||
===========================================
|
||
Cypher 执行模块 - 拆分 + 拆分执行
|
||
===========================================
|
||
|
||
从 GraphRetrieval 抽取出的两个独立纯函数：
|
||
|
||
- split_aggregate_cypher(cypher)
|
||
把一条 Cypher 拆成 (path_cypher, aggregate_cypher)。
|
||
纯字符串处理，只依赖 re 和 logger，不依赖 self。
|
||
|
||
- execute_cypher_with_split(cypher, is_aggregate, driver, split_fn=...)
|
||
执行 Cypher，支持拆分为 path + aggregate 两个查询。
|
||
依赖通过参数注入：driver 负责执行，split_fn 负责拆分
|
||
（默认就是上面的 split_aggregate_cypher）。
|
||
|
||
这样 GraphRetrieval 里的两个方法都可以变成一行委托调用，测试也直接测这两个函数。
|
||
"""
|
||
import re
|
||
import logging
|
||
from typing import Callable, Optional, Tuple, Any,List,Dict
|
||
from neo4j import AsyncGraphDatabase, GraphDatabase
|
||
from graph_search.result_formatter import format_results,serialize_graph_for_llm
|
||
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
def split_aggregate_cypher(cypher: str) -> Tuple[Optional[str], Optional[str]]:
|
||
"""
|
||
这个函数只处理字符串，不需要改为异步。
|
||
"""
|
||
if not cypher or not cypher.strip():
|
||
return None, None
|
||
try:
|
||
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_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_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
|
||
|
||
|
||
# 2. 将执行函数改为 async
|
||
async def execute_cypher_with_split(
|
||
cypher: str,
|
||
is_aggregate: bool,
|
||
driver: Any,
|
||
split_fn: Callable[[str], Tuple[Optional[str], Optional[str]]] = split_aggregate_cypher,
|
||
) -> Tuple[Any, Optional[Any]]:
|
||
"""
|
||
异步执行 Cypher，支持拆分为 path + aggregate 两个查询。
|
||
"""
|
||
path_cypher, aggregate_cypher = split_fn(cypher)
|
||
|
||
if path_cypher and aggregate_cypher:
|
||
logger.info("[Cypher拆分] 执行 path 和 aggregate 查询")
|
||
path_records = None
|
||
|
||
try:
|
||
# 3. 使用 await 执行 path 查询
|
||
path_result = await driver.execute_query(path_cypher)
|
||
path_records = path_result.records
|
||
except Exception as e:
|
||
logger.warning(f"[Cypher拆分] Path 查询失败: {e}")
|
||
path_records = None
|
||
|
||
try:
|
||
# 4. 使用 await 执行 aggregate 查询
|
||
aggregate_result = await driver.execute_query(aggregate_cypher)
|
||
records = aggregate_result.records
|
||
except Exception as e:
|
||
logger.error(f"[Cypher拆分] Aggregate 查询失败: {e}，回退到原始 Cypher")
|
||
# 回退也需异步执行
|
||
result = await driver.execute_query(cypher)
|
||
return result.records, None
|
||
|
||
return records, path_records
|
||
else:
|
||
# 5. 使用 await 执行原始查询
|
||
result = await driver.execute_query(cypher)
|
||
return result.records, None
|
||
|
||
|
||
class ResultProcessor:
|
||
"""查询结果处理器 - 保持不变，因为只是数据处理"""
|
||
def __init__(self):
|
||
pass
|
||
|
||
def process_results(self, records: List, path_records: Optional[List], is_aggregate: bool = False) -> List[Dict[str, Any]]:
|
||
if is_aggregate:
|
||
return self._process_aggregate_results(records, path_records)
|
||
else:
|
||
return self._process_detail_results(records, path_records)
|
||
|
||
def _process_aggregate_results(self, records: List, path_records: Optional[List]) -> List[Dict[str, Any]]:
|
||
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
|
||
logger.debug(f"[结果处理] 聚合查询，返回 {len(raw_results)} 条path记录")
|
||
return raw_results
|
||
else:
|
||
logger.debug(f"[结果处理] 聚合查询，返回 {len(aggregate_results)} 条聚合结果")
|
||
return aggregate_results
|
||
|
||
def _process_detail_results(self, records: List, path_records: Optional[List]) -> List[Dict[str, Any]]:
|
||
value_results = [dict(r) for r in records]
|
||
|
||
if path_records:
|
||
raw_results = [dict(r) for r in path_records]
|
||
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
|
||
logger.debug(f"[结果处理] 明细查询，返回 {len(raw_results)} 条path记录")
|
||
return raw_results
|
||
else:
|
||
logger.debug(f"[结果处理] 明细查询，返回 {len(value_results)} 条明细结果")
|
||
return value_results
|
||
|
||
|
||
def process_neo4j_results(records: List, path_records: Optional[List] = None, is_aggregate: bool = False) -> List[Dict[str, Any]]:
|
||
"""便捷函数"""
|
||
processor = ResultProcessor()
|
||
return processor.process_results(records, path_records, is_aggregate)
|
||
|
||
async def excute_cypher(driver,cypher):
|
||
# driver = AsyncGraphDatabase.driver(URI, auth=AUTH)
|
||
|
||
try:
|
||
query = """
|
||
MATCH (d:设备)-[:发生故障使用维修工作]->(w:维修工作)-[:包含故障]->(f:故障模式)<-[:修理]-(p:维修项目)
|
||
WHERE d.名称 = '喷油器组件'
|
||
AND w.名称 = '喷油器组件维修工作指导'
|
||
AND f.名称 = '喷油器工作不一致导致机座与机架组件功率输出不稳定'
|
||
AND p.名称 = '排除喷油器雾化不良故障'
|
||
RETURN w.名称 AS 维修工作名称, f.名称 AS 故障模式, p.名称 AS 维修项目名称
|
||
"""
|
||
|
||
# 7. 调用异步函数
|
||
records, path_records = await execute_cypher_with_split(cypher, True, driver)
|
||
results = process_neo4j_results(
|
||
records=records,
|
||
path_records=path_records,
|
||
is_aggregate=True
|
||
)
|
||
|
||
# 这里假设 format_results 和 serialize_graph_for_llm 是同步的工具函数
|
||
# 如果它们也是异步的，请使用 await
|
||
path_data = format_results(results)
|
||
|
||
temp_graph_response = {
|
||
"data": {
|
||
"nodes": path_data.get("nodes", []),
|
||
"results": results
|
||
}
|
||
}
|
||
|
||
# 假设 serialize_graph_for_llm 是同步的
|
||
formatted_results_str = serialize_graph_for_llm(temp_graph_response)
|
||
|
||
response = {
|
||
"data": {
|
||
"nodes": path_data.get("nodes", []),
|
||
"links": path_data.get("links", []),
|
||
"results": formatted_results_str
|
||
},
|
||
"meta": {
|
||
"result_count": len(formatted_results_str),
|
||
"query_type": "graph_search",
|
||
"cypher": cypher,
|
||
"query_classification": "graph_search"
|
||
}
|
||
}
|
||
return response
|
||
|
||
finally:
|
||
# 8. 关闭异步驱动
|
||
await driver.close()
|
||
|
||
# ========== 便捷函数和主程序入口 ==========
|
||
if __name__ == "__main__":
|
||
# 6. 配置异步驱动
|
||
URI = "bolt://192.168.0.46:57687"
|
||
AUTH = ("neo4j", "zdht123@")
|
||
|
||
# 定义异步主函数
|
||
async def main():
|
||
# 使用 AsyncGraphDatabase 初始化驱动
|
||
driver = AsyncGraphDatabase.driver(URI, auth=AUTH)
|
||
|
||
try:
|
||
query = """
|
||
MATCH (d:设备)-[:发生故障使用维修工作]->(w:维修工作)-[:包含故障]->(f:故障模式)<-[:修理]-(p:维修项目)
|
||
WHERE d.名称 = '喷油器组件'
|
||
AND w.名称 = '喷油器组件维修工作指导'
|
||
AND f.名称 = '喷油器工作不一致导致机座与机架组件功率输出不稳定'
|
||
AND p.名称 = '排除喷油器雾化不良故障'
|
||
RETURN w.名称 AS 维修工作名称, f.名称 AS 故障模式, p.名称 AS 维修项目名称
|
||
"""
|
||
|
||
# 7. 调用异步函数
|
||
records, path_records = await execute_cypher_with_split(query, True, driver)
|
||
results = process_neo4j_results(
|
||
records=records,
|
||
path_records=path_records,
|
||
is_aggregate=True
|
||
)
|
||
|
||
# 这里假设 format_results 和 serialize_graph_for_llm 是同步的工具函数
|
||
# 如果它们也是异步的，请使用 await
|
||
path_data = format_results(results)
|
||
|
||
temp_graph_response = {
|
||
"data": {
|
||
"nodes": path_data.get("nodes", []),
|
||
"results": results
|
||
}
|
||
}
|
||
|
||
# 假设 serialize_graph_for_llm 是同步的
|
||
formatted_results_str = serialize_graph_for_llm(temp_graph_response)
|
||
|
||
response = {
|
||
"code": 200,
|
||
"message": "success",
|
||
"data": {
|
||
"nodes": path_data.get("nodes", []),
|
||
"links": path_data.get("links", []),
|
||
"results": formatted_results_str
|
||
},
|
||
"meta": {
|
||
"result_count": len(formatted_results_str),
|
||
"query_type": "graph_search",
|
||
"elapsed_ms": 20,
|
||
"cypher": query,
|
||
"query_classification": "graph_search"
|
||
}
|
||
}
|
||
print(response)
|
||
|
||
finally:
|
||
# 8. 关闭异步驱动
|
||
await driver.close()
|
||
|
||
# 9. 运行异步主函数
|
||
import asyncio
|
||
asyncio.run(main())
|