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

Author: xerrors

src/knowledge/adapters/lightrag.py

@@ -42,7 +42,7 @@ async def query_nodes(self, keyword: str, **kwargs) -> dict[str, Any]:
         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,21 +280,28 @@ def _build_subgraph_query(self, limit: int, kb_id: str = None) -> str:
 
         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:
                     continue
 
                 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 @@ def _process_query_result(self, result) -> dict[str, list]:
                     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)