855 lines
34 KiB
Python
855 lines
34 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
from typing import Any
|
||
|
||
from yuxi.knowledge.graphs.extractors import GraphExtractor, GraphExtractorFactory, normalize_extraction_result
|
||
from yuxi.knowledge.graphs.graph_utils import (
|
||
build_graph_payload,
|
||
compute_entity_id,
|
||
compute_triple_id,
|
||
cypher_merge_chunk,
|
||
cypher_merge_entity_mention,
|
||
cypher_merge_relation,
|
||
normalize_entity_name,
|
||
)
|
||
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
|
||
from yuxi.storage.neo4j import (
|
||
Neo4jConnectionManager,
|
||
get_shared_neo4j_connection,
|
||
neo4j_read,
|
||
neo4j_write,
|
||
safe_neo4j_label,
|
||
)
|
||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||
from yuxi.repositories.knowledge_graph_repository import KnowledgeGraphRepository
|
||
from yuxi.utils import logger
|
||
from yuxi.utils.datetime_utils import utc_isoformat
|
||
|
||
GRAPH_CONFIG_KEY = "graph_build_config"
|
||
GRAPH_TASK_TYPE = "knowledge_graph_index"
|
||
|
||
|
||
class MilvusGraphService:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
kb_id: str | None = None,
|
||
kb_repo: KnowledgeBaseRepository | None = None,
|
||
chunk_repo: KnowledgeChunkRepository | None = None,
|
||
graph_repo: KnowledgeGraphRepository | None = None,
|
||
graph_vector_store: MilvusGraphVectorStore | None = None,
|
||
neo4j_connection: Neo4jConnectionManager | None = None,
|
||
):
|
||
self.kb_id = kb_id
|
||
self.kb_repo = kb_repo or KnowledgeBaseRepository()
|
||
self.chunk_repo = chunk_repo or KnowledgeChunkRepository()
|
||
self.graph_repo = graph_repo or KnowledgeGraphRepository()
|
||
self._graph_vector_store = graph_vector_store
|
||
self._connection = neo4j_connection
|
||
|
||
@property
|
||
def connection(self) -> Neo4jConnectionManager:
|
||
if self._connection is None:
|
||
self._connection = get_shared_neo4j_connection()
|
||
return self._connection
|
||
|
||
@property
|
||
def graph_vector_store(self) -> MilvusGraphVectorStore:
|
||
if self._graph_vector_store is None:
|
||
self._graph_vector_store = MilvusGraphVectorStore()
|
||
return self._graph_vector_store
|
||
|
||
@property
|
||
def driver(self):
|
||
return self.connection.driver
|
||
|
||
async def get_status(self, kb_id: str, *, tasker: Any = None) -> dict[str, Any]:
|
||
kb = await self._get_milvus_kb(kb_id)
|
||
params = dict(kb.additional_params or {})
|
||
config = params.get(GRAPH_CONFIG_KEY) or {}
|
||
total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather(
|
||
self.chunk_repo.count_by_kb_id(kb_id),
|
||
self.chunk_repo.count_graph_pending_by_kb_id(kb_id),
|
||
self.chunk_repo.count_graph_indexed_by_kb_id(kb_id),
|
||
self.graph_repo.count_by_kb_id(kb_id),
|
||
)
|
||
entity_count, relationship_count = graph_counts
|
||
|
||
build_task_status = None
|
||
build_task_progress = 0
|
||
if tasker is not None:
|
||
active_task = await tasker.find_task_by_payload(
|
||
task_type=GRAPH_TASK_TYPE,
|
||
payload_match={"kb_id": kb_id},
|
||
statuses={"pending", "running"},
|
||
)
|
||
if active_task:
|
||
build_task_status = active_task.status
|
||
build_task_progress = round(active_task.progress)
|
||
else:
|
||
failed_task = await tasker.find_task_by_payload(
|
||
task_type=GRAPH_TASK_TYPE,
|
||
payload_match={"kb_id": kb_id},
|
||
statuses={"failed", "cancelled"},
|
||
)
|
||
if failed_task:
|
||
build_task_status = "failed"
|
||
build_task_progress = 0
|
||
|
||
return {
|
||
"kb_id": kb_id,
|
||
"kb_type": kb.kb_type,
|
||
"configured": bool(config),
|
||
"locked": bool(config.get("locked")),
|
||
"config": self._public_config(config),
|
||
"total_chunks": total_chunks,
|
||
"pending_chunks": pending_chunks,
|
||
"indexed_chunks": indexed_chunks,
|
||
"entity_count": entity_count,
|
||
"relationship_count": relationship_count,
|
||
"build_task_status": build_task_status,
|
||
"build_task_progress": build_task_progress,
|
||
}
|
||
|
||
async def configure(
|
||
self,
|
||
kb_id: str,
|
||
extractor_type: str,
|
||
extractor_options: dict[str, Any],
|
||
created_by: str,
|
||
) -> dict:
|
||
kb = await self._get_milvus_kb(kb_id)
|
||
additional_params = dict(kb.additional_params or {})
|
||
existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
|
||
normalized_extractor_type = (extractor_type or "").lower()
|
||
if existing_config.get("locked"):
|
||
existing_extractor_type = (existing_config.get("extractor_type") or "").lower()
|
||
if normalized_extractor_type != existing_extractor_type:
|
||
raise ValueError("图谱抽取器类型已锁定，只能修改模型、Schema 等抽取参数")
|
||
|
||
extractor_options = extractor_options or {}
|
||
if normalized_extractor_type == "llm" and extractor_options.get("prompt"):
|
||
raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt，请使用 schema 配置抽取约束")
|
||
GraphExtractorFactory.create(normalized_extractor_type, extractor_options)
|
||
config = {
|
||
"locked": True,
|
||
"extractor_type": normalized_extractor_type,
|
||
"extractor_options": extractor_options or {},
|
||
"created_at": existing_config.get("created_at") or utc_isoformat(),
|
||
"created_by": existing_config.get("created_by") or created_by,
|
||
}
|
||
if existing_config.get("locked"):
|
||
config["updated_at"] = utc_isoformat()
|
||
config["updated_by"] = created_by
|
||
additional_params[GRAPH_CONFIG_KEY] = config
|
||
await self.kb_repo.update(kb_id, {"additional_params": additional_params})
|
||
return config
|
||
|
||
async def build_pending_chunks(self, kb_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
|
||
kb = await self._get_milvus_kb(kb_id)
|
||
config = self._get_locked_config(kb.additional_params or {})
|
||
extractor_options = self._runtime_extractor_options(config)
|
||
extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
|
||
worker_count = self._get_worker_count(config)
|
||
total_pending = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
|
||
processed = 0
|
||
failed = 0
|
||
failed_chunk_ids: set[str] = set()
|
||
write_lock = asyncio.Lock()
|
||
|
||
while True:
|
||
if context is not None:
|
||
await context.raise_if_cancelled()
|
||
chunks = await self.chunk_repo.list_graph_pending_by_kb_id(kb_id, batch_size)
|
||
unprocessed = [c for c in chunks if c.chunk_id not in failed_chunk_ids]
|
||
if not unprocessed:
|
||
break
|
||
|
||
queue: asyncio.Queue[Any] = asyncio.Queue()
|
||
for chunk in unprocessed:
|
||
queue.put_nowait(chunk)
|
||
|
||
async def worker() -> None:
|
||
nonlocal processed, failed
|
||
while True:
|
||
if context is not None:
|
||
await context.raise_if_cancelled()
|
||
try:
|
||
chunk = queue.get_nowait()
|
||
except asyncio.QueueEmpty:
|
||
return
|
||
try:
|
||
extraction_result = await self._get_chunk_extraction_result(kb_id, chunk, extractor)
|
||
async with write_lock:
|
||
entities, triples = await asyncio.to_thread(
|
||
self.write_chunk_graph,
|
||
kb_id,
|
||
chunk,
|
||
extraction_result,
|
||
)
|
||
await self.graph_repo.upsert_chunk_graph(
|
||
kb_id=kb_id,
|
||
file_id=chunk.file_id,
|
||
chunk_id=chunk.chunk_id,
|
||
entities=entities,
|
||
triples=triples,
|
||
)
|
||
await self.graph_vector_store.insert_missing_graph_records(
|
||
kb_id=kb_id,
|
||
embedding_model_spec=kb.embedding_model_spec,
|
||
entities=entities,
|
||
triples=triples,
|
||
)
|
||
await self.chunk_repo.mark_graph_indexed(
|
||
chunk.chunk_id,
|
||
ent_ids=[entity["entity_id"] for entity in entities],
|
||
)
|
||
processed += 1
|
||
except Exception as exc:
|
||
logger.error(f"Chunk 图谱构建失败 chunk_id={chunk.chunk_id}: {exc}")
|
||
failed_chunk_ids.add(chunk.chunk_id)
|
||
failed += 1
|
||
finally:
|
||
queue.task_done()
|
||
|
||
if context is not None:
|
||
completed = processed + failed
|
||
progress = 5.0 + min(90.0, completed / max(total_pending, 1) * 90.0)
|
||
await context.set_progress(progress, f"图谱构建 {completed}/{total_pending}，失败 {failed}")
|
||
|
||
workers = [asyncio.create_task(worker()) for _ in range(min(worker_count, len(unprocessed)))]
|
||
try:
|
||
await asyncio.gather(*workers)
|
||
except Exception:
|
||
for task in workers:
|
||
task.cancel()
|
||
await asyncio.gather(*workers, return_exceptions=True)
|
||
raise
|
||
|
||
remaining = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
|
||
return {"kb_id": kb_id, "success": processed, "failed": failed, "remaining": remaining}
|
||
|
||
@staticmethod
|
||
def _get_worker_count(config: dict[str, Any]) -> int:
|
||
if (config.get("extractor_type") or "").lower() != "llm":
|
||
return 1
|
||
try:
|
||
worker_count = int((config.get("extractor_options") or {}).get("concurrency_count") or 1)
|
||
except (TypeError, ValueError):
|
||
return 1
|
||
return max(1, min(worker_count, 1000))
|
||
|
||
@staticmethod
|
||
def _runtime_extractor_options(config: dict[str, Any]) -> dict[str, Any]:
|
||
options = dict(config.get("extractor_options") or {})
|
||
options.pop("prompt", None)
|
||
return options
|
||
|
||
async def _get_chunk_extraction_result(self, kb_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]:
|
||
extractor_type = extractor.extractor_type
|
||
if chunk.extraction_result:
|
||
return normalize_extraction_result(chunk.extraction_result, extractor_type)
|
||
|
||
extraction_result = await extractor.extract(
|
||
chunk.content,
|
||
chunk_metadata={
|
||
"kb_id": kb_id,
|
||
"chunk_id": chunk.chunk_id,
|
||
"file_id": chunk.file_id,
|
||
"chunk_index": chunk.chunk_index,
|
||
},
|
||
)
|
||
normalized_result = normalize_extraction_result(extraction_result, extractor_type)
|
||
await self.chunk_repo.update_extraction_result(chunk.chunk_id, normalized_result)
|
||
return normalized_result
|
||
|
||
def write_chunk_graph(
|
||
self,
|
||
kb_id: str,
|
||
chunk,
|
||
normalized_result: dict[str, Any],
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
"""将单个 chunk 的抽取结果写入 Neo4j。"""
|
||
label = safe_neo4j_label(kb_id)
|
||
graph_payload = build_graph_payload(normalized_result)
|
||
relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
|
||
entities = graph_payload["entities"]
|
||
relations = graph_payload["relations"]
|
||
entity_by_id = {entity["id"]: entity for entity in entities}
|
||
entity_records = self._build_entity_records(kb_id, entities)
|
||
entity_record_by_local_id = {
|
||
entity["id"]: record for entity, record in zip(entities, entity_records, strict=True)
|
||
}
|
||
triple_records = self._build_triple_records(kb_id, relations, entity_record_by_local_id, graph_payload)
|
||
content_preview = (chunk.content or "")[:300]
|
||
|
||
# 预构建 Cypher 模板（同一 chunk 内复用）
|
||
merge_chunk_cypher = cypher_merge_chunk(label)
|
||
merge_entity_cypher = cypher_merge_entity_mention(label)
|
||
merge_relation_cypher = cypher_merge_relation(label)
|
||
|
||
def query(tx):
|
||
# 1. MERGE Chunk 节点
|
||
tx.run(
|
||
merge_chunk_cypher,
|
||
chunk_id=chunk.chunk_id,
|
||
file_id=chunk.file_id,
|
||
kb_id=kb_id,
|
||
chunk_index=chunk.chunk_index,
|
||
content_preview=content_preview,
|
||
start_char_pos=chunk.start_char_pos,
|
||
end_char_pos=chunk.end_char_pos,
|
||
)
|
||
|
||
# 2. MERGE Entity 节点 + Chunk→Entity (MENTIONS)
|
||
for entity in entities:
|
||
entity_record = entity_record_by_local_id[entity["id"]]
|
||
tx.run(
|
||
merge_entity_cypher,
|
||
chunk_id=chunk.chunk_id,
|
||
file_id=chunk.file_id,
|
||
kb_id=kb_id,
|
||
entity_id=entity_record["entity_id"],
|
||
normalized_name=normalize_entity_name(entity["text"]),
|
||
entity_label=entity.get("label") or "Entity",
|
||
name=entity["text"],
|
||
attributes=json.dumps(entity.get("attributes") or [], ensure_ascii=False),
|
||
)
|
||
|
||
# 3. MERGE Entity→Entity (RELATION) 边
|
||
for relation in relations:
|
||
source = entity_by_id[relation["source"]]
|
||
target = entity_by_id[relation["target"]]
|
||
source_record = entity_record_by_local_id[relation["source"]]
|
||
target_record = entity_record_by_local_id[relation["target"]]
|
||
relation_type = relation.get("label") or "RELATED_TO"
|
||
triple_id = compute_triple_id(
|
||
kb_id,
|
||
source_record["normalized_name"],
|
||
source_record["label"],
|
||
relation_type,
|
||
target_record["normalized_name"],
|
||
target_record["label"],
|
||
)
|
||
tx.run(
|
||
merge_relation_cypher,
|
||
kb_id=kb_id,
|
||
chunk_id=chunk.chunk_id,
|
||
file_id=chunk.file_id,
|
||
source_name=normalize_entity_name(source["text"]),
|
||
source_label=source.get("label") or "Entity",
|
||
target_name=normalize_entity_name(target["text"]),
|
||
target_label=target.get("label") or "Entity",
|
||
relation_type=relation_type,
|
||
triple_id=triple_id,
|
||
text=relation["text"],
|
||
extractor_type=relation_extractor_type,
|
||
)
|
||
|
||
neo4j_write(self.driver, query)
|
||
return entity_records, triple_records
|
||
|
||
def _build_entity_records(self, kb_id: str, entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||
records = []
|
||
for entity in entities:
|
||
label = entity.get("label") or "Entity"
|
||
normalized_name = normalize_entity_name(entity["text"])
|
||
entity_id = compute_entity_id(kb_id, normalized_name, label)
|
||
records.append(
|
||
{
|
||
"entity_id": entity_id,
|
||
"kb_id": kb_id,
|
||
"normalized_name": normalized_name,
|
||
"label": label,
|
||
"name": entity["text"],
|
||
"attributes": entity.get("attributes") or [],
|
||
"content": normalized_name,
|
||
}
|
||
)
|
||
return records
|
||
|
||
def _build_triple_records(
|
||
self,
|
||
kb_id: str,
|
||
relations: list[dict[str, Any]],
|
||
entity_record_by_local_id: dict[str, dict[str, Any]],
|
||
graph_payload: dict[str, Any],
|
||
) -> list[dict[str, Any]]:
|
||
records = []
|
||
seen_triple_ids: set[str] = set()
|
||
extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
|
||
for relation in relations:
|
||
source_record = entity_record_by_local_id[relation["source"]]
|
||
target_record = entity_record_by_local_id[relation["target"]]
|
||
relation_type = relation.get("label") or "RELATED_TO"
|
||
triple_id = compute_triple_id(
|
||
kb_id,
|
||
source_record["normalized_name"],
|
||
source_record["label"],
|
||
relation_type,
|
||
target_record["normalized_name"],
|
||
target_record["label"],
|
||
)
|
||
if triple_id in seen_triple_ids:
|
||
continue
|
||
seen_triple_ids.add(triple_id)
|
||
content = f"{source_record['normalized_name']} → {relation_type} → {target_record['normalized_name']}"
|
||
records.append(
|
||
{
|
||
"triple_id": triple_id,
|
||
"kb_id": kb_id,
|
||
"source_entity_id": source_record["entity_id"],
|
||
"target_entity_id": target_record["entity_id"],
|
||
"relation_type": relation_type,
|
||
"content": content,
|
||
"text": relation["text"],
|
||
"extractor_type": extractor_type,
|
||
}
|
||
)
|
||
return records
|
||
|
||
async def reset(self, kb_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
|
||
kb = await self._get_milvus_kb(kb_id)
|
||
await asyncio.to_thread(self.delete_graph, kb_id)
|
||
await self.graph_repo.delete_by_kb_id(kb_id)
|
||
reset_chunks = await self.chunk_repo.reset_graph_state_by_kb_id(kb_id, clear_extraction_result)
|
||
if clear_config:
|
||
additional_params = dict(kb.additional_params or {})
|
||
additional_params.pop(GRAPH_CONFIG_KEY, None)
|
||
await self.kb_repo.update(kb_id, {"additional_params": additional_params})
|
||
return {
|
||
"message": "图谱构建状态已重置",
|
||
"status": "success",
|
||
"reset_chunks": reset_chunks,
|
||
"clear_extraction_result": clear_extraction_result,
|
||
"clear_config": clear_config,
|
||
}
|
||
|
||
def delete_graph(self, kb_id: str) -> None:
|
||
label = safe_neo4j_label(kb_id)
|
||
|
||
def query(tx):
|
||
tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n")
|
||
|
||
neo4j_write(self.driver, query)
|
||
self.graph_vector_store.drop_graph_collections(kb_id)
|
||
|
||
async def delete_file_graph(self, kb_id: str, file_id: str) -> None:
|
||
orphan_entity_ids, orphan_triple_ids = await self.graph_repo.delete_file_references(file_id)
|
||
await self.graph_vector_store.delete_graph_records(
|
||
kb_id,
|
||
entity_ids=orphan_entity_ids,
|
||
triple_ids=orphan_triple_ids,
|
||
)
|
||
await asyncio.to_thread(self._delete_file_graph_from_neo4j, kb_id, file_id)
|
||
|
||
def _delete_file_graph_from_neo4j(self, kb_id: str, file_id: str) -> None:
|
||
label = safe_neo4j_label(kb_id)
|
||
|
||
def query(tx):
|
||
tx.run(
|
||
f"""
|
||
MATCH (:Chunk:MilvusKB:`{label}`)-[m:MENTIONS {{kb_id: $kb_id, file_id: $file_id}}]->
|
||
(:Entity:MilvusKB:`{label}`)
|
||
DELETE m
|
||
""",
|
||
kb_id=kb_id,
|
||
file_id=file_id,
|
||
)
|
||
tx.run(
|
||
f"""
|
||
MATCH (:Entity:MilvusKB:`{label}`)-[r:RELATION {{kb_id: $kb_id, file_id: $file_id}}]->
|
||
(:Entity:MilvusKB:`{label}`)
|
||
DELETE r
|
||
""",
|
||
kb_id=kb_id,
|
||
file_id=file_id,
|
||
)
|
||
tx.run(
|
||
f"""
|
||
MATCH (c:Chunk:MilvusKB:`{label}` {{kb_id: $kb_id, file_id: $file_id}})
|
||
DETACH DELETE c
|
||
""",
|
||
kb_id=kb_id,
|
||
file_id=file_id,
|
||
)
|
||
tx.run(
|
||
f"""
|
||
MATCH (e:Entity:MilvusKB:`{label}` {{kb_id: $kb_id}})
|
||
WHERE NOT ()-[:MENTIONS]->(e)
|
||
DETACH DELETE e
|
||
""",
|
||
kb_id=kb_id,
|
||
)
|
||
|
||
neo4j_write(self.driver, query)
|
||
|
||
async def query_nodes(
|
||
self,
|
||
kb_id: str | None = None,
|
||
*,
|
||
keyword: str = "",
|
||
max_depth: int = 1,
|
||
max_nodes: int = 50,
|
||
exclude_chunk: bool = False,
|
||
) -> dict[str, Any]:
|
||
effective_kb_id = kb_id or self.kb_id
|
||
if not effective_kb_id:
|
||
return {"nodes": [], "edges": []}
|
||
|
||
label = safe_neo4j_label(effective_kb_id)
|
||
limit = max_nodes
|
||
try:
|
||
with self.driver.session() as session:
|
||
result = session.run(
|
||
self._build_query(label, keyword, limit, max_depth, exclude_chunk),
|
||
keyword=keyword,
|
||
limit=limit,
|
||
)
|
||
return self._process_query_result(result, limit, effective_kb_id, exclude_chunk)
|
||
except Exception as e:
|
||
logger.error(f"Milvus graph query failed: {e}")
|
||
return {"nodes": [], "edges": []}
|
||
|
||
async def query_seed_subgraph(
|
||
self,
|
||
kb_id: str,
|
||
*,
|
||
entity_ids: list[str],
|
||
max_nodes: int,
|
||
) -> dict[str, Any]:
|
||
if not entity_ids:
|
||
return {"nodes": [], "edges": []}
|
||
label = safe_neo4j_label(kb_id)
|
||
cypher = f"""
|
||
MATCH (seed:Entity:MilvusKB:`{label}`)
|
||
WHERE seed.entity_id IN $entity_ids
|
||
MATCH p = (seed)-[*1..2]-(n:MilvusKB:`{label}`)
|
||
WITH p LIMIT $path_limit
|
||
WITH collect(p) AS paths
|
||
UNWIND paths AS node_path
|
||
UNWIND nodes(node_path) AS node
|
||
WITH paths, collect(DISTINCT node) AS graph_nodes
|
||
UNWIND paths AS rel_path
|
||
UNWIND relationships(rel_path) AS rel
|
||
RETURN graph_nodes AS nodes, collect(DISTINCT rel) AS edges
|
||
"""
|
||
try:
|
||
with self.driver.session() as session:
|
||
record = session.run(
|
||
cypher,
|
||
entity_ids=list(dict.fromkeys(entity_ids)),
|
||
path_limit=max(max_nodes, 1) * 4,
|
||
).single()
|
||
if not record:
|
||
return {"nodes": [], "edges": []}
|
||
return self._process_subgraph_record(record, max_nodes, kb_id)
|
||
except Exception as e:
|
||
logger.error(f"Milvus seed subgraph query failed: {e}")
|
||
return {"nodes": [], "edges": []}
|
||
|
||
async def query_and_rank_chunks_by_ppr(
|
||
self,
|
||
kb_id: str,
|
||
seed_weights: dict[str, float],
|
||
*,
|
||
max_nodes: int,
|
||
top_k: int,
|
||
damping: float,
|
||
) -> list[tuple[str, float]]:
|
||
if not seed_weights:
|
||
return []
|
||
subgraph = await self.query_seed_subgraph(
|
||
kb_id,
|
||
entity_ids=list(seed_weights.keys()),
|
||
max_nodes=max_nodes,
|
||
)
|
||
return self.rank_chunks_by_ppr(subgraph, seed_weights, top_k=top_k, damping=damping)
|
||
|
||
@staticmethod
|
||
def rank_chunks_by_ppr(
|
||
subgraph: dict[str, Any],
|
||
seed_weights: dict[str, float],
|
||
*,
|
||
top_k: int,
|
||
damping: float,
|
||
) -> list[tuple[str, float]]:
|
||
nodes = subgraph.get("nodes") or []
|
||
edges = subgraph.get("edges") or []
|
||
if not nodes:
|
||
return []
|
||
|
||
try:
|
||
import igraph as ig
|
||
except ImportError:
|
||
logger.error("Graph retrieval requires python-igraph. Please install igraph.")
|
||
return []
|
||
|
||
node_ids = [node["id"] for node in nodes]
|
||
index_by_id = {node_id: index for index, node_id in enumerate(node_ids)}
|
||
edge_indices = [
|
||
(index_by_id[edge["source_id"]], index_by_id[edge["target_id"]])
|
||
for edge in edges
|
||
if edge.get("source_id") in index_by_id and edge.get("target_id") in index_by_id
|
||
]
|
||
if not edge_indices:
|
||
return []
|
||
|
||
graph = ig.Graph(n=len(nodes), edges=edge_indices, directed=False)
|
||
reset = [0.0] * len(nodes)
|
||
chunk_node_indexes: list[tuple[int, str]] = []
|
||
for index, node in enumerate(nodes):
|
||
properties = node.get("properties") or {}
|
||
if node.get("type") == "Chunk" and properties.get("chunk_id"):
|
||
chunk_node_indexes.append((index, properties["chunk_id"]))
|
||
continue
|
||
entity_id = properties.get("entity_id")
|
||
if entity_id in seed_weights:
|
||
reset[index] = seed_weights[entity_id]
|
||
|
||
reset_total = sum(reset)
|
||
if reset_total <= 0 or not chunk_node_indexes:
|
||
return []
|
||
reset = [value / reset_total for value in reset]
|
||
scores = graph.personalized_pagerank(damping=min(max(damping, 0.1), 0.99), reset=reset)
|
||
ranked = sorted(
|
||
((chunk_id, float(scores[index])) for index, chunk_id in chunk_node_indexes),
|
||
key=lambda item: item[1],
|
||
reverse=True,
|
||
)
|
||
return ranked[:top_k]
|
||
|
||
async def get_labels(self, kb_id: str | None = None) -> list[str]:
|
||
effective_kb_id = kb_id or self.kb_id
|
||
if not effective_kb_id:
|
||
return []
|
||
label = safe_neo4j_label(effective_kb_id)
|
||
|
||
cypher = f"""
|
||
MATCH (n:MilvusKB:`{label}`)
|
||
UNWIND labels(n) AS node_label
|
||
WITH DISTINCT node_label
|
||
WHERE node_label <> 'MilvusKB' AND node_label <> $kb_id
|
||
RETURN node_label
|
||
ORDER BY node_label
|
||
"""
|
||
try:
|
||
records = neo4j_read(self.driver, cypher, kb_id=effective_kb_id)
|
||
return [record["node_label"] for record in records]
|
||
except Exception as e:
|
||
logger.error(f"Failed to get Milvus graph labels: {e}")
|
||
return []
|
||
|
||
async def get_stats(self, kb_id: str | None = None) -> dict[str, Any]:
|
||
effective_kb_id = kb_id or self.kb_id
|
||
if not effective_kb_id:
|
||
return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
|
||
label = safe_neo4j_label(effective_kb_id)
|
||
|
||
stats_cypher = f"""
|
||
MATCH (n:MilvusKB:`{label}`)
|
||
WITH count(n) AS node_count
|
||
OPTIONAL MATCH (:MilvusKB:`{label}`)-[r]->(:MilvusKB:`{label}`)
|
||
RETURN node_count, count(r) AS edge_count
|
||
"""
|
||
label_cypher = f"""
|
||
MATCH (n:Entity:MilvusKB:`{label}`)
|
||
WITH n.label AS entity_label, count(*) AS count
|
||
RETURN entity_label, count
|
||
ORDER BY count DESC
|
||
"""
|
||
try:
|
||
with self.driver.session() as session:
|
||
stats = session.run(stats_cypher).single()
|
||
label_stats = session.run(label_cypher)
|
||
return {
|
||
"total_nodes": stats["node_count"] if stats else 0,
|
||
"total_edges": stats["edge_count"] if stats else 0,
|
||
"entity_types": [{"type": row["entity_label"], "count": row["count"]} for row in label_stats],
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Failed to get Milvus graph stats: {e}")
|
||
return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
|
||
|
||
async def _get_milvus_kb(self, kb_id: str):
|
||
kb = await self.kb_repo.get_by_kb_id(kb_id)
|
||
if kb is None:
|
||
raise ValueError(f"知识库 {kb_id} 不存在")
|
||
if (kb.kb_type or "").lower() != "milvus":
|
||
raise ValueError("仅 Milvus 知识库支持独立图谱构建")
|
||
return kb
|
||
|
||
def _get_locked_config(self, additional_params: dict[str, Any]) -> dict[str, Any]:
|
||
config = additional_params.get(GRAPH_CONFIG_KEY) or {}
|
||
if not config.get("locked"):
|
||
raise ValueError("请先确认并锁定图谱抽取配置")
|
||
if not config.get("extractor_type"):
|
||
raise ValueError("图谱抽取配置缺少 extractor_type")
|
||
return config
|
||
|
||
def _public_config(self, config: dict[str, Any]) -> dict[str, Any] | None:
|
||
if not config:
|
||
return None
|
||
return {
|
||
"locked": bool(config.get("locked")),
|
||
"extractor_type": config.get("extractor_type"),
|
||
"extractor_options": self._runtime_extractor_options(config),
|
||
"created_at": config.get("created_at"),
|
||
"created_by": config.get("created_by"),
|
||
"updated_at": config.get("updated_at"),
|
||
"updated_by": config.get("updated_by"),
|
||
}
|
||
|
||
@staticmethod
|
||
def _build_where(exclude_chunk: bool, keyword: str) -> str:
|
||
clauses = []
|
||
if exclude_chunk:
|
||
clauses.append("NOT n:Chunk")
|
||
if keyword and keyword != "*":
|
||
clauses.append(
|
||
"(toLower(coalesce(n.name, '')) CONTAINS toLower($keyword)"
|
||
" OR toLower(coalesce(n.content_preview, '')) CONTAINS toLower($keyword)"
|
||
" OR toLower(coalesce(n.chunk_id, '')) CONTAINS toLower($keyword))"
|
||
)
|
||
return "WHERE " + " AND ".join(clauses) if clauses else ""
|
||
|
||
def _build_query(self, label: str, keyword: str, limit: int, max_depth: int, exclude_chunk: bool = False) -> str:
|
||
where = self._build_where(exclude_chunk, keyword)
|
||
m_exclude = " WHERE NOT m:Chunk" if exclude_chunk else ""
|
||
|
||
if max_depth <= 0:
|
||
return f"""
|
||
MATCH (n:MilvusKB:`{label}`)
|
||
{where}
|
||
RETURN n AS h, null AS r, null AS t
|
||
LIMIT $limit
|
||
"""
|
||
|
||
return f"""
|
||
MATCH (n:MilvusKB:`{label}`)
|
||
{where}
|
||
WITH n LIMIT $limit
|
||
OPTIONAL MATCH (n)-[r]-(m:MilvusKB:`{label}`){m_exclude}
|
||
RETURN n AS h, r AS r, m AS t
|
||
LIMIT {limit * 10}
|
||
"""
|
||
|
||
def _process_query_result(self, result, limit: int, kb_id: str, exclude_chunk: bool = False) -> dict[str, Any]:
|
||
nodes = []
|
||
edges = []
|
||
node_ids = set()
|
||
edge_ids = set()
|
||
|
||
for record in result:
|
||
for key in ("h", "t"):
|
||
raw_node = record.get(key)
|
||
if raw_node is None:
|
||
continue
|
||
node = self._normalize_node(raw_node, kb_id)
|
||
if not node or node["id"] in node_ids:
|
||
continue
|
||
if exclude_chunk and node.get("type") == "Chunk":
|
||
continue
|
||
nodes.append(node)
|
||
node_ids.add(node["id"])
|
||
raw_edge = record.get("r")
|
||
if raw_edge is not None:
|
||
edge = self._normalize_edge(raw_edge)
|
||
if edge and edge["id"] not in edge_ids:
|
||
edges.append(edge)
|
||
edge_ids.add(edge["id"])
|
||
if len(nodes) >= limit:
|
||
break
|
||
|
||
return {"nodes": nodes[:limit], "edges": edges[: limit * 2]}
|
||
|
||
def _process_subgraph_record(self, record: Any, limit: int, kb_id: str) -> dict[str, Any]:
|
||
nodes = []
|
||
edges = []
|
||
node_ids = set()
|
||
edge_ids = set()
|
||
|
||
for raw_node in record.get("nodes") or []:
|
||
node = self._normalize_node(raw_node, kb_id)
|
||
if not node or node["id"] in node_ids:
|
||
continue
|
||
nodes.append(node)
|
||
node_ids.add(node["id"])
|
||
if len(nodes) >= limit:
|
||
break
|
||
|
||
for raw_edge in record.get("edges") or []:
|
||
edge = self._normalize_edge(raw_edge)
|
||
if not edge or edge["id"] in edge_ids:
|
||
continue
|
||
if edge["source_id"] not in node_ids or edge["target_id"] not in node_ids:
|
||
continue
|
||
edges.append(edge)
|
||
edge_ids.add(edge["id"])
|
||
|
||
return {"nodes": nodes, "edges": edges}
|
||
|
||
def _normalize_node(self, raw_node: Any, kb_id: str | None = None) -> dict[str, Any]:
|
||
if hasattr(raw_node, "element_id"):
|
||
node_id = raw_node.element_id
|
||
labels = list(raw_node.labels)
|
||
properties = dict(raw_node.items())
|
||
elif isinstance(raw_node, dict):
|
||
node_id = raw_node.get("id") or raw_node.get("element_id")
|
||
labels = raw_node.get("labels", [])
|
||
properties = raw_node.get("properties") or {k: v for k, v in raw_node.items() if k not in {"id", "labels"}}
|
||
else:
|
||
return {}
|
||
|
||
effective_kb_id = kb_id or self.kb_id
|
||
db_label = properties.get("kb_id") or effective_kb_id
|
||
filtered_labels = [label for label in labels if label not in {"MilvusKB", db_label}]
|
||
entity_type = "Chunk" if "Chunk" in labels else properties.get("label", "Entity")
|
||
name = properties.get("name") or properties.get("content_preview") or properties.get("chunk_id") or "Unknown"
|
||
return {
|
||
"id": node_id,
|
||
"name": name,
|
||
"original_id": node_id,
|
||
"type": entity_type,
|
||
"labels": filtered_labels,
|
||
"properties": properties,
|
||
"normalized": {
|
||
"name": name,
|
||
"type": entity_type,
|
||
"source": "milvus",
|
||
},
|
||
"graph_type": "milvus",
|
||
}
|
||
|
||
def _normalize_edge(self, raw_edge: Any) -> dict[str, Any]:
|
||
if hasattr(raw_edge, "element_id"):
|
||
edge_id = raw_edge.element_id
|
||
edge_type = raw_edge.type
|
||
source_id = raw_edge.start_node.element_id
|
||
target_id = raw_edge.end_node.element_id
|
||
properties = dict(raw_edge.items())
|
||
edge_type = properties.get("type") or edge_type
|
||
elif isinstance(raw_edge, dict):
|
||
edge_id = raw_edge.get("id")
|
||
edge_type = raw_edge.get("type")
|
||
source_id = raw_edge.get("source_id")
|
||
target_id = raw_edge.get("target_id")
|
||
properties = raw_edge.get("properties", {})
|
||
else:
|
||
return {}
|
||
|
||
return {
|
||
"id": edge_id,
|
||
"source_id": source_id,
|
||
"target_id": target_id,
|
||
"type": edge_type,
|
||
"properties": properties,
|
||
"normalized": {
|
||
"type": edge_type,
|
||
"direction": "directed",
|
||
},
|
||
}
|