From fe6c9ea364235ba233bed90de5e9ccada867e70c Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 19 Aug 2025 13:08:41 +0800 Subject: [PATCH] =?UTF-8?q?feat(=E7=9F=A5=E8=AF=86=E5=BA=93):=20=E5=AE=9E?= =?UTF-8?q?=E7=8E=B0=E7=9F=A5=E8=AF=86=E5=BA=93=E5=88=A0=E9=99=A4=E6=97=B6?= =?UTF-8?q?=E6=B8=85=E7=90=86Milvus=E5=92=8CNeo4j=E6=95=B0=E6=8D=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 添加delete_database方法，在删除知识库时同步清理Milvus集合和Neo4j节点关系 移除graphbase.py中未使用的txt_add_entity和format_query_result_to_graph方法 更新README.md中关于Neo4j节点标签和关系类型的说明 --- README.md | 2 +- src/knowledge/graphbase.py | 79 ------------------------------------ src/knowledge/lightrag_kb.py | 55 ++++++++++++++++++++++++- src/knowledge/milvus_kb.py | 17 +++++++- 4 files changed, 71 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index 24ac863d..5b7b50e6 100644 --- a/README.md +++ b/README.md @@ -175,7 +175,7 @@ custom-provider-name-here: {"h": "上海", "t": "中国", "r": "直辖市"} ``` -此外，也可以通过修改 `docker-compose.yml` 中的 `NEO4J_URI` 配置来接入已有的 Neo4j 实例，**但是**最好确保每个节点都有 Entity 标签，否则会影响到图的检索与构建。 +此外，也可以通过修改 `docker-compose.yml` 中的 `NEO4J_URI` 配置来接入已有的 Neo4j 实例，**但是**最好确保每个节点都有 Entity 标签，每个关系都有 `RELATION` 类型，否则会影响到图的检索与构建。 注：在“图谱”页面，只能看到上传的节点和边，基于 LightRAG 构建的图谱不会展示在里面，完整的图谱可以去 Neo4j 管理页面查看。 diff --git a/src/knowledge/graphbase.py b/src/knowledge/graphbase.py index 6c1fc15f..f7ff8672 100644 --- a/src/knowledge/graphbase.py +++ b/src/knowledge/graphbase.py @@ -103,24 +103,6 @@ class GraphDatabase: if self.status == "closed": self.start() - def txt_add_entity(self, triples, kgdb_name='neo4j'): - """添加实体三元组""" - assert self.driver is not None, "Database is not connected" - self.use_database(kgdb_name) - def create(tx, triples): - for triple in triples: - h = triple['h'] - t = triple['t'] - r = triple['r'] - query = ( - "MERGE (a:Entity:Upload {name: $h}) " - "MERGE (b:Entity:Upload {name: $t}) " - "MERGE (a)-[:" + r.replace(" ", "_") + "]->(b)" - ) - tx.run(query, h=h, t=t) - - with self.driver.session() as session: - session.execute_write(create, triples) async def txt_add_vector_entity(self, triples, kgdb_name='neo4j'): """添加实体三元组""" @@ -583,67 +565,6 @@ class GraphDatabase: formatted_results = {"nodes": nodes, "edges": edges} return formatted_results - def format_query_result_to_graph(self, query_results): - """将检索到的结果转换为 {"nodes": [], "edges": []} 的格式 - - 例如： - { - "nodes": [ - { - "id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:103", - "name": "张若锦" - }, - { - "id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:20", - "name": "贾宝玉" - }, - .... - ], - "edges": [ - { - "id": "5:5efbff88-72ef-44f9-b867-6c0e164a4a13:71", - "type": "奴仆", - "source_id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:88", - "target_id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:20", - "source_name": "宋嬷嬷", - "target_name": "贾宝玉" - }, - .... - ] - } - """ - formatted_results = {"nodes": [], "edges": []} - node_dict = {} - edge_dict = {} - - for item in query_results: - # 检查数据格式 - if len(item) < 2 or not isinstance(item[1], list): - continue - - node_dict[item[0].element_id] = dict(id=item[0].element_id, name=item[0]._properties.get("name", "Unknown")) - node_dict[item[2].element_id] = dict(id=item[2].element_id, name=item[2]._properties.get("name", "Unknown")) - - # 处理关系列表中的每个关系 - for i, relationship in enumerate(item[1]): - try: - # 提取关系信息 - node_info, edge_info = self._extract_relationship_info(relationship, node_dict=node_dict) - if node_info is None or edge_info is None: - continue - - # 添加边 - edge_dict[edge_info["id"]] = edge_info - except Exception as e: - logger.error(f"处理关系时出错: {e}, 关系: {relationship}, {traceback.format_exc()}") - continue - - # 将节点字典转换为列表 - formatted_results["nodes"] = list(node_dict.values()) - formatted_results["edges"] = list(edge_dict.values()) - - - return formatted_results def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None): """ diff --git a/src/knowledge/lightrag_kb.py b/src/knowledge/lightrag_kb.py index 5fbfc740..174350c3 100644 --- a/src/knowledge/lightrag_kb.py +++ b/src/knowledge/lightrag_kb.py @@ -6,11 +6,13 @@ from lightrag import LightRAG, QueryParam from lightrag.llm.openai import openai_complete_if_cache, openai_embed from lightrag.utils import EmbeddingFunc, setup_logger from lightrag.kg.shared_storage import initialize_pipeline_status +from pymilvus import connections, utility +from neo4j import GraphDatabase from src.knowledge.knowledge_base import KnowledgeBase from src.knowledge.indexing import process_url_to_markdown, process_file_to_markdown from src.knowledge.kb_utils import prepare_item_metadata, get_embedding_config -from src.utils import logger +from src.utils import logger, hashstr LIGHTRAG_LLM_PROVIDER = os.getenv("LIGHTRAG_LLM_PROVIDER", "siliconflow") @@ -46,6 +48,57 @@ class LightRagKB(KnowledgeBase): """知识库类型标识""" return "lightrag" + def delete_database(self, db_id: str) -> dict: + """删除数据库，同时清除Milvus和Neo4j中的数据""" + # Drop Milvus collection + try: + milvus_uri = os.getenv('MILVUS_URI', 'http://localhost:19530') + milvus_token = os.getenv('MILVUS_TOKEN', '') + connection_alias = f"lightrag_{hashstr(db_id, 6)}" + + connections.connect( + alias=connection_alias, + uri=milvus_uri, + token=milvus_token + ) + + # 删除 LightRAG 创建的三个集合 + collection_names = [f"{db_id}_chunks", f"{db_id}_relationships", f"{db_id}_entities"] + for collection_name in collection_names: + if utility.has_collection(collection_name, using=connection_alias): + utility.drop_collection(collection_name, using=connection_alias) + logger.info(f"Dropped Milvus collection {collection_name}") + else: + logger.info(f"Milvus collection {collection_name} does not exist, skipping") + + connections.disconnect(connection_alias) + except Exception as e: + logger.error(f"Failed to drop Milvus collection {db_id}: {e}") + + # Delete Neo4j data + neo4j_uri = os.getenv('NEO4J_URI', 'bolt://localhost:7687') + neo4j_username = os.getenv('NEO4J_USERNAME', 'neo4j') + neo4j_password = os.getenv('NEO4J_PASSWORD', '0123456789') + + try: + driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_username, neo4j_password)) + with driver.session() as session: + # 删除带有特定 db_id 标签的节点和关系 + session.run(""" + MATCH (n:`""" + db_id + """`) + DETACH DELETE n + """) + + logger.info(f"Deleted Neo4j nodes and relationships for workspace {db_id}") + except Exception as e: + logger.error(f"Failed to delete Neo4j data for {db_id}: {e}") + finally: + if 'driver' in locals(): + driver.close() + + # Delete local files and metadata + return super().delete_database(db_id) + async def _create_kb_instance(self, db_id: str, kb_config: dict) -> LightRAG: """创建 LightRAG 实例""" logger.info(f"Creating LightRAG instance for {db_id}") diff --git a/src/knowledge/milvus_kb.py b/src/knowledge/milvus_kb.py index e91a0a31..5755bba9 100644 --- a/src/knowledge/milvus_kb.py +++ b/src/knowledge/milvus_kb.py @@ -439,10 +439,25 @@ class MilvusKB(KnowledgeBase): return {"lines": []} + def delete_database(self, db_id: str) -> dict: + """删除数据库，同时清除Milvus中的集合""" + # Drop Milvus collection + try: + if utility.has_collection(db_id, using=self.connection_alias): + utility.drop_collection(db_id, using=self.connection_alias) + logger.info(f"Dropped Milvus collection for {db_id}") + else: + logger.info(f"Milvus collection {db_id} does not exist, skipping") + except Exception as e: + logger.error(f"Failed to drop Milvus collection {db_id}: {e}") + + # Call base method to delete local files and metadata + return super().delete_database(db_id) + def __del__(self): """清理连接""" try: if hasattr(self, 'connection_alias'): connections.disconnect(self.connection_alias) - except Exception: + except Exception: # noqa: S110 pass