diff --git a/docker/docker-compose.dev.yml b/docker/docker-compose.dev.yml index 0232dc2d..dfd09a6f 100644 --- a/docker/docker-compose.dev.yml +++ b/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 diff --git a/src/core/graphbase.py b/src/core/graphbase.py index c1e6ecc8..193b99a7 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -218,7 +218,19 @@ class GraphDatabase: 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 @@ class GraphDatabase: """, 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] diff --git a/src/core/retriever.py b/src/core/retriever.py index 3181a870..ee52d915 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -76,6 +76,8 @@ class Retriever: 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)