kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix(lightrag): 修复查询结果节点数量未限制的问题

Browse Source
This commit is contained in:
Wenjie Zhang 2025-12-23 23:01:58 +08:00
parent cd34ea5474
commit 4eb2c4977e
1 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/knowledge/adapters/lightrag.py
View File

@ -42,7 +42,7 @@ class LightRAGGraphAdapter(GraphAdapter):
try: try:
with self._db.driver.session() as session: with self._db.driver.session() as session:
result = session.run(query, keyword=keyword, kb_id=kb_id, limit=limit) result = session.run(query, keyword=keyword, kb_id=kb_id, limit=limit)
return self._process_query_result(result) return self._process_query_result(result, limit=limit)
except Exception as e: except Exception as e:
logger.error(f"Neo4j query failed: {e}") logger.error(f"Neo4j query failed: {e}")
return {"nodes": [], "edges": []} return {"nodes": [], "edges": []}
@ -280,14 +280,18 @@ class LightRAGGraphAdapter(GraphAdapter):
return query return query
def _process_query_result(self, result) -> dict[str, list]: def _process_query_result(self, result, limit: int = None) -> dict[str, list]:
"""处理查询结果""" """处理查询结果，并限制节点数量不超过 limit"""
nodes = [] nodes = []
edges = [] edges = []
node_ids = set() node_ids = set()
edge_ids = set() edge_ids = set()
for record in result: for record in result:
# 检查是否已达到节点限制
if limit is not None and len(node_ids) >= limit:
break
for key in record.keys(): for key in record.keys():
val = record[key] val = record[key]
if val is None: if val is None:
@ -295,6 +299,9 @@ class LightRAGGraphAdapter(GraphAdapter):
if hasattr(val, "element_id") and hasattr(val, "labels"): # Node if hasattr(val, "element_id") and hasattr(val, "labels"): # Node
if val.element_id not in node_ids: if val.element_id not in node_ids:
# 再次检查限制
if limit is not None and len(node_ids) >= limit:
break
nodes.append(self.normalize_node(val)) nodes.append(self.normalize_node(val))
node_ids.add(val.element_id) node_ids.add(val.element_id)
elif hasattr(val, "element_id") and hasattr(val, "start_node"): # Relationship elif hasattr(val, "element_id") and hasattr(val, "start_node"): # Relationship
@ -305,6 +312,8 @@ class LightRAGGraphAdapter(GraphAdapter):
for item in val: for item in val:
if hasattr(item, "element_id") and hasattr(item, "labels"): if hasattr(item, "element_id") and hasattr(item, "labels"):
if item.element_id not in node_ids: if item.element_id not in node_ids:
if limit is not None and len(node_ids) >= limit:
break
nodes.append(self.normalize_node(item)) nodes.append(self.normalize_node(item))
node_ids.add(item.element_id) node_ids.add(item.element_id)