You've already forked DataMate
Jerry Yan
39338df808
核心功能：
- 三层检索策略：向量检索（Milvus）+ 图检索（KG 服务）+ 融合排序
- LLM 生成：支持同步和流式（SSE）响应
- 知识库访问控制：knowledge_base_id 归属校验 + collection_name 绑定验证
新增模块（9个文件）：
- models.py: 请求/响应模型（GraphRAGQueryRequest, RetrievalStrategy, GraphContext 等）
- milvus_client.py: Milvus 向量检索客户端（OpenAI Embeddings + asyncio.to_thread）
- kg_client.py: KG 服务 REST 客户端（全文检索 + 子图导出，fail-open）
- context_builder.py: 三元组文本化（10 种关系模板）+ 上下文构建
- generator.py: LLM 生成（ChatOpenAI，支持同步和流式）
- retriever.py: 检索编排（并行检索 + 融合排序）
- kb_access.py: 知识库访问校验（归属验证 + collection 绑定，fail-close）
- interface.py: FastAPI 端点（/query, /retrieve, /query/stream）
- __init__.py: 模块入口
修改文件（3个）：
- app/core/config.py: 添加 13 个 graphrag_* 配置项
- app/module/__init__.py: 注册 kg_graphrag_router
- pyproject.toml: 添加 pymilvus 依赖
测试覆盖（79 tests）：
- test_context_builder.py: 13 tests（三元组文本化 + 上下文构建）
- test_kg_client.py: 14 tests（KG 响应解析 + PagedResponse + 边字段映射）
- test_milvus_client.py: 8 tests（向量检索 + asyncio.to_thread）
- test_retriever.py: 11 tests（并行检索 + 融合排序 + fail-open）
- test_kb_access.py: 18 tests（归属校验 + collection 绑定 + 跨用户负例）
- test_interface.py: 15 tests（端点级回归 + 403 short-circuit）
关键设计：
- Fail-open: Milvus/KG 服务失败不阻塞管道，返回空结果
- Fail-close: 访问控制失败拒绝请求，防止授权绕过
- 并行检索: asyncio.gather() 并发运行向量和图检索
- 融合排序: Min-max 归一化 + 加权融合（vector_weight/graph_weight）
- 延迟初始化: 所有客户端在首次请求时初始化
- 配置回退: graphrag_llm_* 为空时回退到 kg_llm_*
安全修复：
- P1-1: KG 响应解析（PagedResponse.content）
- P1-2: 子图边字段映射（sourceEntityId/targetEntityId）
- P1-3: collection_name 越权风险（归属校验 + 绑定验证）
- P1-4: 同步 Milvus I/O（asyncio.to_thread）
- P1-5: 测试覆盖（79 tests，包括安全负例）
测试结果：79 tests pass ✅
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""三元组文本化 + 上下文构建。
|
|
|
|
将图谱子图（实体 + 关系）转为自然语言描述，
|
|
并与向量检索片段合并为 LLM 可消费的上下文文本。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from app.module.kg_graphrag.models import (
|
|
EntitySummary,
|
|
RelationSummary,
|
|
VectorChunk,
|
|
)
|
|
|
|
# 关系类型 -> 中文模板映射
|
|
RELATION_TEMPLATES: dict[str, str] = {
|
|
"HAS_FIELD": "{source}包含字段{target}",
|
|
"DERIVED_FROM": "{source}来源于{target}",
|
|
"USES_DATASET": "{source}使用了数据集{target}",
|
|
"PRODUCES": "{source}产出了{target}",
|
|
"ASSIGNED_TO": "{source}分配给了{target}",
|
|
"BELONGS_TO": "{source}属于{target}",
|
|
"TRIGGERS": "{source}触发了{target}",
|
|
"DEPENDS_ON": "{source}依赖于{target}",
|
|
"IMPACTS": "{source}影响了{target}",
|
|
"SOURCED_FROM": "{source}的知识来源于{target}",
|
|
}
|
|
|
|
# 通用模板（未在映射中的关系类型）
|
|
_DEFAULT_TEMPLATE = "{source}与{target}存在{relation}关系"
|
|
|
|
|
|
def textualize_subgraph(
|
|
entities: list[EntitySummary],
|
|
relations: list[RelationSummary],
|
|
) -> str:
|
|
"""将图谱子图转为自然语言描述。
|
|
|
|
Args:
|
|
entities: 子图中的实体列表。
|
|
relations: 子图中的关系列表。
|
|
|
|
Returns:
|
|
文本化后的图谱描述，每条关系/实体一行。
|
|
"""
|
|
lines: list[str] = []
|
|
|
|
# 记录有关系的实体名称
|
|
mentioned_entities: set[str] = set()
|
|
|
|
# 1. 对每条关系生成一句话
|
|
for rel in relations:
|
|
source_label = f"{rel.source_type}'{rel.source_name}'"
|
|
target_label = f"{rel.target_type}'{rel.target_name}'"
|
|
template = RELATION_TEMPLATES.get(rel.relation_type, _DEFAULT_TEMPLATE)
|
|
line = template.format(
|
|
source=source_label,
|
|
target=target_label,
|
|
relation=rel.relation_type,
|
|
)
|
|
lines.append(line)
|
|
mentioned_entities.add(rel.source_name)
|
|
mentioned_entities.add(rel.target_name)
|
|
|
|
# 2. 对独立实体（无关系）生成描述句
|
|
for entity in entities:
|
|
if entity.name not in mentioned_entities:
|
|
desc = entity.description or ""
|
|
if desc:
|
|
lines.append(f"{entity.type}'{entity.name}': {desc}")
|
|
else:
|
|
lines.append(f"存在{entity.type}'{entity.name}'")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def build_context(
|
|
vector_chunks: list[VectorChunk],
|
|
graph_text: str,
|
|
vector_weight: float = 0.6,
|
|
graph_weight: float = 0.4,
|
|
) -> str:
|
|
"""合并向量检索片段和图谱文本化内容为 LLM 上下文。
|
|
|
|
Args:
|
|
vector_chunks: 向量检索到的文档片段列表。
|
|
graph_text: 文本化后的图谱描述。
|
|
vector_weight: 向量分数权重（当前用于日志/调试，不影响上下文排序）。
|
|
graph_weight: 图谱相关性权重。
|
|
|
|
Returns:
|
|
合并后的上下文文本，分为「相关文档」和「知识图谱上下文」两个部分。
|
|
"""
|
|
sections: list[str] = []
|
|
|
|
# 向量检索片段
|
|
if vector_chunks:
|
|
doc_lines = ["## 相关文档"]
|
|
for i, chunk in enumerate(vector_chunks, 1):
|
|
doc_lines.append(f"[{i}] {chunk.text}")
|
|
sections.append("\n".join(doc_lines))
|
|
|
|
# 图谱文本化内容
|
|
if graph_text:
|
|
sections.append(f"## 知识图谱上下文\n{graph_text}")
|
|
|
|
if not sections:
|
|
return "（未检索到相关上下文信息）"
|
|
|
|
return "\n\n".join(sections)
|