343 lines
11 KiB
Python
343 lines
11 KiB
Python
"""
|
||
知识图谱服务 - 实体识别、关系抽取、图谱查询
|
||
"""
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select, func
|
||
from app.models.knowledge_graph import KGNode, KGEdge
|
||
from app.models.document import Document, DocumentChunk
|
||
from app.services.llm_service import get_llm
|
||
from langchain_core.messages import HumanMessage
|
||
import json
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ENTITY_EXTRACTION_PROMPT = """从以下文本中提取实体和关系，返回 JSON 格式。
|
||
|
||
实体类型：
|
||
- person（人物）：人名、角色
|
||
- concept（概念）：抽象概念、理论、方法
|
||
- topic（主题）：话题、领域
|
||
- task（任务）：要做的事情
|
||
- event（事件）：发生的事件
|
||
- document（文档）：文件、资料
|
||
|
||
关系类型：
|
||
- related_to（相关于）
|
||
- part_of（隶属于）
|
||
- caused_by（由...导致）
|
||
- depends_on（取决于）
|
||
- contains（包含）
|
||
- located_in（位于）
|
||
- works_on（从事）
|
||
|
||
要求：
|
||
1. 识别文本中所有有意义的实体（不超过10个）
|
||
2. 识别实体之间的关系（每个实体至少一条关系）
|
||
3. 每个实体要有 name、type、description（1-2句话）
|
||
4. 关系要有 source、target、relation_type
|
||
|
||
文本内容：
|
||
{text}
|
||
|
||
请只返回 JSON，不要有其他内容：
|
||
{{
|
||
"entities": [
|
||
{{"name": "实体名", "type": "类型", "description": "描述"}}
|
||
],
|
||
"relations": [
|
||
{{"source": "实体A", "target": "实体B", "relation_type": "关系类型"}}
|
||
]
|
||
}}
|
||
"""
|
||
|
||
|
||
RELATION_INFERENCE_PROMPT = """根据以下实体列表和用户的问题，推断相关实体之间的关系。
|
||
|
||
用户问题：{question}
|
||
|
||
已知实体：
|
||
{entities}
|
||
|
||
请推断这些实体之间的隐含关系，返回 JSON：
|
||
{{
|
||
"inferred_relations": [
|
||
{{"source": "实体A", "target": "实体B", "relation_type": "关系类型", "confidence": 0.9}}
|
||
]
|
||
}}
|
||
|
||
关系类型：related_to / part_of / caused_by / depends_on / contains / works_on / located_in
|
||
confidence: 0.0-1.0，表示推断置信度
|
||
"""
|
||
|
||
|
||
class GraphService:
|
||
def __init__(self, db: AsyncSession):
|
||
self.db = db
|
||
self.llm = get_llm()
|
||
|
||
async def build_graph(self, user_id: str, document_ids: list[str] | None = None):
|
||
"""
|
||
从文档构建/更新知识图谱
|
||
- 遍历所有 chunk
|
||
- LLM 实体识别
|
||
- LLM 关系抽取
|
||
- 去重合并
|
||
"""
|
||
query = (
|
||
select(DocumentChunk)
|
||
.join(Document)
|
||
.where(Document.user_id == user_id)
|
||
.where(Document.is_indexed == True)
|
||
)
|
||
if document_ids:
|
||
query = query.where(DocumentChunk.document_id.in_(document_ids))
|
||
|
||
result = await self.db.execute(query)
|
||
chunks = list(result.scalars().all())
|
||
|
||
logger.info(f"[GraphService] 开始构建图谱，共 {len(chunks)} 个 chunks")
|
||
|
||
for chunk in chunks:
|
||
try:
|
||
await self._process_chunk(chunk, user_id)
|
||
except Exception as e:
|
||
logger.error(f"[GraphService] 处理 chunk {chunk.id} 失败: {e}")
|
||
continue
|
||
|
||
logger.info(f"[GraphService] 图谱构建完成")
|
||
|
||
async def _process_chunk(self, chunk: DocumentChunk, user_id: str):
|
||
"""处理单个 chunk，提取实体和关系"""
|
||
prompt = ENTITY_EXTRACTION_PROMPT.format(text=chunk.content[:2000])
|
||
response = await self.llm.invoke([HumanMessage(content=prompt)])
|
||
|
||
try:
|
||
data = json.loads(response.content)
|
||
except json.JSONDecodeError:
|
||
return
|
||
|
||
entities = data.get("entities", [])
|
||
relations = data.get("relations", [])
|
||
|
||
if not entities:
|
||
return
|
||
|
||
# 先查找已存在的节点
|
||
existing_nodes = {}
|
||
for entity_data in entities:
|
||
name = entity_data["name"]
|
||
result = await self.db.execute(
|
||
select(KGNode)
|
||
.where(KGNode.user_id == user_id)
|
||
.where(KGNode.name == name)
|
||
)
|
||
node = result.scalar_one_or_none()
|
||
if node:
|
||
existing_nodes[name] = node
|
||
|
||
# 插入新节点
|
||
entity_map = {}
|
||
for entity_data in entities:
|
||
name = entity_data["name"]
|
||
if name in existing_nodes:
|
||
entity_map[name] = existing_nodes[name].id
|
||
else:
|
||
node = KGNode(
|
||
user_id=user_id,
|
||
name=name,
|
||
entity_type=entity_data["type"],
|
||
description=entity_data.get("description", ""),
|
||
source_document_id=chunk.document_id,
|
||
)
|
||
self.db.add(node)
|
||
await self.db.flush()
|
||
entity_map[name] = node.id
|
||
|
||
# 插入关系（去重）
|
||
for rel in relations:
|
||
src, tgt = rel["source"], rel["target"]
|
||
if src not in entity_map or tgt not in entity_map:
|
||
continue
|
||
|
||
# 检查关系是否已存在
|
||
result = await self.db.execute(
|
||
select(KGEdge).where(
|
||
KGEdge.source_id == entity_map[src],
|
||
KGEdge.target_id == entity_map[tgt],
|
||
KGEdge.relation_type == rel["relation_type"],
|
||
)
|
||
)
|
||
existing = result.scalar_one_or_none()
|
||
if not existing:
|
||
edge = KGEdge(
|
||
source_id=entity_map[src],
|
||
target_id=entity_map[tgt],
|
||
relation_type=rel["relation_type"],
|
||
)
|
||
self.db.add(edge)
|
||
|
||
await self.db.commit()
|
||
|
||
async def get_graph_summary(self, user_id: str) -> str:
|
||
"""获取用户图谱的整体摘要"""
|
||
# 统计
|
||
node_count = await self.db.execute(
|
||
select(func.count()).select_from(KGNode).where(KGNode.user_id == user_id)
|
||
)
|
||
edge_count = await self.db.execute(
|
||
select(func.count()).select_from(KGEdge)
|
||
.select_from(KGEdge)
|
||
.join(KGNode, KGNode.id == KGEdge.source_id)
|
||
.where(KGNode.user_id == user_id)
|
||
)
|
||
|
||
node_total = node_count.scalar() or 0
|
||
edge_total = edge_count.scalar() or 0
|
||
|
||
if node_total == 0:
|
||
return "知识图谱为空，请先上传文档并构建图谱。"
|
||
|
||
# 按类型统计节点
|
||
type_result = await self.db.execute(
|
||
select(KGNode.entity_type, func.count())
|
||
.where(KGNode.user_id == user_id)
|
||
.group_by(KGNode.entity_type)
|
||
)
|
||
type_stats = type_result.all()
|
||
|
||
# 关系类型统计
|
||
rel_result = await self.db.execute(
|
||
select(KGEdge.relation_type, func.count())
|
||
.join(KGNode, KGNode.id == KGEdge.source_id)
|
||
.where(KGNode.user_id == user_id)
|
||
.group_by(KGEdge.relation_type)
|
||
)
|
||
rel_stats = rel_result.all()
|
||
|
||
# 列出最重要的节点（按 importance）
|
||
top_nodes_result = await self.db.execute(
|
||
select(KGNode)
|
||
.where(KGNode.user_id == user_id)
|
||
.order_by(KGNode.importance.desc())
|
||
.limit(10)
|
||
)
|
||
top_nodes = list(top_nodes_result.scalars().all())
|
||
|
||
lines = [
|
||
f"## 知识图谱摘要",
|
||
f"",
|
||
f"**总节点数**: {node_total}",
|
||
f"**总关系数**: {edge_total}",
|
||
f"",
|
||
f"### 节点类型分布",
|
||
]
|
||
for etype, count in type_stats:
|
||
lines.append(f"- {etype}: {count} 个")
|
||
|
||
lines.append(f"\n### 关系类型分布")
|
||
for rtype, count in rel_stats:
|
||
lines.append(f"- {rtype}: {count} 条")
|
||
|
||
lines.append(f"\n### 核心实体 (Top 10)")
|
||
for node in top_nodes:
|
||
lines.append(f"- [{node.entity_type}] {node.name}: {node.description[:50]}...")
|
||
|
||
return "\n".join(lines)
|
||
|
||
async def get_entity_context(self, entity: str, user_id: str) -> str:
|
||
"""获取某个实体的详细上下文"""
|
||
# 查找节点
|
||
result = await self.db.execute(
|
||
select(KGNode).where(
|
||
KGNode.user_id == user_id,
|
||
KGNode.name.contains(entity),
|
||
).limit(5)
|
||
)
|
||
nodes = list(result.scalars().all())
|
||
|
||
if not nodes:
|
||
return f"未找到实体: {entity}"
|
||
|
||
lines = []
|
||
for node in nodes:
|
||
lines.append(f"### {node.name} [{node.entity_type}]")
|
||
lines.append(f"描述: {node.description or '无描述'}")
|
||
|
||
# 获取该节点的关系
|
||
edges_result = await self.db.execute(
|
||
select(KGEdge, KGNode)
|
||
.join(KGNode, KGNode.id == KGEdge.target_id)
|
||
.where(KGEdge.source_id == node.id)
|
||
.limit(10)
|
||
)
|
||
out_edges = list(edges_result.all())
|
||
|
||
in_edges_result = await self.db.execute(
|
||
select(KGEdge, KGNode)
|
||
.join(KGNode, KGNode.id == KGEdge.source_id)
|
||
.where(KGEdge.target_id == node.id)
|
||
.limit(10)
|
||
)
|
||
in_edges = list(in_edges_result.all())
|
||
|
||
if out_edges:
|
||
lines.append("**关联到**:")
|
||
for edge, target in out_edges:
|
||
lines.append(f" - {node.name} --[{edge.relation_type}]--> {target.name}")
|
||
|
||
if in_edges:
|
||
lines.append("**被关联于**:")
|
||
for edge, source in in_edges:
|
||
lines.append(f" - {source.name} --[{edge.relation_type}]--> {node.name}")
|
||
|
||
lines.append("")
|
||
|
||
return "\n".join(lines)
|
||
|
||
async def get_neighbors(self, node_id: str, depth: int = 1) -> dict:
|
||
"""获取节点的邻居节点（用于图谱可视化）"""
|
||
visited = set()
|
||
current_level = {node_id}
|
||
all_nodes = []
|
||
all_edges = []
|
||
|
||
for _ in range(depth):
|
||
if not current_level:
|
||
break
|
||
next_level = set()
|
||
|
||
for nid in current_level:
|
||
if nid in visited:
|
||
continue
|
||
visited.add(nid)
|
||
|
||
# 获取节点
|
||
node_result = await self.db.execute(
|
||
select(KGNode).where(KGNode.id == nid)
|
||
)
|
||
node = node_result.scalar_one_or_none()
|
||
if node:
|
||
all_nodes.append(node)
|
||
|
||
# 获取出边
|
||
out_result = await self.db.execute(
|
||
select(KGEdge).where(KGEdge.source_id == nid)
|
||
)
|
||
for edge in out_result.scalars().all():
|
||
all_edges.append(edge)
|
||
next_level.add(edge.target_id)
|
||
|
||
# 获取入边
|
||
in_result = await self.db.execute(
|
||
select(KGEdge).where(KGEdge.target_id == nid)
|
||
)
|
||
for edge in in_result.scalars().all():
|
||
all_edges.append(edge)
|
||
next_level.add(edge.source_id)
|
||
|
||
current_level = next_level
|
||
|
||
return {"nodes": all_nodes, "edges": all_edges}
|