Merge pull request #117 from xerrors/112-kg-retrieve-error

fix：检索前检查是否存在 entityEmbeddings 这个索引

Author: xerrors

docker/docker-compose.dev.yml

@@ -60,6 +60,7 @@ services:
       - NEO4J_AUTH=neo4j/0123456789
       - NEO4J_server_bolt_listen__address=0.0.0.0:7687
       - NEO4J_server_http_listen__address=0.0.0.0:7474
+      - ENTITY_EMBEDDING=true
     networks:
       - app-network
 

src/core/graphbase.py

@@ -218,7 +218,19 @@ def query_node(self, entity_name, hops=2, **kwargs):
 
     def query_by_vector(self, entity_name, threshold=0.9, kgdb_name='neo4j', hops=2, num_of_res=5):
         self.use_database(kgdb_name)
+        def _index_exists(tx, index_name):
+            """检查索引是否存在"""
+            result = tx.run("SHOW INDEXES")
+            for record in result:
+                if record["name"] == index_name:
+                    return True
+            return False
+
         def query(tx, text):
+            # 首先检查索引是否存在
+            if not _index_exists(tx, "entityEmbeddings"):
+                raise Exception("向量索引不存在，请先创建索引")
+
             embedding = self.get_embedding(text)
             result = tx.run("""
             CALL db.index.vector.queryNodes('entityEmbeddings', 10, $embedding)
@@ -227,9 +239,14 @@ def query(tx, text):
             """, embedding=embedding)
             return result.values()
 
-        with self.driver.session() as session:
-            results = session.execute_read(query, entity_name)
-
+        try:
+            with self.driver.session() as session:
+                results = session.execute_read(query, entity_name)
+        except Exception as e:
+            if "向量索引不存在" in str(e):
+                logger.error(f"向量索引不存在，请先创建索引: {e}, {traceback.format_exc()}")
+                return []
+            raise e
 
         # 筛选出分数高于阈值的实体
         qualified_entities = [result[0] for result in results[:num_of_res] if result[1] > threshold]

src/core/retriever.py

@@ -76,6 +76,8 @@ def query_graph(self, query, history, refs):
         results = []
         if refs["meta"].get("use_graph") and config.enable_knowledge_base:
             for entity in refs["entities"]:
+                if entity == "":
+                    continue
                 result = graph_base.query_by_vector(entity)
                 if result != []:
                     results.extend(result)