From e3fd24f383893b765806d388e74bebf9a2abd6ad Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Sat, 8 Mar 2025 22:49:13 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20indexing=20nodes?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/core/database.py       | 14 ++++++++++++
 src/core/graphbase.py      | 47 ++++++++++++++++++++++++++++++++++++++
 src/routers/data_router.py | 24 +++++++++++++++++++
 3 files changed, 85 insertions(+)

diff --git a/src/core/database.py b/src/core/database.py
index 0c190d0f..5676d4fb 100644
--- a/src/core/database.py
+++ b/src/core/database.py
@@ -87,6 +87,20 @@ class DataBaseManager:
         else:
             return {"message": "Graph base not enabled", "graph": {}}
 
+    def is_graph_running(self):
+        """检查图数据库是否正在运行
+
+        Returns:
+            bool: 图数据库是否正在运行
+        """
+        # 检查是否启用了图数据库
+        if not self.config.enable_knowledge_graph or not hasattr(self, 'graph_base') or self.graph_base is None:
+            return False
+
+        # 获取图数据库信息，检查状态
+        graph_info = self.graph_base.get_database_info("neo4j")
+        return graph_info.get("status") == "open"
+
     def create_database(self, database_name, description, db_type, dimension):
         from src.config import EMBED_MODEL_INFO
         dimension = dimension or EMBED_MODEL_INFO[self.config.embed_model]["dimension"]
diff --git a/src/core/graphbase.py b/src/core/graphbase.py
index 126ca5cc..65f1d7c5 100644
--- a/src/core/graphbase.py
+++ b/src/core/graphbase.py
@@ -377,6 +377,25 @@ class GraphDatabase:
             logger.error(f"保存图数据库信息失败：{e}")
             return False
 
+    def query_nodes_without_embedding(self, kgdb_name='neo4j'):
+        """查询没有嵌入向量的节点
+
+        Returns:
+            list: 没有嵌入向量的节点列表
+        """
+        self.use_database(kgdb_name)
+
+        def query(tx):
+            result = tx.run("""
+            MATCH (n:Entity)
+            WHERE n.embedding IS NULL
+            RETURN n.name AS name
+            """)
+            return [record["name"] for record in result]
+
+        with self.driver.session() as session:
+            return session.execute_read(query)
+
     def load_graph_info(self):
         """
         从工作目录中的JSON文件加载图数据库的基本信息
@@ -404,6 +423,34 @@ class GraphDatabase:
             logger.error(f"加载图数据库信息失败：{e}")
             return False
 
+    def add_embedding_to_nodes(self, node_names=None, kgdb_name='neo4j'):
+        """为节点添加嵌入向量
+
+        Args:
+            node_names (list, optional): 要添加嵌入向量的节点名称列表，None表示所有没有嵌入向量的节点
+            kgdb_name (str, optional): 图数据库名称，默认为'neo4j'
+
+        Returns:
+            int: 成功添加嵌入向量的节点数量
+        """
+        self.use_database(kgdb_name)
+
+        # 如果node_names为None，则获取所有没有嵌入向量的节点
+        if node_names is None:
+            node_names = self.query_nodes_without_embedding(kgdb_name)
+
+        count = 0
+        with self.driver.session() as session:
+            for node_name in node_names:
+                try:
+                    embedding = self.get_embedding(node_name)
+                    session.execute_write(self.set_embedding, node_name, embedding)
+                    count += 1
+                except Exception as e:
+                    logger.error(f"为节点 '{node_name}' 添加嵌入向量失败: {e}")
+
+        return count
+
 
 if __name__ == "__main__":
     pass
\ No newline at end of file
diff --git a/src/routers/data_router.py b/src/routers/data_router.py
index 06f24403..d2bb3417 100644
--- a/src/routers/data_router.py
+++ b/src/routers/data_router.py
@@ -105,8 +105,32 @@ async def upload_file(file: UploadFile = File(...)):
 @data.get("/graph")
 async def get_graph_info():
     graph_info = startup.dbm.get_graph()
+
+    # 获取未索引节点数量
+    unindexed_count = 0
+    if startup.dbm.is_graph_running():
+        # 调用GraphDatabase的query_nodes_without_embedding方法
+        unindexed_nodes = startup.dbm.graph_base.query_nodes_without_embedding()
+        unindexed_count = len(unindexed_nodes) if unindexed_nodes else 0
+
+    # 将未索引节点数量添加到返回结果中
+    graph_info["graph"]["unindexed_node_count"] = unindexed_count
+
     return graph_info
 
+@data.post("/graph/index-nodes")
+async def index_nodes(data: dict = Body(default={})):
+    if not startup.dbm.is_graph_running():
+        raise HTTPException(status_code=400, detail="图数据库未启动")
+
+    # 获取参数或使用默认值
+    kgdb_name = data.get('kgdb_name', 'neo4j')
+
+    # 调用GraphDatabase的add_embedding_to_nodes方法
+    count = startup.dbm.graph_base.add_embedding_to_nodes(kgdb_name=kgdb_name)
+
+    return {"status": "success", "message": f"已成功为{count}个节点添加嵌入向量", "indexed_count": count}
+
 @data.get("/graph/node")
 async def get_graph_node(entity_name: str):
     result = startup.dbm.graph_base.query_node(entity_name=entity_name)