feat(知识库): 实现知识库删除时清理Milvus和Neo4j数据
添加delete_database方法，在删除知识库时同步清理Milvus集合和Neo4j节点关系 移除graphbase.py中未使用的txt_add_entity和format_query_result_to_graph方法 更新README.md中关于Neo4j节点标签和关系类型的说明
This commit is contained in:
Wenjie Zhang
parent
e6145657c5
commit
fe6c9ea364
@ -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 管理页面查看。
|
||||
|
||||
|
||||
@ -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):
|
||||
"""
|
||||
|
||||
@ -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}")
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user