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 ✅
51 lines
1.4 KiB
TOML
51 lines
1.4 KiB
TOML
[project]
|
|
name = "datamate-python"
|
|
version = "0.1.0"
|
|
description = "This is the Python backend of DataMate."
|
|
authors = [
|
|
{name = "Jason Wang",email = "jasonwong2019@outlook.com"}
|
|
]
|
|
license = {text = "MIT"}
|
|
readme = "README.md"
|
|
requires-python =">=3.12,<4.0.0"
|
|
dependencies = [
|
|
"uvicorn[standard] (>=0.38.0,<0.39.0)",
|
|
"aiomysql (>=0.3.2,<0.4.0)",
|
|
"pymysql (>=1.1.2,<2.0.0)",
|
|
"aiosqlite (>=0.21.0,<0.22.0)",
|
|
"httpx (>=0.28.1,<0.29.0)",
|
|
"pydantic-settings (>=2.12.0,<3.0.0)",
|
|
"python-multipart (>=0.0.20,<0.0.21)",
|
|
"python-dotenv (>=1.2.1,<2.0.0)",
|
|
"python-dateutil (>=2.9.0.post0,<3.0.0)",
|
|
"pyyaml (>=6.0.3,<7.0.0)",
|
|
"unstructured (>=0.18.21,<0.19.0)",
|
|
"markdown (>=3.10,<4.0)",
|
|
"langchain-community (>=0.4.1,<0.5.0)",
|
|
"jsonschema (>=4.25.1,<5.0.0)",
|
|
"greenlet (>=3.3.0,<4.0.0)",
|
|
"docx2txt (>=0.9,<0.10)",
|
|
"openpyxl (>=3.1.5,<4.0.0)",
|
|
"xlrd (>=2.0.1,<3.0.0)",
|
|
"jq (>=1.10.0,<2.0.0)",
|
|
"openai (>=2.9.0,<3.0.0)",
|
|
"langchain-openai (>=1.1.1,<2.0.0)",
|
|
"langchain (>=1.1.3,<2.0.0)",
|
|
"langchain-experimental (>=0.3.0,<1.0.0)",
|
|
"pydantic (>=2.12.5,<3.0.0)",
|
|
"sqlalchemy (>=2.0.45,<3.0.0)",
|
|
"fastapi (>=0.124.0,<0.125.0)",
|
|
"Pillow (>=11.0.0,<12.0.0)",
|
|
"pymilvus (>=2.5.0,<3.0.0)",
|
|
]
|
|
|
|
|
|
[build-system]
|
|
requires = ["poetry-core>=2.0.0,<3.0.0"]
|
|
build-backend = "poetry.core.masonry.api"
|
|
|
|
[tool.poetry]
|
|
packages = [
|
|
{ include = "app" }
|
|
]
|