From 4eb2c4977e3acded1bd9b0995eb92328ed73cca5 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 23 Dec 2025 23:01:58 +0800 Subject: [PATCH] =?UTF-8?q?fix(lightrag):=20=E4=BF=AE=E5=A4=8D=E6=9F=A5?= =?UTF-8?q?=E8=AF=A2=E7=BB=93=E6=9E=9C=E8=8A=82=E7=82=B9=E6=95=B0=E9=87=8F?= =?UTF-8?q?=E6=9C=AA=E9=99=90=E5=88=B6=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/knowledge/adapters/lightrag.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/knowledge/adapters/lightrag.py b/src/knowledge/adapters/lightrag.py index 5afcfd5a..e3c7ae71 100644 --- a/src/knowledge/adapters/lightrag.py +++ b/src/knowledge/adapters/lightrag.py @@ -42,7 +42,7 @@ class LightRAGGraphAdapter(GraphAdapter): try: with self._db.driver.session() as session: 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: logger.error(f"Neo4j query failed: {e}") return {"nodes": [], "edges": []} @@ -280,14 +280,18 @@ class LightRAGGraphAdapter(GraphAdapter): 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 = [] edges = [] node_ids = set() edge_ids = set() for record in result: + # 检查是否已达到节点限制 + if limit is not None and len(node_ids) >= limit: + break + for key in record.keys(): val = record[key] if val is None: @@ -295,6 +299,9 @@ class LightRAGGraphAdapter(GraphAdapter): if hasattr(val, "element_id") and hasattr(val, "labels"): # Node 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)) node_ids.add(val.element_id) elif hasattr(val, "element_id") and hasattr(val, "start_node"): # Relationship @@ -305,6 +312,8 @@ class LightRAGGraphAdapter(GraphAdapter): for item in val: if hasattr(item, "element_id") and hasattr(item, "labels"): 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)) node_ids.add(item.element_id)