Initial commit: QMDSearch 分层信息检索服务
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索(dense + sparse RRF 融合) - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
# ===========================================
|
||||
# QMDSearch 环境变量配置
|
||||
# ===========================================
|
||||
|
||||
# --- 应用 ---
|
||||
APP_PORT=8000
|
||||
LOG_LEVEL=info
|
||||
|
||||
# --- 嵌入模型 ---
|
||||
# openai | local
|
||||
EMBEDDING_PROVIDER=openai
|
||||
OPENAI_API_KEY=sk-xxx
|
||||
OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
EMBEDDING_MODEL=text-embedding-3-small
|
||||
EMBEDDING_DIMENSION=1536
|
||||
|
||||
# --- Qdrant ---
|
||||
QDRANT_PORT=6333
|
||||
QDRANT_DASHBOARD_PORT=6334
|
||||
|
||||
# --- Redis ---
|
||||
REDIS_PORT=6379
|
||||
|
||||
# --- Ollama 本地模型(文档三级总结)---
|
||||
OLLAMA_PORT=11434
|
||||
OLLAMA_MODEL=qwen2.5:1.5b
|
||||
# 备选模型: qwen2.5:3b(更好的总结质量,需更多内存)
|
||||
# EMBEDDING_PROVIDER=local 时使用的 Ollama 嵌入模型
|
||||
OLLAMA_EMBEDDING_MODEL=bge-m3
|
||||
|
||||
# --- NAS 持久化 ---
|
||||
NAS_DATA_DIR=./data
|
||||
|
||||
# --- 检索参数 ---
|
||||
# L2 语义检索召回数量
|
||||
RETRIEVAL_TOP_K=20
|
||||
# L3 重排后返回数量
|
||||
RETRIEVAL_FINAL_K=5
|
||||
|
||||
# --- 文档入库参数 ---
|
||||
# 低于此字符数触发 2.5 级回退
|
||||
SUMMARY_MIN_TEXT_LENGTH=500
|
||||
|
||||
# --- 入库异步任务 ---
|
||||
# 入库后台任务并发上限
|
||||
INGEST_MAX_CONCURRENCY=2
|
||||
# 任务状态 Redis 保留秒数(进行中与已完成,默认 24 小时)
|
||||
INGEST_TASK_TTL_DONE=86400
|
||||
# 失败任务状态 Redis 保留秒数(默认 7 天)
|
||||
INGEST_TASK_TTL_FAILED=604800
|
||||
|
||||
# --- 知识分类(taxonomy)---
|
||||
# taxonomy JSON 文件路径,留空使用内置默认类目集
|
||||
TAXONOMY_PATH=
|
||||
# 分类置信度低于此值归入 uncategorized
|
||||
CLASSIFY_CONFIDENCE_THRESHOLD=0.6
|
||||
# query 路由命中类目数上限,超过走全库兜底
|
||||
CLASSIFY_MAX_CATEGORIES=3
|
||||
|
||||
# --- 分层检索参数 ---
|
||||
# L1 层候选文档数
|
||||
L1_DOC_TOP_N=10
|
||||
# L2 层候选 section 数
|
||||
L2_SECTION_TOP_N=5
|
||||
# L3 层定位数
|
||||
L3_TOP_N=10
|
||||
# 是否启用稀疏检索
|
||||
SPARSE_ENABLED=true
|
||||
# Redis 缓存秒数
|
||||
CACHE_TTL=300
|
||||
|
||||
# --- 分块参数 ---
|
||||
# chunk 超长二次切分阈值(字符数)
|
||||
CHUNK_MAX_CHARS=800
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
.venv/
|
||||
|
||||
# Env
|
||||
.env
|
||||
.env.local
|
||||
|
||||
# Data
|
||||
data/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Docker
|
||||
docker-compose.override.yml
|
||||
|
||||
# Logs
|
||||
logs/
|
||||
*.log
|
||||
|
||||
# Test
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
@@ -0,0 +1,31 @@
|
||||
# Checklist
|
||||
|
||||
## 配置与任务管理器
|
||||
|
||||
- [x] config 含 ingest_max_concurrency(默认 2)/ingest_task_ttl_done(默认 86400)/ingest_task_ttl_failed(默认 604800),.env.example 同步
|
||||
- [x] submit 返回 task_id 且状态初始 pending;后台任务按信号量限流(超额保持 pending)
|
||||
- [x] 状态机按流水线阶段推进:pending→summarizing→classifying→embedding→writing→done
|
||||
- [x] 成功后 result 为完整 IngestionResult;失败 error 含 stage 与 message,partial_summary(已产出总结)保留
|
||||
- [x] 状态写 Redis:done TTL 24h、failed TTL 7 天;Redis 故障降级内存字典且接口不报错
|
||||
|
||||
## API
|
||||
|
||||
- [x] POST /api/v1/documents 合法请求返回 HTTP 202 + {task_id, status:"pending"},不再同步等待流水线
|
||||
- [x] POST 空文本等校验失败仍同步返回 code=1001(不进任务队列)
|
||||
- [x] GET /api/v1/documents/tasks/{task_id} 返回状态/时间戳;done 附 result;failed 附 error.stage/message
|
||||
- [x] 查询不存在 task_id 返回 code=1004
|
||||
- [x] 既有入库相关测试全部适配新契约且通过
|
||||
|
||||
## 管理页面
|
||||
|
||||
- [x] 入库提交后展示 task_id 与状态,2s 轮询直至 done/failed
|
||||
- [x] done 展示 category/置信度/tags/总结层级/chunks_count;failed 展示 stage 与 message
|
||||
- [x] 轮询期间禁止重复提交;5 分钟超时停止轮询并提示
|
||||
- [x] 页面测试覆盖 tasks 轮询路径与新交互标记
|
||||
|
||||
## 集成与文档
|
||||
|
||||
- [x] 内存闭环:提交→轮询 done→GET /documents/{doc_id} 可查→检索命中→删除 全通
|
||||
- [x] 失败路径:FakeOllama 抛错 → 任务 failed、stage 正确、failed TTL=7 天
|
||||
- [x] CLAUDE.md API 清单反映 POST /documents 202 异步与任务查询端点
|
||||
- [x] `uv run pytest` 全部通过;`uv run ruff check app tests scripts` 无错误
|
||||
@@ -0,0 +1,72 @@
|
||||
# 入库异步任务化 Spec
|
||||
|
||||
## Why
|
||||
|
||||
当前 `POST /api/v1/documents` 同步等待整条入库流水线(总结→分类→切分→向量化→写库),真实 Ollama 冒烟单篇耗时 39s+,长文档分钟级,客户端/网关易超时,批量导入不可用。改为后台任务模式:提交即返回 task_id,状态可查,失败任务长期保留便于归因。
|
||||
|
||||
## What Changes
|
||||
|
||||
- **BREAKING**:`POST /api/v1/documents` 由同步返回 IngestionResult(200)改为立即返回 `{task_id, status}`(HTTP 202),流水线在后台 asyncio 任务中执行
|
||||
- **新增** `GET /api/v1/documents/tasks/{task_id}`:查询任务状态(pending/各阶段/done/failed)、完成后的 IngestionResult、失败时的 stage 与错误信息
|
||||
- **新增** `app/core/ingest_tasks.py` IngestTaskManager:提交登记 → 信号量限流并发 → 后台执行 → 状态机推进 → 结果/错误落储
|
||||
- **状态存储**:Redis 为主(key `ingest_task:{task_id}`,JSON);done TTL 24h,**failed TTL 7 天**(便于失败归因);Redis 不可用时降级进程内字典(重启丢失,日志告警),服务不因此拒绝入库
|
||||
- **管理页**:入库区块改为「提交 → 轮询任务状态 → 展示进度/结果/失败详情」
|
||||
- **配置新增**:`ingest_max_concurrency`(默认 2)、`ingest_task_ttl_done`(默认 86400s)、`ingest_task_ttl_failed`(默认 604800s,7 天)
|
||||
- CLAUDE.md API 清单同步更新(标注 POST /documents 行为变更)
|
||||
|
||||
## Impact
|
||||
|
||||
- Affected specs: 文档入库 API、管理页面、任务状态存储
|
||||
- Affected code:
|
||||
- 修改:[document.py](file:///Users/kplam/coding/QMDSearch/app/api/v1/document.py)(POST 改 202 + 新增任务查询端点)、[config.py](file:///Users/kplam/coding/QMDSearch/app/config.py)、[.env.example](file:///Users/kplam/coding/QMDSearch/.env.example)、[admin.html](file:///Users/kplam/coding/QMDSearch/app/static/admin.html)(入库区块轮询)、[CLAUDE.md](file:///Users/kplam/coding/QMDSearch/CLAUDE.md)
|
||||
- 新增:`app/core/ingest_tasks.py`、`tests/test_ingest_tasks.py`、`tests/test_ingest_task_api.py`
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: 异步入库提交
|
||||
|
||||
系统 SHALL 将 `POST /api/v1/documents` 改为异步任务模式:校验通过后立即返回 HTTP 202 与 `{task_id, status: "pending"}`,入库流水线在后台任务执行;并发后台任务数受 `ingest_max_concurrency` 信号量限制,超限任务保持 pending 排队。请求体与校验规则(text 非空等)与现状一致,校验失败仍同步返回 1001。
|
||||
|
||||
#### Scenario: 提交即返回
|
||||
|
||||
- **WHEN** 提交合法文档
|
||||
- **THEN** 响应 202 + task_id,无需等待流水线完成;后台任务按并发额度开始执行
|
||||
|
||||
### Requirement: 任务状态查询
|
||||
|
||||
系统 SHALL 提供 `GET /api/v1/documents/tasks/{task_id}`:返回 `{task_id, status, created_at, updated_at}`;status 取值 `pending | summarizing | classifying | embedding | writing | done | failed`(与流水线阶段一致推进);done 时附 `result`(完整 IngestionResult:document_id/category/tags/总结层级/chunks_count);failed 时附 `error: {stage, message}`(stage 来自 IngestionError,已产出总结不丢失一并附在 error.partial_summary)。task_id 不存在返回 code=1004。
|
||||
|
||||
#### Scenario: 生命周期
|
||||
|
||||
- **WHEN** 提交后轮询 task_id
|
||||
- **THEN** 依次观察到 pending→阶段状态→done(result 可查,文档已可检索);失败任务观察到 failed + error.stage
|
||||
|
||||
### Requirement: 状态存储与保留
|
||||
|
||||
任务状态 SHALL 写入 Redis(`ingest_task:{task_id}`,JSON):done TTL=ingest_task_ttl_done(默认 24h),failed TTL=ingest_task_ttl_failed(默认 7 天),进行中状态 TTL 取 done 值。Redis 不可用时降级进程内字典并日志告警,入库与查询照常(重启后历史任务丢失可接受,接口语义不变)。
|
||||
|
||||
#### Scenario: Redis 宕机降级
|
||||
|
||||
- **WHEN** Redis 连接失败
|
||||
- **THEN** 提交与查询仍正常,仅日志 warning;进程重启后任务状态丢失
|
||||
|
||||
### Requirement: 管理页入库区块
|
||||
|
||||
管理页入库表单 SHALL 改为:提交后展示 task_id 与状态进度条/文本,每 2s 轮询任务接口,done 展示分类/标签/总结层级/chunks_count,failed 展示 error.stage 与 message;轮询期间禁止重复提交,超时(5 分钟)停止轮询并提示可稍后手动查询。
|
||||
|
||||
### Requirement: 并发控制
|
||||
|
||||
后台入库 SHALL 用 asyncio 信号量限制并发(默认 2),避免多任务同时打满 Ollama/embedding;排队任务状态保持 pending。
|
||||
|
||||
## MODIFIED Requirements
|
||||
|
||||
### Requirement: POST /api/v1/documents(行为变更)
|
||||
|
||||
原同步返回 IngestionResult 的行为废弃,改为 202 + task_id。所有既有调用方(管理页、既有测试)同步适配;响应体统一响应包装不变(`{code:0, data:{task_id,status}, message:"ok"}`,HTTP 状态码 202)。
|
||||
|
||||
## REMOVED Requirements
|
||||
|
||||
### Requirement: 同步入库等待
|
||||
|
||||
**Reason**:长耗时流水线导致客户端超时、无法批量导入。
|
||||
**Migration**:调用方改为提交后轮询 `GET /api/v1/documents/tasks/{task_id}`;管理页已内置轮询。
|
||||
@@ -0,0 +1,18 @@
|
||||
# Tasks
|
||||
|
||||
- [x] Task 1: 配置 + IngestTaskManager:config 新增 ingest_max_concurrency/ingest_task_ttl_done/ingest_task_ttl_failed 与 .env.example 同步;新建 app/core/ingest_tasks.py——submit(doc)→task_id、asyncio 信号量限流后台执行、状态机(pending→summarizing→classifying→embedding→writing→done/failed)、Redis 主存储(failed TTL 7 天)+ Redis 故障降级进程内字典、get(task_id) 查询;单测(FakeIngester 成功/失败/IngestionError.stage 透传/partial_summary 保留/并发限流/Redis 降级)
|
||||
- [x] SubTask 1.1: 配置项 + .env.example
|
||||
- [x] SubTask 1.2: IngestTaskManager 实现 + 单测
|
||||
- [x] Task 2: API 改造:POST /api/v1/documents 改 202 返回 {task_id,status}(校验失败仍同步 1001);新增 GET /api/v1/documents/tasks/{task_id}(不存在 1004);适配既有 test_document_api.py 与 test_e2e_integration.py 中入库断言为新契约;新增 tests/test_ingest_task_api.py
|
||||
- [x] Task 3: 管理页入库区块改造:提交→展示 task_id 与状态→2s 轮询→done 展示结果/failed 展示 stage+message→轮询期禁重复提交→5 分钟超时提示;更新 test_admin_page.py 断言(含 tasks 轮询路径)
|
||||
- [x] Task 4: 集成验证 + 文档:内存 Qdrant + FakeOllama 走通 提交→轮询至 done→文档可检索→删除 闭环;失败路径(FakeOllama 抛错)状态 failed 且 stage 正确;Redis 降级路径可用;CLAUDE.md API 清单更新(POST /documents 标注 202 异步、新增任务查询端点);`uv run pytest` 全绿 + ruff 通过
|
||||
|
||||
# Task Dependencies
|
||||
|
||||
- [Task 2] depends on [Task 1]
|
||||
- [Task 3] depends on [Task 2]
|
||||
- [Task 4] depends on [Task 3]
|
||||
|
||||
# Parallelizable
|
||||
|
||||
- 各任务串行依赖,无并行项
|
||||
@@ -0,0 +1,45 @@
|
||||
# Checklist
|
||||
|
||||
## 配置与模型
|
||||
|
||||
- [x] config 含 taxonomy 路径、分类置信度阈值、各层 top-k(l1/l2/l3)、sparse 开关、缓存 TTL,.env.example 同步更新
|
||||
- [x] taxonomy 默认配置文件存在且含 uncategorized 兜底类,加载校验函数可用
|
||||
- [x] CategoryResult(主类+多标签+置信度)、ChunkModel、SearchRequest/SearchResponse/SearchHit 模型齐备且有类型注解
|
||||
|
||||
## 入库链路
|
||||
|
||||
- [x] L2 大纲优先使用文档原生标题树,无结构文本回退 LLM 生成,2.5 级回退逻辑不受影响
|
||||
- [x] 分类器输出主类+多标签+置信度,低置信文档归入 uncategorized,不再硬编码 default
|
||||
- [x] chunk 按标题树切分并记录 section_path,超长 section 二次切分,短文本单 chunk
|
||||
- [x] 入库后 Qdrant 四层集合均有数据:doc_l1/doc_l2/doc_l3/chunks,payload 含 doc_id/category/tags/section_path
|
||||
- [x] chunks 与 doc_l1 同时携带 dense 与 sparse 向量,sparse 由本地分词+BM25 生成、无外部模型依赖
|
||||
- [x] 入库失败时返回明确错误码,已生成总结不丢失(可重试)
|
||||
|
||||
## 检索链路
|
||||
|
||||
- [x] query 解析以 JSON 约束输出意图类目+置信度+rewrite+关键词,解析失败自动降级全库检索
|
||||
- [x] 分类路由:高置信按主类硬过滤,多标签软召回生效,低置信/超类目上限走全库兜底
|
||||
- [x] 三级文档检索沿 L1→L2→L3→chunk 逐层收敛,chunk 候选仅来自 L3 命中范围
|
||||
- [x] 2.5 级文档 L1 命中后直进 L3/chunk 层
|
||||
- [x] L2/L3 空召回时回退上一层范围直搜 chunk,不返回空结果
|
||||
- [x] chunk 层 dense+sparse 双路召回经 RRF 融合,返回 final_k 个结果
|
||||
- [x] 检索结果 text 字段为原文 chunk,摘要仅以 doc_summary 上下文标注出现
|
||||
- [x] Redis 缓存命中时重复 query 不重复调用 Ollama/Qdrant;Redis 宕机检索仍可用
|
||||
|
||||
## API
|
||||
|
||||
- [x] POST /api/v1/documents 入库成功返回 document_id、分类结果、总结层级
|
||||
- [x] POST /api/v1/search 返回统一格式 {"code":0,"data":...,"message":"ok"}
|
||||
- [x] GET /api/v1/knowledge/categories 返回 taxonomy 类目列表
|
||||
- [x] 错误码符合 0=成功、1xxx=客户端错误、2xxx=服务端错误规范
|
||||
|
||||
## 评测
|
||||
|
||||
- [x] 回归集样例 ≥5 篇文档,每篇含应检出/不应检出 query 及 golden 标注
|
||||
- [x] 评测脚本输出 Entity Recall(L1/L3)、Hallucination Rate、Routing F1、Pruning Loss、Precision@5/Recall@10
|
||||
- [x] 报告含平铺 chunk baseline 对比,未达门槛项显式标出
|
||||
|
||||
## 端到端
|
||||
|
||||
- [x] docker compose 环境下 入库→检索→评测 全链路跑通(本机以内存 Qdrant 集成测试 + 真实 Ollama 冒烟等效验证,NAS docker 部署待用户侧执行)
|
||||
- [x] `uv run pytest` 全部通过
|
||||
@@ -0,0 +1,144 @@
|
||||
# 分层摘要索引 + 分类路由检索全链路 Spec
|
||||
|
||||
## Why
|
||||
|
||||
当前 QMDSearch 仅实现了三级总结器(summarizer),入库链路的分类/向量化/写入为 TODO,检索侧整体缺失,无法兑现 CLAUDE.md 中「面向 AI Agent 的分层信息检索服务」定位。本变更按「分层预摘要索引(Hierarchical Summarization Index)+ 分类路由 + 自顶向下剪枝检索」架构补齐全链路:离线侧把 Agentic RAG 中 router/planner 需要的结构信息预先算好,在线侧用「分类过滤 + 摘要树逐层剪枝 + hybrid 检索」实现大规模知识库下又省又准的检索,并附评测脚本骨架守护摘要质量与剪枝召回。
|
||||
|
||||
## What Changes
|
||||
|
||||
- **入库链路补全**:taxonomy 驱动的分类判定(主类 + 多标签 + 置信度,`uncategorized` 兜底)、dense + sparse(BM25)双向量生成、Qdrant 四层集合写入(L1 总结 / L2 大纲节点 / L3 内容大纲节点 / 原文 chunk)
|
||||
- **summarizer 改造**:L2 大纲优先解析文档原生标题树(Markdown 标题/编号标题),无结构文本回退 LLM 生成;新增 chunk 切分器,chunk 关联所属 section 路径
|
||||
- **分层检索引擎(新增)**:query 解析(Ollama 小模型 JSON 约束输出:意图分类 + query rewrite + 关键词)→ 分类路由(主类硬过滤 + 多标签软召回 + 低置信全库兜底)→ L1→L2→L3 摘要树逐层剪枝 → chunk 层 dense+sparse hybrid 召回(RRF 融合)→ 重排返回原文 chunk
|
||||
- **检索边界**:摘要仅用于路由与上下文标注,返回给调用方 Agent 的答案语料只含原文 chunk(防摘要幻觉进入生成)
|
||||
- **降级路径**:分类低置信跳过路由;任一层召回为空回退上一层范围直搜 chunk;2.5 级文档 L1 命中后直进 L3/chunk 层
|
||||
- **API(新增)**:`POST /api/v1/documents` 入库、`POST /api/v1/search` 分层检索、`GET /api/v1/knowledge/categories` 类目查询
|
||||
- **Redis 缓存**:query 解析结果与检索结果缓存(短 TTL)
|
||||
- **评测脚本骨架(新增)**:离线回归集 + Entity Recall / Hallucination Rate / Routing F1 / Pruning Loss / Precision@5 指标,支持与平铺 chunk baseline 对比
|
||||
|
||||
**存储设计说明**:不按类目建集合,采用全局四层集合 + category payload 索引过滤(Qdrant keyword payload index)。理由:支持多标签软召回与 uncategorized 兜底跨类检索,taxonomy 调整无需迁库;此设计取代 CLAUDE.md 中「按分类映射写入对应集合」的设想(该设想从未实现,非破坏性变更)。
|
||||
|
||||
**命名对齐**:用户方案中 L0/L1/L2 对应代码库既有命名 L1 总结 / L2 大纲 / L3 内容大纲(+ L2.5 回退),本 spec 沿用代码库命名。
|
||||
|
||||
## Impact
|
||||
|
||||
- Affected specs: 文档入库(ingestion)、文档三级总结(summarizer)、分层检索(新增)、评测(新增)
|
||||
- Affected code:
|
||||
- 修改:[summarizer.py](file:///Users/kplam/coding/QMDSearch/app/core/summarizer.py)、[ingestion.py](file:///Users/kplam/coding/QMDSearch/app/core/ingestion.py)、[document.py](file:///Users/kplam/coding/QMDSearch/app/models/document.py)、[config.py](file:///Users/kplam/coding/QMDSearch/app/config.py)、[main.py](file:///Users/kplam/coding/QMDSearch/app/main.py)
|
||||
- 新增:`app/core/embeddings.py`、`app/core/sparse.py`、`app/core/chunker.py`、`app/core/classifier.py`、`app/core/retriever.py`、`app/core/ranker.py`、`app/core/query_parser.py`、`app/services/qdrant.py`、`app/services/redis.py`、`app/models/search.py`、`app/models/knowledge.py`、`app/api/v1/document.py`、`app/api/v1/search.py`、`app/api/v1/knowledge.py`、`scripts/eval/`(评测脚本与回归集样例)
|
||||
- 基础设施:Qdrant 集合初始化(4 集合 + payload 索引 + sparse 配置),无需新增容器
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: taxonomy 分类判定
|
||||
|
||||
系统 SHALL 提供可配置的类目体系(taxonomy),入库时基于 L1 总结对文档分类,输出主类、多标签与置信度;置信度低于阈值时归入 `uncategorized` 兜底类。taxonomy 通过配置文件定义,支持环境变量指定路径,未配置时使用内置默认类目集。
|
||||
|
||||
#### Scenario: 正常分类
|
||||
|
||||
- **WHEN** 入库一篇内容明确的文档
|
||||
- **THEN** 返回 `main_category`(主类)、`tags`(0~3 个附加标签)、`confidence`(0~1),且主类属于 taxonomy 已定义类目
|
||||
|
||||
#### Scenario: 低置信兜底
|
||||
|
||||
- **WHEN** 文档内容跨类目或无法明确归类(分类置信度低于阈值)
|
||||
- **THEN** `main_category` 为 `uncategorized`,多标签保留候选类目,检索时该文档仅在全库兜底通道与标签软召回中可见
|
||||
|
||||
### Requirement: 双向量索引(dense + sparse)
|
||||
|
||||
系统 SHALL 为 L1 总结、L2 大纲节点、L3 内容大纲节点、原文 chunk 生成 dense 向量;并至少为原文 chunk 与 L1 总结生成 sparse(BM25)向量。dense 向量提供方可配置(openai / local),sparse 向量使用本地分词 + 特征哈希 BM25 实现,不依赖外部模型下载。
|
||||
|
||||
#### Scenario: 四层集合写入
|
||||
|
||||
- **WHEN** 文档完成总结与分类
|
||||
- **THEN** L1/L2/L3 各层节点与原文 chunk 分别写入对应集合,每条数据携带 `doc_id`、`category`、`tags`、`section_path`(chunk/L2/L3 层)payload,chunk 同时携带 dense 与 sparse 向量
|
||||
|
||||
### Requirement: chunk 切分与 section 关联
|
||||
|
||||
系统 SHALL 将原文按结构(标题树)优先、长度兜底切分为 chunk,每个 chunk 记录所属 section 路径,与 L2/L3 节点可互相定位。
|
||||
|
||||
#### Scenario: 结构化文档切分
|
||||
|
||||
- **WHEN** 文档含标题结构
|
||||
- **THEN** chunk 按 section 边界切分,`section_path` 与 L2 大纲节点对应;超长 section 内部按长度二次切分
|
||||
|
||||
### Requirement: query 解析与分类路由
|
||||
|
||||
系统 SHALL 在检索前用 Ollama 小模型以 JSON 约束输出解析 query:意图命中的类目集合 + 置信度、rewrite 后 query、关键词。解析失败或分类置信度低时 SHALL 跳过路由进入全库检索。
|
||||
|
||||
#### Scenario: 明确意图路由
|
||||
|
||||
- **WHEN** query 意图明确命中 1~2 个类目且置信度达标
|
||||
- **THEN** 后续检索仅在命中类目 payload 过滤范围内进行
|
||||
|
||||
#### Scenario: 跨类/模糊兜底
|
||||
|
||||
- **WHEN** 分类置信度低于阈值或命中类目数超过上限
|
||||
- **THEN** 放弃类目过滤,全库检索(不丢召回)
|
||||
|
||||
### Requirement: 自顶向下摘要树剪枝检索
|
||||
|
||||
系统 SHALL 按 L1→L2→L3→chunk 顺序逐层缩小检索范围:L1 层检索落候选文档(top-N),候选文档内检索 L2 落候选 section,候选 section 内检索 L3 定位 chunk 范围,最终在范围内 hybrid 检索 chunk。2.5 级文档无 L2 层,L1 命中后 SHALL 直进 L3/chunk 层。任一层召回为空时 SHALL 回退到上一层范围直接检索 chunk。
|
||||
|
||||
#### Scenario: 三级文档逐层命中
|
||||
|
||||
- **WHEN** query 在路由范围内有匹配文档
|
||||
- **THEN** 检索沿 摘要树逐层收敛,最终 chunk 候选仅来自 L3 命中的 section 范围
|
||||
|
||||
#### Scenario: 剪枝为空回退
|
||||
|
||||
- **WHEN** L2 或 L3 层在候选范围内召回为空
|
||||
- **THEN** 回退到上一层候选文档范围直接 hybrid 检索 chunk,不返回空结果
|
||||
|
||||
### Requirement: hybrid 召回与重排
|
||||
|
||||
chunk 层 SHALL 同时执行 dense 与 sparse 检索,使用 RRF 融合排序,按 `retrieval_top_k` 召回、`retrieval_final_k` 返回。
|
||||
|
||||
#### Scenario: RRF 融合
|
||||
|
||||
- **WHEN** chunk 层执行检索
|
||||
- **THEN** dense 与 sparse 两路结果经 RRF 融合后排序,返回 final_k 个原文 chunk
|
||||
|
||||
### Requirement: 检索结果只含原文
|
||||
|
||||
检索 API SHALL 仅返回原文 chunk 及引用元数据(doc_id、标题、section_path、score、所属文档 L1 总结作为上下文标注),摘要不作为可答题语料返回字段的主体。
|
||||
|
||||
#### Scenario: 结果结构
|
||||
|
||||
- **WHEN** 检索成功
|
||||
- **THEN** 每个 hit 含 `text`(原文)、`doc_id`、`title`、`section_path`、`score`、`doc_summary`(L1,仅上下文标注)
|
||||
|
||||
### Requirement: 检索缓存
|
||||
|
||||
系统 SHALL 使用 Redis 缓存 query 解析结果与检索结果(短 TTL,可配置),缓存键包含 query 文本与路由类目;缓存不可用时不影响主流程。
|
||||
|
||||
### Requirement: 入库与检索 API
|
||||
|
||||
系统 SHALL 提供 `POST /api/v1/documents`(入库,返回 document_id/分类/总结层级)、`POST /api/v1/search`(分层检索)、`GET /api/v1/knowledge/categories`(taxonomy 类目列表),统一响应格式 `{"code": 0, "data": {...}, "message": "ok"}`,错误码 0=成功、1xxx=客户端错误、2xxx=服务端错误。
|
||||
|
||||
### Requirement: 评测脚本骨架
|
||||
|
||||
系统 SHALL 提供离线评测脚本:回归集格式(文档 + golden query 含应检出/不应检出 + golden doc/section/chunk 标注),实现 Entity Recall、Hallucination Rate(LLM-as-judge 反查)、Routing F1、Pruning Loss、Precision@5 / Recall@10 指标,并支持同一 query 集对「平铺 chunk baseline」与本方案做 A/B 对比输出报告。
|
||||
|
||||
#### Scenario: 摘要质量门禁
|
||||
|
||||
- **WHEN** 对回归集运行评测脚本
|
||||
- **THEN** 输出各指标:L1 Entity Recall ≥ 0.85、L3 ≥ 0.9、Hallucination Rate < 2%、Pruning Loss < 8% 作为参考门槛,不达标项在报告中显式标红
|
||||
|
||||
## MODIFIED Requirements
|
||||
|
||||
### Requirement: L2 大纲生成(原 LLM 自由生成)
|
||||
|
||||
L2 大纲 SHALL 优先解析文档原生标题树(Markdown `#`/`##`/`###` 及常见编号标题模式),直接以标题层级作为大纲节点;仅当文档无可解析结构时回退为 LLM 生成大纲。L2 限定为「结构导航」,不写入结论性压缩语句,避免与 L1 总结语义重叠。原文无标题结构且为短文本时仍走既有 2.5 级回退。
|
||||
|
||||
**理由**:用户方案评审结论——LLM 自由生成的三级粒度易出现层级语义重叠、检索两头命中同一段;原生标题树更稳且零成本。
|
||||
|
||||
### Requirement: 分类判定(原 stub 返回 "default")
|
||||
|
||||
原 `Ingester._classify` 占位实现 SHALL 替换为独立 classifier 模块:基于 L1 总结调用 Ollama 小模型 JSON 约束输出主类 + 多标签 + 置信度,对齐 taxonomy;不再硬编码返回 `default`。
|
||||
|
||||
## REMOVED Requirements
|
||||
|
||||
### Requirement: 按分类映射写入对应集合(CLAUDE.md 设想)
|
||||
|
||||
**Reason**:固定 taxonomy 的按类目分集合会使新业务线/跨类目问题被路由砍死,且多标签软召回与 uncategorized 兜底无法跨集合实现;taxonomy 变更需迁库。
|
||||
**Migration**:从未实现,无需迁移;以全局四层集合 + category payload 索引过滤替代(见存储设计说明)。
|
||||
@@ -0,0 +1,54 @@
|
||||
# Tasks
|
||||
|
||||
- [x] Task 1: 扩展配置与数据模型:config 增加 taxonomy/各层 top-k/sparse 开关/缓存 TTL 等参数;新建 `app/models/knowledge.py`(TaxonomyCategory、CategoryResult)与 `app/models/search.py`(SearchRequest/SearchResponse/SearchHit);`app/models/document.py` 增加 ChunkModel(doc_id、text、section_path、chunk_index)
|
||||
- [x] SubTask 1.1: config.py 增加 taxonomy_path、classify_confidence_threshold、l1_doc_top_n、l2_section_top_n、l3_top_n、sparse_enabled、cache_ttl 等配置项,同步更新 .env.example
|
||||
- [x] SubTask 1.2: 新建 taxonomy 默认配置文件(内置通用类目集 + uncategorized),提供加载与校验函数
|
||||
- [x] SubTask 1.3: 新建 knowledge.py / search.py 模型,document.py 增加 ChunkModel 与 CategoryResult 引用
|
||||
- [x] Task 2: 向量服务:`app/core/embeddings.py`(dense,openai | local 双 provider 统一异步接口,批量编码)+ `app/core/sparse.py`(本地分词 + 特征哈希 BM25 稀疏向量,输出 Qdrant sparse vector 格式,不依赖外部模型)
|
||||
- [x] SubTask 2.1: EmbeddingService 抽象 + openai provider 实现 + local provider(Ollama embedding 接口)实现
|
||||
- [x] SubTask 2.2: SparseEncoder(中文分词 + BM25 权重 + 特征哈希到固定维度),单测验证输出稀疏格式合法
|
||||
- [x] Task 3: Qdrant 服务封装 `app/services/qdrant.py`:启动时初始化 4 个集合(doc_l1 / doc_l2 / doc_l3 / chunks,chunks 与 doc_l1 带 sparse 向量配置),建立 category/tags/doc_id/section_path payload 索引;提供 upsert 与按层检索(支持 payload 过滤 + dense/sparse/hybrid 查询)接口
|
||||
- [x] SubTask 3.1: 集合初始化与 payload 索引(幂等,应用启动时执行)
|
||||
- [x] SubTask 3.2: upsert 接口(按层写入,chunk 携带 dense+sparse)
|
||||
- [x] SubTask 3.3: 分层查询接口(dense 过滤查询 / hybrid 查询)
|
||||
- [x] Task 4: summarizer 改造 + chunk 切分器:L2 大纲优先解析原生标题树(Markdown 标题 + 编号标题正则),无结构回退现有 LLM 生成;新建 `app/core/chunker.py` 按标题树切分 chunk(超长 section 按长度二次切分,短文本整篇单 chunk),chunk 记录 section_path 与 L2 节点对应关系
|
||||
- [x] SubTask 4.1: 标题树解析器(Markdown ATX 标题 + 中文编号标题),输出 section 树
|
||||
- [x] SubTask 4.2: summarizer 接入标题树:有结构时 L2 直接用标题树生成大纲文本,无结构走原 prompt;L3 保持不变
|
||||
- [x] SubTask 4.3: Chunker 实现与单测(结构化/非结构化/短文本三种输入)
|
||||
- [x] Task 5: 分类器 + ingestion 全链路补全:`app/core/classifier.py`(基于 L1 总结,Ollama JSON 约束输出主类+多标签+置信度,低置信归 uncategorized);改造 `app/core/ingestion.py` 串起 总结→分类→chunk 切分→双向量化→Qdrant 写入,返回真实 document_id 与集合信息
|
||||
- [x] SubTask 5.1: Classifier 实现(taxonomy 注入 prompt,JSON 解析容错,置信度阈值判兜底)
|
||||
- [x] SubTask 5.2: Ingester 全链路串联(总结→分类→切 chunk→embedding→sparse→upsert 四层)
|
||||
- [x] SubTask 5.3: 入库失败处理:Qdrant 写入失败不丢已生成总结,返回明确错误码与可重试标识
|
||||
- [x] Task 6: 入库与类目 API:`app/api/v1/document.py`(POST /api/v1/documents)、`app/api/v1/knowledge.py`(GET /api/v1/knowledge/categories),统一响应包装与错误码,main.py 注册路由
|
||||
- [x] Task 7: query 解析与分类路由:`app/core/query_parser.py`(Ollama 小模型 JSON 约束输出:命中类目+置信度、rewrite query、关键词;解析失败/低置信/命中过多类目→全库兜底)
|
||||
- [x] SubTask 7.1: QueryParser 实现与 JSON 容错解析
|
||||
- [x] SubTask 7.2: 路由决策函数(主类硬过滤 + 多标签软召回合并为 Qdrant payload filter + 兜底判定),单测覆盖三类分支
|
||||
- [x] Task 8: 分层检索引擎 + 重排 + 检索 API:`app/core/retriever.py`(L1→L2→L3→chunk 逐层剪枝,2.5 级文档跳过 L2,空召回回退上一层直搜 chunk)+ `app/core/ranker.py`(RRF 融合 + final_k 截断)+ `app/api/v1/search.py`(POST /api/v1/search)
|
||||
- [x] SubTask 8.1: Retriever 逐层 drill-down 主流程
|
||||
- [x] SubTask 8.2: 降级路径(2.5 级跳层、空召回回退、路由兜底直通)
|
||||
- [x] SubTask 8.3: Ranker RRF 融合与单测
|
||||
- [x] SubTask 8.4: 检索 API 与统一响应
|
||||
- [x] Task 9: Redis 缓存 `app/services/redis.py`:query 解析结果与检索结果缓存(短 TTL 可配置),缓存键含 query 与路由类目;Redis 不可用时降级直连不报错
|
||||
- [x] Task 10: 评测脚本骨架 `scripts/eval/`:回归集 JSON 格式定义与样例(≥5 篇文档、每篇 3~5 条应检出 query + 1~2 条不应检出 query,标注 golden doc/section/chunk);指标实现 Entity Recall、Routing F1、Pruning Loss、Precision@5/Recall@10;Hallucination Rate 用 Ollama LLM-as-judge 反查;支持平铺 chunk baseline 对比并输出 Markdown 报告(含门槛标红:L1 Entity Recall≥0.85、L3≥0.9、幻觉率<2%、Pruning Loss<8%)
|
||||
- [x] SubTask 10.1: 回归集 schema + 样例数据
|
||||
- [x] SubTask 10.2: 摘要质量指标(Entity Recall / Hallucination Rate)
|
||||
- [x] SubTask 10.3: 检索效用指标(Routing F1 / Pruning Loss / Precision@5)+ baseline 对比报告
|
||||
- [x] Task 11: 端到端验证:docker compose 起 Qdrant/Redis/Ollama,走通 入库→检索→评测 全链路;`uv run pytest` 全绿;修复发现的问题
|
||||
|
||||
# Task Dependencies
|
||||
|
||||
- [Task 2] depends on [Task 1]
|
||||
- [Task 3] depends on [Task 1]
|
||||
- [Task 4] depends on [Task 1]
|
||||
- [Task 5] depends on [Task 2, Task 3, Task 4]
|
||||
- [Task 6] depends on [Task 5]
|
||||
- [Task 7] depends on [Task 1]
|
||||
- [Task 8] depends on [Task 3, Task 7]
|
||||
- [Task 9] depends on [Task 7, Task 8]
|
||||
- [Task 10] depends on [Task 5, Task 8]
|
||||
- [Task 11] depends on [Task 6, Task 9, Task 10]
|
||||
|
||||
# Parallelizable
|
||||
|
||||
- Task 2 / Task 3 / Task 4 / Task 7 互相独立,可并行
|
||||
- Task 6 与 Task 8 可并行(依赖均已满足后)
|
||||
@@ -0,0 +1,41 @@
|
||||
# Checklist
|
||||
|
||||
## 文档管理 API
|
||||
|
||||
- [x] GET /api/v1/documents 分页返回 items(doc_id/title/category/tags/summary)与 next_offset,空库返回空列表
|
||||
- [x] GET /api/v1/documents/{doc_id} 返回 L1 全部字段 + l2_nodes/l3_nodes + chunks_count
|
||||
- [x] 不存在 doc_id 详情返回 code=1004
|
||||
- [x] DELETE /api/v1/documents/{doc_id} 后四层集合该文档点全部清除,再查详情返回 1004
|
||||
- [x] DELETE 不存在 doc_id 幂等成功且 deleted 统计为 0
|
||||
|
||||
## 统计 API
|
||||
|
||||
- [x] GET /api/v1/knowledge/stats 返回四层集合点数、类目分布、uncategorized_count
|
||||
- [x] stats 数值与内存 Qdrant 实际写入一致(测试验证)
|
||||
|
||||
## 管理页面
|
||||
|
||||
- [x] GET /admin 返回 200 HTML
|
||||
- [x] 页面含 概览/文档管理/入库/检索测试台/类目 五个区块
|
||||
- [x] 页面无外部 CDN/外链资源(无 http(s):// src 或 link)
|
||||
- [x] 所有 fetch 指向 /api/v1/ 路径且与后端路由契约一致(路径、方法、字段名)
|
||||
- [x] 删除操作有二次确认逻辑(confirm 或等价交互)
|
||||
- [x] code≠0 时页面有错误提示处理
|
||||
|
||||
## 单测补全
|
||||
|
||||
- [x] judge.py:正常幻觉率计算、部分断言不支持、解析失败返回 0.0
|
||||
- [x] response.py:ok/error/ApiError 结构与 code
|
||||
- [x] ollama.py:generate 正常/json_mode payload 含 format:is_available 可达与异常分支
|
||||
- [x] run_eval.py:门槛判定(✅/❌)与数值格式化纯函数
|
||||
|
||||
## 部署与文档
|
||||
|
||||
- [x] Dockerfile 含 COPY scripts/ scripts/
|
||||
- [x] CLAUDE.md 项目结构、API 清单(含 documents 管理端点/stats//admin)与实际一致
|
||||
|
||||
## 集成
|
||||
|
||||
- [x] 内存 Qdrant 管理闭环冒烟通过:入库→列表→详情→检索→删除→统计
|
||||
- [x] `uv run pytest` 全部通过(含新增测试)
|
||||
- [x] `uv run ruff check app tests scripts` 无错误
|
||||
@@ -0,0 +1,83 @@
|
||||
# 项目完整性补全:管理 API + 管理页面 + 单测补全 Spec
|
||||
|
||||
## Why
|
||||
|
||||
分层 RAG 全链路已验收(149 测试全绿),但系统对外闭环不完整:已入库文档无法查看/管理(无列表/详情/删除 API),运维人员无 UI 可手工验证检索与入库,judge/response/ollama 客户端等模块无单测覆盖,Dockerfile 未打包评测脚本,CLAUDE.md 项目结构已过时。本变更补齐管理闭环与测试覆盖,使系统达到可运维状态。
|
||||
|
||||
## What Changes
|
||||
|
||||
- **文档管理 API(新增)**:`GET /api/v1/documents`(分页列表,源自 doc_l1 scroll)、`GET /api/v1/documents/{doc_id}`(详情:L1 总结 + L2/L3 节点 + chunk 数)、`DELETE /api/v1/documents/{doc_id}`(按 doc_id 过滤删除四层集合所有点,幂等)
|
||||
- **统计 API(新增)**:`GET /api/v1/knowledge/stats`(四层集合点数、类目分布、uncategorized 数量)
|
||||
- **管理页面(新增)**:FastAPI 托管静态单页 `/admin`,原生 JS + fetch 单文件实现(无 Node 构建链、镜像零新增依赖),含 概览 / 文档管理 / 文档入库 / 检索测试台 / 类目列表 五个区块;删除操作前端二次确认
|
||||
- **单元测试补全**:scripts/eval/judge.py、app/api/response.py、app/services/ollama.py(HTTP 行为)、scripts/eval/run_eval.py 可测纯函数、新增管理 API 端点测试、/admin 页面存在性测试
|
||||
- **完整性修补**:Dockerfile 增加 `COPY scripts/ scripts/`;CLAUDE.md 项目结构与 API 清单更新至当前实现
|
||||
|
||||
无破坏性变更(现有 API 路径与响应结构不变,仅新增)。
|
||||
|
||||
## Impact
|
||||
|
||||
- Affected specs: 文档管理(新增)、统计概览(新增)、管理页面(新增)、测试覆盖、部署
|
||||
- Affected code:
|
||||
- 修改:[qdrant.py](file:///Users/kplam/coding/QMDSearch/app/services/qdrant.py)(scroll/delete/count 管理操作)、[main.py](file:///Users/kplam/coding/QMDSearch/app/main.py)(挂载 /admin 静态页)、[document.py](file:///Users/kplam/coding/QMDSearch/app/api/v1/document.py)(新增端点)、[knowledge.py](file:///Users/kplam/coding/QMDSearch/app/api/v1/knowledge.py)(stats)、[Dockerfile](file:///Users/kplam/coding/QMDSearch/Dockerfile)、[CLAUDE.md](file:///Users/kplam/coding/QMDSearch/CLAUDE.md)
|
||||
- 新增:`app/static/admin.html`、`tests/test_judge.py`、`tests/test_response.py`、`tests/test_ollama_client.py`、`tests/test_run_eval.py`、`tests/test_document_admin_api.py`、`tests/test_admin_page.py`
|
||||
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: 文档列表与详情
|
||||
|
||||
系统 SHALL 提供 `GET /api/v1/documents`:基于 doc_l1 集合 scroll 分页返回文档摘要列表(doc_id、title、category、tags、总结层级推断、L1 总结摘要),支持 `limit`/`offset` 游标分页;并提供 `GET /api/v1/documents/{doc_id}` 返回单文档详情:L1 payload 全部字段、L2/L3 节点列表、chunk 数量。doc_id 不存在时返回 code=1004。
|
||||
|
||||
#### Scenario: 分页列表
|
||||
|
||||
- **WHEN** 请求 `GET /api/v1/documents?limit=20`
|
||||
- **THEN** 返回 code=0,data 含 items(每项 doc_id/title/category/tags/summary)与 next_offset(无更多为 null)
|
||||
|
||||
#### Scenario: 详情与不存在
|
||||
|
||||
- **WHEN** 请求已入库 doc_id 的详情
|
||||
- **THEN** 返回 L1 字段 + l2_nodes/l3_nodes 列表 + chunks_count;请求不存在的 doc_id 返回 code=1004
|
||||
|
||||
### Requirement: 文档删除
|
||||
|
||||
系统 SHALL 提供 `DELETE /api/v1/documents/{doc_id}`:按 doc_id payload 过滤删除 doc_l1/doc_l2/doc_l3/chunks 四个集合中的全部点;删除不存在 doc_id 幂等返回成功(deleted_points=0);返回各集合删除点数统计。
|
||||
|
||||
#### Scenario: 删除已入库文档
|
||||
|
||||
- **WHEN** 对已入库 doc_id 执行 DELETE
|
||||
- **THEN** 四层集合该 doc_id 的点全部被清除,再次 GET 详情返回 1004
|
||||
|
||||
### Requirement: 统计概览
|
||||
|
||||
系统 SHALL 提供 `GET /api/v1/knowledge/stats`:返回四层集合各自点数、按 category 的文档分布、uncategorized 文档数。统计基于 Qdrant count 与 doc_l1 轻量 scroll 聚合(仅取 category 字段),属管理端低频接口。
|
||||
|
||||
#### Scenario: 概览数据
|
||||
|
||||
- **WHEN** 请求 stats
|
||||
- **THEN** data 含 collections(四层点数)、categories(类目→文档数)、uncategorized_count
|
||||
|
||||
### Requirement: 管理页面
|
||||
|
||||
系统 SHALL 在 `/admin` 提供静态单页管理界面(单 HTML 文件,内联 CSS/JS,无外部 CDN 依赖),包含五个区块:概览(stats 展示)、文档管理(列表 + 详情查看 + 删除,删除需二次确认)、文档入库(表单提交 text/title/source,展示分类与总结结果)、检索测试台(输入 query 展示 hits/routed_categories/fallback)、类目列表。页面调用同-origin `/api/v1/*`,统一处理 code≠0 的错误提示。
|
||||
|
||||
#### Scenario: 页面可用性
|
||||
|
||||
- **WHEN** 浏览器访问 `/admin`
|
||||
- **THEN** 返回 200 HTML,包含五个功能区块与全部 fetch 调用指向 `/api/v1/` 路径,无外部资源引用
|
||||
|
||||
### Requirement: 单元测试补全
|
||||
|
||||
系统 SHALL 为以下模块补齐单测:judge.py(mock Ollama:断言抽取/支持判定/解析失败返回 0.0)、response.py(ok/error/ApiError 结构)、ollama.py(mock httpx:generate/json_mode payload/is_available 可达与异常分支)、run_eval.py 可测纯函数(格式化与门槛判定)、新增管理 API 与 /admin 页面存在性(TestClient 200 + 关键区块标记)。
|
||||
|
||||
### Requirement: 部署与文档完整性
|
||||
|
||||
Dockerfile SHALL 打包 scripts/ 目录(镜像内可执行评测);CLAUDE.md SHALL 更新项目结构、API 清单与管理页面说明至当前实现。
|
||||
|
||||
## MODIFIED Requirements
|
||||
|
||||
### Requirement: QdrantService(新增管理操作)
|
||||
|
||||
QdrantService SHALL 新增:`scroll_l1(limit, offset) -> (items, next_offset)`(仅取列表所需 payload 字段)、`get_doc_detail(doc_id) -> dict | None`(L1 + L2/L3 节点 + chunk 计数)、`delete_by_doc_id(doc_id) -> dict[str, int]`(四集合按过滤删除,返回各集合删除数)、`count(collection) -> int`。均复用现有客户端与集合常量,不改变既有 upsert/查询行为。
|
||||
|
||||
## REMOVED Requirements
|
||||
|
||||
无。
|
||||
@@ -0,0 +1,24 @@
|
||||
# Tasks
|
||||
|
||||
- [x] Task 1: QdrantService 管理操作扩展:scroll_l1 分页(仅取 doc_id/title/category/tags/text/level 所需字段)、get_doc_detail(L1+L2/L3 节点+chunk 计数)、delete_by_doc_id(四集合过滤删除返回各集合删除数)、count;内存 Qdrant 单测覆盖
|
||||
- [x] SubTask 1.1: scroll_l1 与 count 实现 + 单测(分页游标、空集合)
|
||||
- [x] SubTask 1.2: get_doc_detail 与 delete_by_doc_id 实现 + 单测(存在/不存在/删除后四层清空)
|
||||
- [x] Task 2: 文档管理 API:`GET /api/v1/documents`(limit/offset 分页)、`GET /api/v1/documents/{doc_id}`(详情,不存在 code=1004)、`DELETE /api/v1/documents/{doc_id}`(幂等,返回各集合删除数);沿用统一响应与懒加载单例模式;TestClient 测试(mock QdrantService)
|
||||
- [x] Task 3: 统计 API:`GET /api/v1/knowledge/stats`(四层点数 + 类目分布 + uncategorized 数);TestClient 测试
|
||||
- [x] Task 4: 管理页面 `app/static/admin.html` 单文件(内联 CSS/JS,零外部依赖)+ main.py 挂载 /admin:概览/文档管理(列表/详情/删除二次确认)/入库表单/检索测试台/类目列表五区块,统一 code≠0 错误提示;TestClient 存在性测试(200 + 五区块标记 + 无 http(s) 外链资源)
|
||||
- [x] Task 5: 单测补全:tests/test_judge.py(mock Ollama:正常判定/部分断言不支持/解析失败返回 0.0)、tests/test_response.py、tests/test_ollama_client.py(mock httpx:generate/json_mode/is_available 分支)、tests/test_run_eval.py(门槛判定与格式化纯函数)
|
||||
- [x] Task 6: 部署与文档修补:Dockerfile 增加 `COPY scripts/ scripts/`;CLAUDE.md 更新项目结构(scripts/eval、static)、API 清单(documents 管理端点、stats、/admin)与管理页面使用说明
|
||||
- [x] Task 7: 集成验证:`uv run pytest` 全绿 + `uv run ruff check app tests scripts` 通过;内存 Qdrant 走通 入库→列表→详情→检索→删除→统计 管理闭环冒烟;确认页面 fetch 路径与实际 API 契约一致
|
||||
- [x] Task 8: 修复验收发现:CLAUDE.md 项目结构中 tests 目录注释「18 个测试文件」与实际 26 个不符,改为不写死数量防止再次漂移
|
||||
|
||||
# Task Dependencies
|
||||
|
||||
- [Task 2] depends on [Task 1]
|
||||
- [Task 3] depends on [Task 1]
|
||||
- [Task 4] depends on [Task 2, Task 3]
|
||||
- [Task 7] depends on [Task 4, Task 5, Task 6]
|
||||
|
||||
# Parallelizable
|
||||
|
||||
- Task 1 / Task 5 / Task 6 互相独立,可并行
|
||||
- Task 2 与 Task 3 可并行(Task 1 完成后)
|
||||
@@ -0,0 +1,160 @@
|
||||
# QMDSearch - AI Agent 分层信息检索服务
|
||||
|
||||
## 项目概述
|
||||
|
||||
QMDSearch 是面向 AI Agent 的分层信息检索服务,支持多层级知识库检索、向量语义搜索和结构化数据查询。基于 Docker 部署在 NAS 上,为 AI Agent 提供高效、精准的信息检索能力。
|
||||
|
||||
## 技术栈
|
||||
|
||||
- **语言**: Python 3.12+
|
||||
- **Web 框架**: FastAPI
|
||||
- **向量数据库**: Qdrant (Docker)
|
||||
- **缓存**: Redis (Docker)
|
||||
- **关系数据库**: PostgreSQL (Docker, 可选)
|
||||
- **嵌入模型**: OpenAI / 本地模型 (通过配置切换)
|
||||
- **本地推理模型**: Ollama (Docker) — 用于文档三级总结,默认 qwen2.5:1.5b
|
||||
- **部署**: Docker Compose on NAS
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
QMDSearch/
|
||||
├── app/ # 应用主目录
|
||||
│ ├── main.py # FastAPI 入口(含 /admin 管理页面挂载)
|
||||
│ ├── config.py # 配置管理
|
||||
│ ├── api/ # API 路由层
|
||||
│ │ ├── response.py # 统一响应格式 (ok/error/ApiError)
|
||||
│ │ └── v1/ # API v1 版本
|
||||
│ │ ├── search.py # 检索接口
|
||||
│ │ ├── document.py # 文档入库/管理接口
|
||||
│ │ └── knowledge.py # 知识库接口
|
||||
│ ├── core/ # 核心业务逻辑
|
||||
│ │ ├── retriever.py # 分层检索引擎 (L1→L2→L3→chunk)
|
||||
│ │ ├── query_parser.py # query 解析与分类路由
|
||||
│ │ ├── embeddings.py # 向量嵌入
|
||||
│ │ ├── sparse.py # 稀疏向量编码 (BM25 近似)
|
||||
│ │ ├── ranker.py # RRF 融合与结果截断
|
||||
│ │ ├── summarizer.py # 文档三级总结 (Ollama)
|
||||
│ │ ├── headings.py # 原生标题树解析
|
||||
│ │ ├── classifier.py # 文档分类 (主类+标签+置信度)
|
||||
│ │ ├── chunker.py # 标题树感知 chunk 切分
|
||||
│ │ └── ingestion.py # 文档入库 (总结→分类→写入)
|
||||
│ ├── models/ # 数据模型
|
||||
│ │ ├── search.py # 检索请求/响应模型
|
||||
│ │ ├── knowledge.py # 知识库/taxonomy 模型
|
||||
│ │ └── document.py # 文档/总结模型
|
||||
│ ├── services/ # 服务层 (外部交互)
|
||||
│ │ ├── qdrant.py # Qdrant 客户端 (四层集合)
|
||||
│ │ ├── redis.py # Redis 缓存客户端
|
||||
│ │ └── ollama.py # Ollama 客户端
|
||||
│ ├── static/ # 静态资源
|
||||
│ │ └── admin.html # 管理页面 (单文件)
|
||||
│ └── utils/ # 工具函数
|
||||
├── scripts/ # 运维与评测脚本
|
||||
│ ├── eval/ # 回归评测 (run_eval.py / metrics.py / judge.py / regression_set.json)
|
||||
│ └── smoke_live.py # 在线冒烟脚本
|
||||
├── tests/ # 测试
|
||||
├── docker-compose.yml # Docker 编排
|
||||
├── Dockerfile # 应用镜像
|
||||
├── .env.example # 环境变量模板
|
||||
├── pyproject.toml # 项目配置
|
||||
└── CLAUDE.md # 本文件
|
||||
```
|
||||
|
||||
## 分层检索架构
|
||||
|
||||
检索链路:query 解析路由 → L1→L2→L3 摘要树逐层剪枝 → chunk 层 hybrid 检索 (dense + sparse,RRF 融合)。
|
||||
|
||||
0. **query 解析路由**: 用 Ollama 小模型将 query 解析为结构化结果(rewrite、关键词、命中类目+置信度);高置信且类目数不超上限时按类目过滤,低置信/解析失败/类目过多则全库兜底(不丢召回)
|
||||
1. **L1 - 文档定位层**: 在文档 L1 总结集合中检索,产出候选文档集合;无命中时直接全库 chunk 兜底(fallback=True)
|
||||
2. **L2 - 章节大纲层**: 在候选文档内检索 L2 章节大纲,按 section 剪枝定位;未命中的文档(含 2.5 级文档)落入 L3 b 路(仅按 doc 过滤)
|
||||
3. **L3 - 小节内容层**: 两路查询(a:L2 命中文档按 section_path 过滤;b:其余文档仅按 doc_id 过滤)后 RRF 融合;无命中时 chunk 层回退为 L1 候选文档级检索
|
||||
4. **chunk 层**: 在 L3 收窄的范围内对原文 chunk 做 hybrid 检索(dense 向量 + sparse 稀疏向量 RRF 融合),截断后返回最终 Top-K
|
||||
|
||||
## API 设计规范
|
||||
|
||||
- RESTful 风格,版本化路径: `/api/v1/`
|
||||
- 请求/响应使用 Pydantic 模型校验
|
||||
- 统一响应格式: `{"code": 0, "data": {...}, "message": "ok"}`
|
||||
- 错误码: 0=成功, 1xxx=客户端错误, 2xxx=服务端错误
|
||||
|
||||
## API 清单
|
||||
|
||||
| 方法 | 路径 | 说明 |
|
||||
|------|------|------|
|
||||
| GET | `/api/v1/health` | 健康检查 |
|
||||
| POST | `/api/v1/search` | 分层检索 |
|
||||
| POST | `/api/v1/documents` | 文档入库(202 异步入库,返回 task_id) |
|
||||
| GET | `/api/v1/documents/tasks/{task_id}` | 入库任务状态查询(done 附 result,failed 附 error) |
|
||||
| GET | `/api/v1/knowledge/categories` | 知识分类类目集 |
|
||||
| GET | `/api/v1/knowledge/stats` | 统计(四层点数 + 类目分布 + uncategorized 数) |
|
||||
| GET | `/api/v1/documents` | 文档列表(limit/offset 分页) |
|
||||
| GET | `/api/v1/documents/{doc_id}` | 文档详情 |
|
||||
| DELETE | `/api/v1/documents/{doc_id}` | 删除文档(幂等) |
|
||||
| GET | `/admin` | 管理页面 |
|
||||
|
||||
入库任务状态持久化在 Redis(key: `ingest_task:{task_id}`):进行中与 done 保留 24h,failed 保留 7 天;Redis 不可用时降级为纯内存。
|
||||
|
||||
## 管理页面
|
||||
|
||||
浏览器访问 `/admin`,单页面含概览、文档管理(列表/详情/删除)、文档入库、检索测试台、类目列表五个区块。
|
||||
|
||||
## 编码规范
|
||||
|
||||
- 使用 uv 管理依赖
|
||||
- 类型注解必须 (Python 3.12+ 语法)
|
||||
- 异步优先 (async/await)
|
||||
- 配置通过环境变量注入,使用 pydantic-settings
|
||||
- 日志使用 structlog 结构化日志
|
||||
|
||||
## 文档入库与三级总结
|
||||
|
||||
文档入库时通过 Ollama 本地小模型对文档内容进行三级总结,然后根据总结进行分类入库。
|
||||
|
||||
### 总结层级
|
||||
|
||||
| 层级 | 名称 | 说明 | 存储用途 |
|
||||
|------|------|------|----------|
|
||||
| L1 | 总结 | 对整篇文档的一句话高度概括 | 快速分类、知识库路由 |
|
||||
| L2 | 大纲 | 提取文档主要章节和关键主题(优先使用文档原生标题树,无结构文本回退 LLM 生成) | 检索召回、上下文概览 |
|
||||
| L3 | 内容大纲 | 每个章节的详细内容摘要 | 精确匹配、深度检索 |
|
||||
|
||||
**2.5 级回退**: 当文档内容不足以支撑三级总结时(如短文本、简单说明),自动降级为二级总结:
|
||||
- L1: 总结(同上)
|
||||
- L2.5: 内容大纲(跳过大纲层,直接输出详细摘要)
|
||||
|
||||
### 入库流程
|
||||
|
||||
```
|
||||
文档输入 → 文本提取 → 三级总结(Ollama) → 分类判定(L1总结) → 向量化 → 写入Qdrant
|
||||
↓
|
||||
不足三级 → 2.5级回退
|
||||
```
|
||||
|
||||
1. **文本提取**: 从文件中提取纯文本内容
|
||||
2. **三级总结**: 调用 Ollama 依次生成 L1/L2/L3 总结;其中 L2 大纲优先使用文档原生标题树(不调 LLM),无结构文本(标题数 < 2)回退 LLM 生成
|
||||
3. **分类判定**: 根据 L1 总结将文档分配到 taxonomy 类目,输出主类目 + 附加标签 + 置信度;LLM 输出解析失败或置信度低于阈值时归 uncategorized 兜底(低置信时候选类目名保留进 tags 供软召回)
|
||||
4. **向量化**: 对原文 chunk 和各级总结分别生成 embedding(dense + sparse)
|
||||
5. **写入 Qdrant**: 将文档元数据、各级总结、向量写入对应四层集合(L1/L2/L3/chunks)
|
||||
|
||||
### 本地模型 (Ollama)
|
||||
|
||||
- 服务: Ollama 容器,默认模型 `qwen2.5:1.5b`(约 1GB,适合 NAS 低资源环境)
|
||||
- 备选模型: `qwen2.5:3b`(更好的总结质量,约 2GB)
|
||||
- 模型通过环境变量 `OLLAMA_MODEL` 配置
|
||||
- 首次启动时自动拉取模型,需 NAS 可访问外网
|
||||
|
||||
## Docker 部署说明
|
||||
|
||||
- 目标环境: NAS (ARM64/AMD64)
|
||||
- 镜像: python:3.12-slim
|
||||
- 持久化: NAS 本地目录挂载
|
||||
- 端口: 默认 8000 (API), 6333 (Qdrant), 6379 (Redis), 11434 (Ollama)
|
||||
- 健康检查: `/api/v1/health`
|
||||
|
||||
## 开发流程
|
||||
|
||||
1. 修改代码后本地测试: `uv run pytest`
|
||||
2. 构建镜像: `docker compose build`
|
||||
3. 启动服务: `docker compose up -d`
|
||||
4. 查看日志: `docker compose logs -f app`
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
FROM python:3.12-slim AS base
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# 安装 uv
|
||||
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uv/bin/uv
|
||||
|
||||
# 依赖层
|
||||
COPY pyproject.toml uv.lock ./
|
||||
RUN uv sync --frozen --no-dev
|
||||
|
||||
# 应用层
|
||||
COPY app/ app/
|
||||
COPY scripts/ scripts/
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=5s --retries=3 \
|
||||
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/api/v1/health')"
|
||||
|
||||
CMD ["uv", "run", "uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -0,0 +1,26 @@
|
||||
"""统一 API 响应包装与业务异常
|
||||
|
||||
响应格式:{"code": 0, "data": ..., "message": "ok"}
|
||||
错误码约定:0 成功;1001 请求参数校验失败;2000 服务器内部错误。
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def ok(data: Any) -> dict[str, Any]:
|
||||
"""成功响应包装"""
|
||||
return {"code": 0, "data": data, "message": "ok"}
|
||||
|
||||
|
||||
def error(code: int, message: str) -> dict[str, Any]:
|
||||
"""错误响应包装"""
|
||||
return {"code": code, "data": None, "message": message}
|
||||
|
||||
|
||||
class ApiError(Exception):
|
||||
"""业务异常:携带错误码与消息,由全局异常处理器转为统一错误响应"""
|
||||
|
||||
def __init__(self, code: int, message: str) -> None:
|
||||
self.code = code
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
@@ -0,0 +1,111 @@
|
||||
"""文档 API:POST /api/v1/documents 异步入库 + 任务查询 + 文档管理(列表/详情/删除)"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
from fastapi import APIRouter, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from app.api.response import ApiError, ok
|
||||
from app.config import Settings
|
||||
from app.core.ingest_tasks import IngestTaskManager
|
||||
from app.core.ingestion import Ingester
|
||||
from app.models.document import DocumentInput
|
||||
from app.services.qdrant import QdrantService
|
||||
from app.services.redis import RedisCache, get_cache
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
router = APIRouter(prefix="/api/v1", tags=["document"])
|
||||
|
||||
# 模块级懒加载单例,避免每请求重建 Ingester 及其下游依赖
|
||||
_ingester: Ingester | None = None
|
||||
_qdrant: QdrantService | None = None
|
||||
_task_manager: IngestTaskManager | None = None
|
||||
|
||||
|
||||
def _get_ingester() -> Ingester:
|
||||
global _ingester
|
||||
if _ingester is None:
|
||||
_ingester = Ingester()
|
||||
return _ingester
|
||||
|
||||
|
||||
def _get_qdrant() -> QdrantService:
|
||||
global _qdrant
|
||||
if _qdrant is None:
|
||||
_qdrant = QdrantService()
|
||||
return _qdrant
|
||||
|
||||
|
||||
def _get_task_manager() -> IngestTaskManager:
|
||||
"""入库任务管理器懒加载单例
|
||||
|
||||
Redis 沿用全局缓存单例(RedisCache 读写全容错,不可用时镜像写失败仅告警);
|
||||
获取缓存实例异常时传 None,退化为纯内存模式。
|
||||
"""
|
||||
global _task_manager
|
||||
if _task_manager is None:
|
||||
try:
|
||||
redis: RedisCache | None = get_cache()
|
||||
except Exception:
|
||||
logger.warning("Redis 缓存不可用,入库任务状态仅保留在内存", exc_info=True)
|
||||
redis = None
|
||||
_task_manager = IngestTaskManager(ingester=_get_ingester(), redis=redis, settings=Settings())
|
||||
return _task_manager
|
||||
|
||||
|
||||
@router.post("/documents")
|
||||
async def ingest_document(doc: DocumentInput) -> JSONResponse:
|
||||
"""文档入库入口:登记异步任务并返回 202 + task_id,入库结果经任务查询端点获取"""
|
||||
if not doc.text.strip():
|
||||
raise ApiError(1001, "文档内容不能为空")
|
||||
task_id = await _get_task_manager().submit(doc)
|
||||
return JSONResponse(status_code=202, content=ok({"task_id": task_id, "status": "pending"}))
|
||||
|
||||
|
||||
@router.get("/documents/tasks/{task_id}")
|
||||
async def get_ingest_task(task_id: str) -> dict[str, Any]:
|
||||
"""查询入库任务状态:含 task_id/status/created_at/updated_at,done 附 result,failed 附 error"""
|
||||
task = await _get_task_manager().get(task_id)
|
||||
if task is None:
|
||||
raise ApiError(1004, "任务不存在")
|
||||
return ok(task)
|
||||
|
||||
|
||||
@router.get("/documents")
|
||||
async def list_documents(
|
||||
limit: int = Query(default=20, ge=1, le=100),
|
||||
offset: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""分页列出文档(L1 摘要),返回 items 与下一页游标 next_offset"""
|
||||
try:
|
||||
items, next_offset = await _get_qdrant().scroll_l1(limit=limit, offset=offset)
|
||||
except Exception as exc:
|
||||
logger.error("文档列表查询失败", error=str(exc))
|
||||
raise ApiError(2000, f"文档列表查询失败: {exc}") from exc
|
||||
return ok({"items": items, "next_offset": next_offset})
|
||||
|
||||
|
||||
@router.get("/documents/{doc_id}")
|
||||
async def get_document(doc_id: str) -> dict[str, Any]:
|
||||
"""获取文档详情:L1 记录 + L2/L3 节点 + chunks 数量"""
|
||||
try:
|
||||
detail = await _get_qdrant().get_doc_detail(doc_id)
|
||||
except Exception as exc:
|
||||
logger.error("文档详情查询失败", doc_id=doc_id, error=str(exc))
|
||||
raise ApiError(2000, f"文档详情查询失败: {exc}") from exc
|
||||
if detail is None:
|
||||
raise ApiError(1004, "文档不存在")
|
||||
return ok(detail)
|
||||
|
||||
|
||||
@router.delete("/documents/{doc_id}")
|
||||
async def delete_document(doc_id: str) -> dict[str, Any]:
|
||||
"""删除文档:四层集合中该 doc_id 的所有点;幂等,不存在也返回成功(删除数全 0)"""
|
||||
try:
|
||||
deleted = await _get_qdrant().delete_by_doc_id(doc_id)
|
||||
except Exception as exc:
|
||||
logger.error("文档删除失败", doc_id=doc_id, error=str(exc))
|
||||
raise ApiError(2000, f"文档删除失败: {exc}") from exc
|
||||
return ok({"doc_id": doc_id, "deleted": deleted, "deleted_total": sum(deleted.values())})
|
||||
@@ -0,0 +1,75 @@
|
||||
"""知识分类 API:GET /api/v1/knowledge/categories、GET /api/v1/knowledge/stats"""
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.response import ApiError, ok
|
||||
from app.config import settings
|
||||
from app.models.knowledge import UNCATEGORIZED, TaxonomyCategory, load_taxonomy
|
||||
from app.services.qdrant import ALL_COLLECTIONS, COLLECTION_L1, QdrantService
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
router = APIRouter(prefix="/api/v1", tags=["knowledge"])
|
||||
|
||||
# 类目分布统计的分页大小
|
||||
_STATS_SCROLL_PAGE_SIZE = 100
|
||||
|
||||
# 模块级懒加载单例,避免每请求重建 QdrantService
|
||||
_qdrant: QdrantService | None = None
|
||||
|
||||
|
||||
def _get_qdrant() -> QdrantService:
|
||||
global _qdrant
|
||||
if _qdrant is None:
|
||||
_qdrant = QdrantService()
|
||||
return _qdrant
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_taxonomy() -> list[TaxonomyCategory]:
|
||||
"""加载并缓存 taxonomy 类目集(进程内只加载一次)"""
|
||||
return load_taxonomy(settings.taxonomy_path)
|
||||
|
||||
|
||||
@router.get("/knowledge/categories")
|
||||
async def list_categories() -> dict[str, Any]:
|
||||
"""返回完整知识分类类目集"""
|
||||
categories = _get_taxonomy()
|
||||
return ok({"categories": [c.model_dump() for c in categories], "count": len(categories)})
|
||||
|
||||
|
||||
@router.get("/knowledge/stats")
|
||||
async def knowledge_stats() -> dict[str, Any]:
|
||||
"""返回四层集合规模与 L1 类目分布统计"""
|
||||
service = _get_qdrant()
|
||||
try:
|
||||
collections = {collection: await service.count(collection) for collection in ALL_COLLECTIONS}
|
||||
categories = await _aggregate_l1_categories(service)
|
||||
except Exception as exc:
|
||||
logger.error("获取知识库统计失败", error=str(exc))
|
||||
raise ApiError(2000, f"获取知识库统计失败: {exc}") from exc
|
||||
return ok(
|
||||
{
|
||||
"collections": collections,
|
||||
"categories": categories,
|
||||
"uncategorized_count": categories.get(UNCATEGORIZED, 0),
|
||||
"documents_total": collections[COLLECTION_L1],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _aggregate_l1_categories(service: QdrantService) -> dict[str, int]:
|
||||
"""分页遍历 L1 文档,按 category 聚合文档数(空类目归入 uncategorized)"""
|
||||
categories: dict[str, int] = {}
|
||||
offset: str | None = None
|
||||
while True:
|
||||
items, offset = await service.scroll_l1(limit=_STATS_SCROLL_PAGE_SIZE, offset=offset)
|
||||
for item in items:
|
||||
category = item.get("category") or UNCATEGORIZED
|
||||
categories[category] = categories.get(category, 0) + 1
|
||||
if offset is None:
|
||||
return categories
|
||||
@@ -0,0 +1,58 @@
|
||||
"""检索 API:POST /api/v1/search"""
|
||||
|
||||
from hashlib import sha256
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.response import ApiError, ok
|
||||
from app.core.retriever import Retriever
|
||||
from app.models.search import SearchRequest
|
||||
from app.services.redis import get_cache
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
router = APIRouter(prefix="/api/v1", tags=["search"])
|
||||
|
||||
# 模块级懒加载单例,避免每请求重建 Qdrant/Embedding client
|
||||
_retriever: Retriever | None = None
|
||||
|
||||
|
||||
def _get_retriever() -> Retriever:
|
||||
global _retriever
|
||||
if _retriever is None:
|
||||
_retriever = Retriever()
|
||||
return _retriever
|
||||
|
||||
|
||||
def _cache_key(request: SearchRequest) -> str:
|
||||
"""检索缓存键:query + top_k 的短哈希(top_k 影响结果集,需参与键计算)"""
|
||||
digest = sha256((request.query + "|" + str(request.top_k)).encode()).hexdigest()[:16]
|
||||
return f"search:{digest}"
|
||||
|
||||
|
||||
@router.post("/search")
|
||||
async def search(request: SearchRequest) -> dict[str, Any]:
|
||||
"""分层检索入口,返回统一包装的 SearchResponse
|
||||
|
||||
先查 Redis 缓存:命中直接返回缓存的响应;未命中走检索流程并回写缓存。
|
||||
缓存读写失败均降级为无缓存行为,不影响检索。
|
||||
"""
|
||||
cache_key = _cache_key(request)
|
||||
cached = await get_cache().get_json(cache_key)
|
||||
if cached is not None:
|
||||
logger.info("检索缓存命中", query=request.query, cache_key=cache_key)
|
||||
return cached
|
||||
|
||||
try:
|
||||
response = await _get_retriever().search(request)
|
||||
except ApiError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("检索失败", query=request.query)
|
||||
raise ApiError(2000, f"检索失败: {exc}") from exc
|
||||
|
||||
result = ok(response.model_dump())
|
||||
await get_cache().set_json(cache_key, result)
|
||||
return result
|
||||
@@ -0,0 +1,60 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""应用配置,通过环境变量注入"""
|
||||
|
||||
# 应用
|
||||
app_name: str = "QMDSearch"
|
||||
log_level: str = "info"
|
||||
|
||||
# 嵌入模型
|
||||
embedding_provider: str = "openai" # openai | local
|
||||
openai_api_key: str = ""
|
||||
openai_base_url: str = "https://api.openai.com/v1"
|
||||
embedding_model: str = "text-embedding-3-small"
|
||||
embedding_dimension: int = 1536
|
||||
|
||||
# Ollama 本地模型(用于文档三级总结)
|
||||
ollama_base_url: str = "http://localhost:11434"
|
||||
ollama_model: str = "qwen2.5:1.5b" # 备选: qwen2.5:3b
|
||||
ollama_embedding_model: str = "bge-m3" # embedding_provider=local 时使用的嵌入模型
|
||||
|
||||
# Qdrant
|
||||
qdrant_host: str = "localhost"
|
||||
qdrant_port: int = 6333
|
||||
|
||||
# Redis
|
||||
redis_url: str = "redis://localhost:6379/0"
|
||||
|
||||
# 检索参数
|
||||
retrieval_top_k: int = 20 # L2 语义检索召回数
|
||||
retrieval_final_k: int = 5 # L3 重排后返回数
|
||||
|
||||
# 文档入库参数
|
||||
summary_min_text_length: int = 500 # 低于此字符数触发 2.5 级回退
|
||||
|
||||
# 入库异步任务
|
||||
ingest_max_concurrency: int = 2 # 入库后台任务并发上限
|
||||
ingest_task_ttl_done: int = 86400 # 任务状态 Redis 保留秒数(进行中与已完成,24h)
|
||||
ingest_task_ttl_failed: int = 604800 # 失败任务状态 Redis 保留秒数(7 天)
|
||||
|
||||
# 知识分类(taxonomy)
|
||||
taxonomy_path: str = "" # taxonomy JSON 文件路径,为空用内置默认
|
||||
classify_confidence_threshold: float = 0.6 # 低于此值归 uncategorized
|
||||
classify_max_categories: int = 3 # query 路由命中类目数上限,超过走全库兜底
|
||||
|
||||
# 分层检索参数
|
||||
l1_doc_top_n: int = 10 # L1 层候选文档数
|
||||
l2_section_top_n: int = 5 # L2 层候选 section 数
|
||||
l3_top_n: int = 10 # L3 层定位数
|
||||
sparse_enabled: bool = True # 是否启用稀疏检索
|
||||
cache_ttl: int = 300 # Redis 缓存秒数
|
||||
|
||||
# 分块参数
|
||||
chunk_max_chars: int = 800 # chunk 超长二次切分阈值
|
||||
|
||||
model_config = {"env_prefix": "", "case_sensitive": False}
|
||||
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1,121 @@
|
||||
"""文档 chunk 切分器
|
||||
|
||||
按文档原生标题树切分 chunk:
|
||||
- 结构化文本:每个标题起点切分 section(标题行到下一个标题前),
|
||||
超长 section 按空行段落二次切分,单段落仍超长则硬切
|
||||
- 无结构文本:直接按空行段落累加切分
|
||||
|
||||
每个 chunk 记录 section_path(祖先标题链," / " 连接),与 L2 大纲节点互相定位。
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
import structlog
|
||||
|
||||
from app.config import settings
|
||||
from app.core.headings import Heading, parse_headings
|
||||
from app.models.document import ChunkModel
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 段落分隔:一个或多个空行
|
||||
_PARAGRAPH_SPLIT_PATTERN = re.compile(r"\n\s*\n")
|
||||
|
||||
|
||||
class Chunker:
|
||||
"""按标题树切分文档 chunk"""
|
||||
|
||||
def __init__(self, max_chars: int = settings.chunk_max_chars) -> None:
|
||||
self.max_chars = max_chars
|
||||
|
||||
def chunk(self, text: str, doc_id: str) -> list[ChunkModel]:
|
||||
"""将文档文本切分为 chunk 列表
|
||||
|
||||
Args:
|
||||
text: 文档纯文本内容
|
||||
doc_id: 文档 ID
|
||||
|
||||
Returns:
|
||||
list[ChunkModel]: 切分结果,chunk_index 从 0 递增
|
||||
"""
|
||||
stripped = text.strip()
|
||||
if not stripped:
|
||||
return []
|
||||
|
||||
# 全文不超长:整篇单 chunk
|
||||
if len(stripped) <= self.max_chars:
|
||||
return [ChunkModel(doc_id=doc_id, chunk_index=0, text=stripped)]
|
||||
|
||||
headings = parse_headings(stripped)
|
||||
chunks: list[ChunkModel] = []
|
||||
|
||||
if headings:
|
||||
# 结构化:先按标题切分 section,再按长度二次切分
|
||||
for section_text, section_path in self._split_sections(stripped, headings):
|
||||
for piece in self._split_by_length(section_text):
|
||||
chunks.append(
|
||||
ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece, section_path=section_path)
|
||||
)
|
||||
else:
|
||||
# 无结构:直接按段落累加切分
|
||||
for piece in self._split_by_length(stripped):
|
||||
chunks.append(ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece))
|
||||
|
||||
logger.info("文档切分完成", doc_id=doc_id, chunks_count=len(chunks), has_headings=bool(headings))
|
||||
return chunks
|
||||
|
||||
def _split_sections(self, text: str, headings: list[Heading]) -> list[tuple[str, str]]:
|
||||
"""按标题树切分 section,返回 (section 文本, section_path) 列表
|
||||
|
||||
每个标题起点切分一个 section,section 文本含标题行本身;
|
||||
section_path 为祖先标题链(含自身标题),用 " / " 连接;
|
||||
首个标题前的引导正文归入无前缀 section(section_path 为空)。
|
||||
"""
|
||||
lines = text.splitlines()
|
||||
sections: list[tuple[str, str]] = []
|
||||
|
||||
# 首个标题前的引导内容
|
||||
preamble = "\n".join(lines[: headings[0].line_index]).strip()
|
||||
if preamble:
|
||||
sections.append((preamble, ""))
|
||||
|
||||
# 维护祖先标题栈:遇到同级或更高级标题时弹栈
|
||||
stack: list[Heading] = []
|
||||
for i, heading in enumerate(headings):
|
||||
while stack and stack[-1].level >= heading.level:
|
||||
stack.pop()
|
||||
stack.append(heading)
|
||||
|
||||
end = headings[i + 1].line_index if i + 1 < len(headings) else len(lines)
|
||||
section_text = "\n".join(lines[heading.line_index : end]).strip()
|
||||
section_path = " / ".join(h.title for h in stack)
|
||||
sections.append((section_text, section_path))
|
||||
|
||||
return sections
|
||||
|
||||
def _split_by_length(self, text: str) -> list[str]:
|
||||
"""按 max_chars 切分文本:先按空行段落累加,单段落超长则硬切"""
|
||||
if len(text) <= self.max_chars:
|
||||
return [text]
|
||||
|
||||
pieces: list[str] = []
|
||||
current = ""
|
||||
for paragraph in _PARAGRAPH_SPLIT_PATTERN.split(text):
|
||||
paragraph = paragraph.strip()
|
||||
if not paragraph:
|
||||
continue
|
||||
candidate = f"{current}\n\n{paragraph}" if current else paragraph
|
||||
if len(candidate) <= self.max_chars:
|
||||
current = candidate
|
||||
continue
|
||||
if current:
|
||||
pieces.append(current)
|
||||
current = ""
|
||||
# 单段落仍超长:按 max_chars 硬切
|
||||
if len(paragraph) > self.max_chars:
|
||||
pieces.extend(paragraph[i : i + self.max_chars] for i in range(0, len(paragraph), self.max_chars))
|
||||
else:
|
||||
current = paragraph
|
||||
if current:
|
||||
pieces.append(current)
|
||||
return pieces
|
||||
@@ -0,0 +1,139 @@
|
||||
"""文档分类器
|
||||
|
||||
入库链路第二步:基于 L1 总结,用 Ollama 小模型将文档判定为 taxonomy 中的
|
||||
主类目 + 附加标签:
|
||||
- LLM 输出解析失败 / 类目名不在 taxonomy → 归 uncategorized(confidence=0.0)
|
||||
- 置信度低于阈值 → 主类目归 uncategorized,候选类目名保留进 tags(软召回用)
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
|
||||
import structlog
|
||||
|
||||
from app.config import settings
|
||||
from app.models.knowledge import UNCATEGORIZED, CategoryResult, TaxonomyCategory, load_taxonomy
|
||||
from app.services.ollama import OllamaClient
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 从 LLM 输出中提取第一个 {...} JSON 块(贪婪匹配到最后的 },兼容嵌套对象)
|
||||
_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
|
||||
|
||||
|
||||
def _extract_json(raw: str) -> dict | None:
|
||||
"""从 LLM 输出中提取 JSON 对象
|
||||
|
||||
先尝试直接解析;失败则用正则提取第一个 {...} 块再解析。
|
||||
返回 None 表示无法提取出合法的 JSON 对象。
|
||||
"""
|
||||
text = raw.strip()
|
||||
try:
|
||||
data = json.loads(text)
|
||||
return data if isinstance(data, dict) else None
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
match = _JSON_BLOCK_RE.search(text)
|
||||
if not match:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(match.group(0))
|
||||
return data if isinstance(data, dict) else None
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
|
||||
class Classifier:
|
||||
"""文档分类器:将 L1 总结判定为 taxonomy 主类目 + 附加标签"""
|
||||
|
||||
def __init__(self, ollama: OllamaClient | None = None, taxonomy: list[TaxonomyCategory] | None = None) -> None:
|
||||
self.ollama = ollama or OllamaClient()
|
||||
self.taxonomy = taxonomy if taxonomy is not None else load_taxonomy()
|
||||
# 合法类目名集合(含 uncategorized)
|
||||
self._valid_names = {c.name for c in self.taxonomy}
|
||||
|
||||
async def classify(self, l1_summary: str, title: str = "") -> CategoryResult:
|
||||
"""对文档进行分类判定
|
||||
|
||||
Args:
|
||||
l1_summary: 文档 L1 总结
|
||||
title: 文档标题(可选,辅助判定)
|
||||
|
||||
Returns:
|
||||
CategoryResult: 主类目 / 附加标签 / 置信度
|
||||
"""
|
||||
prompt = self._build_prompt(l1_summary, title)
|
||||
raw = await self.ollama.generate(prompt, json_mode=True)
|
||||
|
||||
data = _extract_json(raw)
|
||||
if data is None:
|
||||
logger.warning("分类失败:LLM 输出非合法 JSON", title=title, output=raw[:200])
|
||||
return CategoryResult(main_category=UNCATEGORIZED, tags=[], confidence=0.0)
|
||||
|
||||
main_category = data.get("main_category")
|
||||
confidence = data.get("confidence")
|
||||
if (
|
||||
not isinstance(main_category, str)
|
||||
or main_category not in self._valid_names
|
||||
or not isinstance(confidence, (int, float))
|
||||
or not 0 <= confidence <= 1
|
||||
):
|
||||
logger.warning(
|
||||
"分类失败:类目名不在 taxonomy 或 confidence 非法",
|
||||
title=title,
|
||||
main_category=main_category,
|
||||
confidence=confidence,
|
||||
)
|
||||
return CategoryResult(main_category=UNCATEGORIZED, tags=[], confidence=0.0)
|
||||
|
||||
tags = self._clean_tags(data.get("tags"), main_category)
|
||||
confidence = float(confidence)
|
||||
|
||||
# 低置信度软召回:主类目归 uncategorized,候选类目名保留进 tags
|
||||
if confidence < settings.classify_confidence_threshold:
|
||||
logger.info(
|
||||
"分类置信度低于阈值,归入 uncategorized",
|
||||
title=title,
|
||||
candidate=main_category,
|
||||
confidence=confidence,
|
||||
threshold=settings.classify_confidence_threshold,
|
||||
)
|
||||
if main_category != UNCATEGORIZED and main_category not in tags:
|
||||
tags.insert(0, main_category)
|
||||
return CategoryResult(main_category=UNCATEGORIZED, tags=tags, confidence=confidence)
|
||||
|
||||
return CategoryResult(main_category=main_category, tags=tags, confidence=confidence)
|
||||
|
||||
def _build_prompt(self, l1_summary: str, title: str) -> str:
|
||||
"""构造分类 prompt:列出全部 taxonomy 类目(含 uncategorized 及其用途说明)"""
|
||||
category_lines = []
|
||||
for c in self.taxonomy:
|
||||
if c.name == UNCATEGORIZED:
|
||||
category_lines.append(f"- {c.name}: 当文档跨多个类目或无法明确归入其他类目时选择此类目")
|
||||
else:
|
||||
category_lines.append(f"- {c.name}: {c.description}")
|
||||
category_block = "\n".join(category_lines)
|
||||
|
||||
prompt = (
|
||||
"你是知识库分类助手。请根据文档标题和总结,判断文档最适合归入以下哪个类目。\n\n"
|
||||
f"可选类目:\n{category_block}\n\n"
|
||||
"要求:\n"
|
||||
"- main_category 只能从上面的类目名中选择,不要输出其他名称;\n"
|
||||
"- tags 为 0~3 个附加标签(词或短语),且不能包含 main_category 本身;\n"
|
||||
"- confidence 为 0~1 之间的小数,表示对主类目判断的置信度;\n"
|
||||
"- 只输出 JSON,不要输出任何其他内容。\n\n"
|
||||
'输出格式:{"main_category": "类目名", "tags": ["标签"], "confidence": 0.0}\n\n'
|
||||
)
|
||||
if title:
|
||||
prompt += f"文档标题:{title}\n"
|
||||
prompt += f"文档总结:{l1_summary}"
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
def _clean_tags(raw_tags: object, main_category: str) -> list[str]:
|
||||
"""清洗 LLM 输出的 tags:仅保留非空字符串、剔除主类目名、最多 3 个"""
|
||||
if not isinstance(raw_tags, list):
|
||||
return []
|
||||
tags = [t for t in raw_tags if isinstance(t, str) and t and t != main_category]
|
||||
return tags[:3]
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Dense 向量嵌入服务
|
||||
|
||||
提供统一的 EmbeddingService 接口,支持两种 provider:
|
||||
- openai:OpenAI 兼容 API(AsyncOpenAI)
|
||||
- local:本地 Ollama /api/embed 接口
|
||||
"""
|
||||
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
import httpx
|
||||
import structlog
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EmbeddingService(Protocol):
|
||||
"""统一嵌入服务接口"""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
"""批量生成文本向量
|
||||
|
||||
Args:
|
||||
texts: 待嵌入文本列表,空列表时直接返回空列表
|
||||
|
||||
Returns:
|
||||
与输入等长的向量列表
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
def _check_dimension(vectors: list[list[float]], provider: str) -> None:
|
||||
"""返回维度与配置不一致时告警(不抛错),每次调用最多提示一次"""
|
||||
for vec in vectors:
|
||||
if len(vec) != settings.embedding_dimension:
|
||||
logger.warning(
|
||||
"嵌入向量维度与配置不一致",
|
||||
provider=provider,
|
||||
actual=len(vec),
|
||||
expected=settings.embedding_dimension,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
class OpenAIEmbeddingService:
|
||||
"""OpenAI 兼容 API 嵌入服务"""
|
||||
|
||||
def __init__(self, api_key: str, base_url: str, model: str) -> None:
|
||||
self._client = AsyncOpenAI(api_key=api_key, base_url=base_url)
|
||||
self._model = model
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
resp = await self._client.embeddings.create(model=self._model, input=texts)
|
||||
vectors = [item.embedding for item in resp.data]
|
||||
_check_dimension(vectors, provider="openai")
|
||||
logger.debug("OpenAI 嵌入完成", model=self._model, count=len(vectors))
|
||||
return vectors
|
||||
|
||||
|
||||
class LocalEmbeddingService:
|
||||
"""本地 Ollama 嵌入服务(/api/embed)"""
|
||||
|
||||
def __init__(self, base_url: str, model: str, timeout: float = 60.0) -> None:
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
url = f"{self.base_url}/api/embed"
|
||||
payload = {"model": self.model, "input": texts}
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
resp = await client.post(url, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
vectors: list[list[float]] = data.get("embeddings", [])
|
||||
_check_dimension(vectors, provider="local")
|
||||
logger.debug("Ollama 嵌入完成", model=self.model, count=len(vectors))
|
||||
return vectors
|
||||
|
||||
|
||||
def create_embedding_service() -> EmbeddingService:
|
||||
"""按 settings.embedding_provider 创建嵌入服务实例"""
|
||||
if settings.embedding_provider == "local":
|
||||
return LocalEmbeddingService(
|
||||
base_url=settings.ollama_base_url,
|
||||
model=settings.ollama_embedding_model,
|
||||
)
|
||||
return OpenAIEmbeddingService(
|
||||
api_key=settings.openai_api_key,
|
||||
base_url=settings.openai_base_url,
|
||||
model=settings.embedding_model,
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
"""标题树解析器
|
||||
|
||||
从纯文本中解析文档原生标题结构,支持两类模式:
|
||||
- Markdown ATX 标题(# ~ ######,# 数量即层级)
|
||||
- 中文编号标题(第X章/节/篇、一、1.1 等编号,行长度不超过 60 字符)
|
||||
|
||||
解析结果用于 L2 大纲生成与 chunk 的 section 切分。
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class Heading(BaseModel):
|
||||
"""文档标题节点"""
|
||||
|
||||
title: str = Field(description="标题文本")
|
||||
level: int = Field(description="标题层级,1 为最顶层")
|
||||
line_index: int = Field(description="标题所在行号(从 0 开始)")
|
||||
|
||||
|
||||
# Markdown ATX 标题:1~6 个 # 后跟空白
|
||||
_ATX_PATTERN = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
|
||||
|
||||
# 中文篇章节编号:第X章/第X节/第X篇(章/篇=1 级,节=2 级)
|
||||
_CN_CHAPTER_PATTERN = re.compile(r"^第[一二三四五六七八九十百\d]+([章节篇])")
|
||||
|
||||
# 中文序号:一、二、……,固定 1 级
|
||||
_CN_ENUM_PATTERN = re.compile(r"^[一二三四五六七八九十]+、")
|
||||
|
||||
# 数字编号:1. / 1、/ 1.1 / 1.1.1 等,按点分段数定层级
|
||||
_NUM_PATTERN = re.compile(r"^(\d+(?:\.\d+)*)[、.\s]")
|
||||
|
||||
# 编号类标题行的最大长度,超过则视为正文
|
||||
_MAX_HEADING_LINE_LENGTH = 60
|
||||
|
||||
# 第X[章节篇] 后缀对应的层级
|
||||
_CN_CHAPTER_LEVELS = {"章": 1, "节": 2, "篇": 1}
|
||||
|
||||
|
||||
def parse_headings(text: str) -> list[Heading]:
|
||||
"""解析文本中的标题,按行号升序返回
|
||||
|
||||
Args:
|
||||
text: 文档纯文本内容
|
||||
|
||||
Returns:
|
||||
list[Heading]: 标题列表,无标题时返回空列表
|
||||
"""
|
||||
headings: list[Heading] = []
|
||||
for line_index, line in enumerate(text.splitlines()):
|
||||
heading = _match_heading(line, line_index)
|
||||
if heading is not None:
|
||||
headings.append(heading)
|
||||
return headings
|
||||
|
||||
|
||||
def render_outline(headings: list[Heading]) -> str:
|
||||
"""将标题树渲染为大纲文本,每行一个节点,按层级缩进"""
|
||||
return "\n".join(f"{' ' * (h.level - 1)}- {h.title}" for h in headings)
|
||||
|
||||
|
||||
def _match_heading(line: str, line_index: int) -> Heading | None:
|
||||
"""匹配单行是否为标题,是则返回 Heading,否则返回 None"""
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
return None
|
||||
|
||||
# Markdown ATX 标题(无长度限制)
|
||||
match = _ATX_PATTERN.match(stripped)
|
||||
if match:
|
||||
return Heading(title=match.group(2), level=len(match.group(1)), line_index=line_index)
|
||||
|
||||
# 编号类标题有行长度限制,过长视为正文
|
||||
if len(stripped) > _MAX_HEADING_LINE_LENGTH:
|
||||
return None
|
||||
|
||||
# 第X章/节/篇
|
||||
match = _CN_CHAPTER_PATTERN.match(stripped)
|
||||
if match:
|
||||
return Heading(title=stripped, level=_CN_CHAPTER_LEVELS[match.group(1)], line_index=line_index)
|
||||
|
||||
# 一、二、……
|
||||
if _CN_ENUM_PATTERN.match(stripped):
|
||||
return Heading(title=stripped, level=1, line_index=line_index)
|
||||
|
||||
# 数字编号,层级 = 点分段数(1.=1、1.1=2、1.1.1=3)
|
||||
match = _NUM_PATTERN.match(stripped)
|
||||
if match:
|
||||
level = match.group(1).count(".") + 1
|
||||
return Heading(title=stripped, level=level, line_index=line_index)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,172 @@
|
||||
"""入库异步任务管理器
|
||||
|
||||
将文档入库包装为后台异步任务:submit 登记任务并立即返回 task_id,
|
||||
后台受并发上限控制执行 Ingester.ingest,并按阶段推进任务状态。
|
||||
|
||||
内存注册表为主(记录 status/created_at/updated_at/result/error),
|
||||
Redis 为持久镜像(key: ingest_task:{task_id}),每次状态迁移同步写入;
|
||||
Redis 不可用或写入失败仅记录 warning,不影响任务执行。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
|
||||
from app.config import Settings
|
||||
from app.core.ingestion import Ingester, IngestionError
|
||||
from app.models.document import DocumentInput
|
||||
from app.services.redis import RedisCache
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# Redis 任务状态 key 前缀
|
||||
REDIS_KEY_PREFIX = "ingest_task:"
|
||||
|
||||
|
||||
class IngestTaskStatus(StrEnum):
|
||||
"""入库任务状态"""
|
||||
|
||||
PENDING = "pending" # 已登记,排队等待执行
|
||||
SUMMARIZING = "summarizing" # 三级总结中
|
||||
CLASSIFYING = "classifying" # 分类判定中
|
||||
EMBEDDING = "embedding" # 向量化中
|
||||
WRITING = "writing" # 写入 Qdrant 中
|
||||
DONE = "done" # 入库完成
|
||||
FAILED = "failed" # 入库失败
|
||||
|
||||
|
||||
# 终态集合
|
||||
TERMINAL_STATUSES: frozenset[str] = frozenset({IngestTaskStatus.DONE, IngestTaskStatus.FAILED})
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
"""当前 UTC 时间的 ISO8601 字符串"""
|
||||
return datetime.now(UTC).isoformat()
|
||||
|
||||
|
||||
class IngestTaskManager:
|
||||
"""入库异步任务管理器:登记、后台执行、状态查询与 Redis 持久镜像"""
|
||||
|
||||
def __init__(self, ingester: Ingester, redis: RedisCache | None, settings: Settings) -> None:
|
||||
self._ingester = ingester
|
||||
self._redis = redis
|
||||
self._settings = settings
|
||||
self._tasks: dict[str, dict[str, Any]] = {}
|
||||
self._semaphore = asyncio.Semaphore(settings.ingest_max_concurrency)
|
||||
# 持有后台任务与镜像任务引用,避免被 GC 提前回收
|
||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
||||
self._mirror_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
async def submit(self, doc: DocumentInput) -> str:
|
||||
"""登记入库任务并后台执行,立即返回 task_id"""
|
||||
task_id = uuid.uuid4().hex
|
||||
now = _utc_now_iso()
|
||||
self._tasks[task_id] = {
|
||||
"task_id": task_id,
|
||||
"status": IngestTaskStatus.PENDING,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"result": None,
|
||||
"error": None,
|
||||
}
|
||||
self._schedule_mirror(task_id)
|
||||
background = asyncio.create_task(self._run(task_id, doc))
|
||||
self._background_tasks.add(background)
|
||||
background.add_done_callback(self._background_tasks.discard)
|
||||
logger.info("入库任务已登记", task_id=task_id, title=doc.title)
|
||||
return task_id
|
||||
|
||||
async def get(self, task_id: str) -> dict[str, Any] | None:
|
||||
"""查询任务状态:先查内存注册表,miss 再查 Redis 镜像,都没有返回 None"""
|
||||
record = self._tasks.get(task_id)
|
||||
if record is not None:
|
||||
return record
|
||||
if self._redis is None:
|
||||
return None
|
||||
try:
|
||||
return await self._redis.get_json(f"{REDIS_KEY_PREFIX}{task_id}")
|
||||
except Exception:
|
||||
logger.warning("入库任务状态读取 Redis 失败,降级为未命中", task_id=task_id, exc_info=True)
|
||||
return None
|
||||
|
||||
async def wait_done(self, task_id: str, timeout: float = 30.0) -> dict[str, Any]:
|
||||
"""轮询内存注册表直到任务进入终态(done/failed)或超时
|
||||
|
||||
进入终态后会等待已调度的 Redis 镜像写完再返回;超时抛 TimeoutError。
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + timeout
|
||||
while True:
|
||||
record = self._tasks.get(task_id)
|
||||
if record is not None and record["status"] in TERMINAL_STATUSES:
|
||||
if self._mirror_tasks:
|
||||
await asyncio.gather(*self._mirror_tasks, return_exceptions=True)
|
||||
return record
|
||||
if loop.time() >= deadline:
|
||||
raise TimeoutError(f"入库任务 {task_id} 在 {timeout}s 内未进入终态")
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def _run(self, task_id: str, doc: DocumentInput) -> None:
|
||||
"""后台执行入库:并发限流 + 阶段状态推进 + 结果/错误落账"""
|
||||
async with self._semaphore:
|
||||
try:
|
||||
result = await self._ingester.ingest(doc, progress_cb=lambda stage: self._on_progress(task_id, stage))
|
||||
except IngestionError as exc:
|
||||
# 入库已知失败:透传阶段与已产出的部分总结
|
||||
self._finish_failed(
|
||||
task_id,
|
||||
{
|
||||
"stage": exc.stage,
|
||||
"message": str(exc),
|
||||
"partial_summary": exc.summary.model_dump(mode="json") if exc.summary is not None else None,
|
||||
},
|
||||
)
|
||||
except Exception as exc:
|
||||
self._finish_failed(task_id, {"stage": "unknown", "message": str(exc), "partial_summary": None})
|
||||
else:
|
||||
self._tasks[task_id].update(
|
||||
status=IngestTaskStatus.DONE,
|
||||
updated_at=_utc_now_iso(),
|
||||
result=result.model_dump(mode="json"),
|
||||
)
|
||||
self._schedule_mirror(task_id)
|
||||
logger.info("入库任务完成", task_id=task_id)
|
||||
|
||||
def _finish_failed(self, task_id: str, error: dict[str, Any]) -> None:
|
||||
"""将任务置为 failed 并记录错误信息"""
|
||||
self._tasks[task_id].update(status=IngestTaskStatus.FAILED, updated_at=_utc_now_iso(), error=error)
|
||||
self._schedule_mirror(task_id)
|
||||
logger.error("入库任务失败", task_id=task_id, stage=error["stage"], error=error["message"])
|
||||
|
||||
def _on_progress(self, task_id: str, stage: str) -> None:
|
||||
"""Ingester 阶段回调:推进任务状态并同步镜像(同步函数,供 progress_cb 使用)"""
|
||||
record = self._tasks.get(task_id)
|
||||
if record is None:
|
||||
return
|
||||
record["status"] = stage
|
||||
record["updated_at"] = _utc_now_iso()
|
||||
self._schedule_mirror(task_id)
|
||||
|
||||
def _schedule_mirror(self, task_id: str) -> None:
|
||||
"""将当前任务状态快照异步镜像到 Redis(同步上下文也可调用)"""
|
||||
if self._redis is None:
|
||||
return
|
||||
mirror = asyncio.create_task(self._mirror_to_redis(task_id, dict(self._tasks[task_id])))
|
||||
self._mirror_tasks.add(mirror)
|
||||
mirror.add_done_callback(self._mirror_tasks.discard)
|
||||
|
||||
async def _mirror_to_redis(self, task_id: str, snapshot: dict[str, Any]) -> None:
|
||||
"""写入 Redis 镜像:进行中与 done 用 ttl_done,failed 用 ttl_failed;写失败仅告警"""
|
||||
ttl = (
|
||||
self._settings.ingest_task_ttl_failed
|
||||
if snapshot["status"] == IngestTaskStatus.FAILED
|
||||
else self._settings.ingest_task_ttl_done
|
||||
)
|
||||
try:
|
||||
await self._redis.set_json(f"{REDIS_KEY_PREFIX}{task_id}", snapshot, ttl=ttl)
|
||||
except Exception:
|
||||
logger.warning("入库任务状态镜像 Redis 失败", task_id=task_id, exc_info=True)
|
||||
@@ -0,0 +1,306 @@
|
||||
"""文档入库模块
|
||||
|
||||
入库流程:文档输入 → 三级总结(Ollama) → 分类判定(L1总结) → 切分 chunk
|
||||
→ 构建 L2/L3 大纲节点 → 批量向量化(dense + sparse)→ 写入 Qdrant 四层集合
|
||||
|
||||
L2/L3 大纲节点的构建策略见 _build_l2_nodes / _build_l3_nodes。
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
|
||||
from app.config import settings
|
||||
from app.core.chunker import Chunker
|
||||
from app.core.classifier import Classifier
|
||||
from app.core.embeddings import EmbeddingService, create_embedding_service
|
||||
from app.core.headings import Heading, parse_headings
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.core.summarizer import Summarizer
|
||||
from app.models.document import ChunkModel, DocumentInput, DocumentSummary, IngestionResult, SummaryLevel
|
||||
from app.models.knowledge import CategoryResult
|
||||
from app.services.qdrant import COLLECTION_L2, COLLECTION_L3, QdrantService, SparseVectorTuple
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# IngestionError 阶段标识
|
||||
STAGE_SUMMARIZE = "summarize"
|
||||
STAGE_CLASSIFY = "classify"
|
||||
STAGE_EMBED = "embed"
|
||||
STAGE_QDRANT = "qdrant"
|
||||
|
||||
|
||||
class IngestionError(Exception):
|
||||
"""入库失败异常
|
||||
|
||||
携带失败阶段(stage)与已产出的总结(summary,如有),
|
||||
Qdrant 写入失败时总结不丢,上层可按阶段重试。
|
||||
"""
|
||||
|
||||
def __init__(self, stage: str, message: str, summary: DocumentSummary | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.stage = stage
|
||||
self.summary = summary
|
||||
|
||||
|
||||
def _heading_paths(headings: list[Heading]) -> list[tuple[str, str]]:
|
||||
"""按文档顺序计算每个标题的 (标题文本, 祖先标题链含自身),链用 " / " 连接"""
|
||||
paths: list[tuple[str, str]] = []
|
||||
stack: list[Heading] = []
|
||||
for heading in headings:
|
||||
# 遇到同级或更高级标题时弹栈,维护当前祖先链
|
||||
while stack and stack[-1].level >= heading.level:
|
||||
stack.pop()
|
||||
stack.append(heading)
|
||||
paths.append((heading.title, " / ".join(h.title for h in stack)))
|
||||
return paths
|
||||
|
||||
|
||||
def _split_l3_blocks(outline: str) -> list[tuple[str, str]]:
|
||||
"""将内容大纲按 "## " 行分块,返回 (块标题, 块文本) 列表
|
||||
|
||||
无 "## " 行时整块作为一个节点(块标题为空);
|
||||
首个 "## " 之前的引导内容直接忽略。
|
||||
"""
|
||||
stripped = outline.strip()
|
||||
if not stripped:
|
||||
return []
|
||||
|
||||
blocks: list[tuple[str, list[str]]] = []
|
||||
for line in stripped.splitlines():
|
||||
if line.startswith("## "):
|
||||
blocks.append((line[3:].strip(), [line]))
|
||||
elif blocks:
|
||||
blocks[-1][1].append(line)
|
||||
if not blocks:
|
||||
return [("", stripped)]
|
||||
return [(title, "\n".join(lines).strip()) for title, lines in blocks]
|
||||
|
||||
|
||||
class Ingester:
|
||||
"""文档入库器:编排总结、分类、切分、向量化与 Qdrant 写入全链路"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
summarizer: Summarizer | None = None,
|
||||
classifier: Classifier | None = None,
|
||||
chunker: Chunker | None = None,
|
||||
embedding: EmbeddingService | None = None,
|
||||
sparse: SparseEncoder | None = None,
|
||||
qdrant: QdrantService | None = None,
|
||||
) -> None:
|
||||
self.summarizer = summarizer or Summarizer()
|
||||
self.classifier = classifier or Classifier()
|
||||
self.chunker = chunker or Chunker()
|
||||
self.embedding = embedding or create_embedding_service()
|
||||
self.sparse = sparse or SparseEncoder()
|
||||
self.qdrant = qdrant or QdrantService()
|
||||
|
||||
async def ingest(self, doc: DocumentInput, progress_cb: Callable[[str], None] | None = None) -> IngestionResult:
|
||||
"""执行文档入库
|
||||
|
||||
Args:
|
||||
doc: 文档输入(文本内容 + 元数据)
|
||||
progress_cb: 可选的阶段进度回调(同步函数),在各阶段边界以
|
||||
"summarizing" / "classifying" / "embedding" / "writing" 调用
|
||||
|
||||
Returns:
|
||||
IngestionResult: 入库结果
|
||||
|
||||
Raises:
|
||||
IngestionError: 任一阶段失败时抛出,携带 stage 与已产出总结
|
||||
"""
|
||||
|
||||
def _report(stage: str) -> None:
|
||||
if progress_cb is not None:
|
||||
progress_cb(stage)
|
||||
|
||||
logger.info("开始文档入库", title=doc.title, text_length=len(doc.text))
|
||||
doc_id = uuid.uuid4().hex
|
||||
|
||||
# 1. 三级总结
|
||||
_report("summarizing")
|
||||
try:
|
||||
summary = await self.summarizer.summarize(doc.text, title=doc.title)
|
||||
except Exception as exc:
|
||||
logger.error("入库失败:三级总结", stage=STAGE_SUMMARIZE, error=str(exc))
|
||||
raise IngestionError(STAGE_SUMMARIZE, f"三级总结失败: {exc}") from exc
|
||||
logger.info("三级总结完成", doc_id=doc_id, level=summary.level.value)
|
||||
|
||||
# 2. 分类判定(基于 L1 总结)
|
||||
_report("classifying")
|
||||
try:
|
||||
category = await self.classifier.classify(summary.l1_summary, title=doc.title)
|
||||
except Exception as exc:
|
||||
logger.error("入库失败:分类判定", stage=STAGE_CLASSIFY, error=str(exc))
|
||||
raise IngestionError(STAGE_CLASSIFY, f"分类判定失败: {exc}", summary=summary) from exc
|
||||
logger.info("分类判定完成", doc_id=doc_id, category=category.main_category, confidence=category.confidence)
|
||||
|
||||
# 3. 切分 chunk 并构建 L2/L3 大纲节点((text, section_path) 列表)
|
||||
chunks = self.chunker.chunk(doc.text, doc_id)
|
||||
l2_nodes = self._build_l2_nodes(doc.text, summary)
|
||||
l3_nodes = self._build_l3_nodes(doc.text, summary)
|
||||
|
||||
# 4. 批量 embedding:L1 + L2 + L3 + chunks 一次调用,按序切片取向量
|
||||
texts = [
|
||||
summary.l1_summary,
|
||||
*(node_text for node_text, _ in l2_nodes),
|
||||
*(node_text for node_text, _ in l3_nodes),
|
||||
*(c.text for c in chunks),
|
||||
]
|
||||
try:
|
||||
_report("embedding")
|
||||
vectors = await self.embedding.embed(texts)
|
||||
except Exception as exc:
|
||||
logger.error("入库失败:向量化", stage=STAGE_EMBED, error=str(exc))
|
||||
raise IngestionError(STAGE_EMBED, f"向量化失败: {exc}", summary=summary) from exc
|
||||
l1_vector = vectors[0]
|
||||
l2_vectors = vectors[1 : 1 + len(l2_nodes)]
|
||||
l3_vectors = vectors[1 + len(l2_nodes) : 1 + len(l2_nodes) + len(l3_nodes)]
|
||||
chunk_vectors = vectors[1 + len(l2_nodes) + len(l3_nodes) :]
|
||||
|
||||
# 5. sparse 向量(仅 L1 与 chunks 需要)
|
||||
l1_sparse: SparseVectorTuple | None = None
|
||||
chunk_sparses: list[SparseVectorTuple | None] = [None] * len(chunks)
|
||||
if settings.sparse_enabled:
|
||||
l1_sparse = self.sparse.encode(summary.l1_summary)
|
||||
chunk_sparses = [self.sparse.encode(c.text) for c in chunks]
|
||||
|
||||
# 6. 写入 Qdrant 四层集合
|
||||
_report("writing")
|
||||
try:
|
||||
await self._write_qdrant(
|
||||
doc_id,
|
||||
doc,
|
||||
summary,
|
||||
category,
|
||||
chunks,
|
||||
l2_nodes,
|
||||
l3_nodes,
|
||||
l1_vector,
|
||||
l2_vectors,
|
||||
l3_vectors,
|
||||
chunk_vectors,
|
||||
l1_sparse,
|
||||
chunk_sparses,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("入库失败:Qdrant 写入", stage=STAGE_QDRANT, doc_id=doc_id, error=str(exc))
|
||||
raise IngestionError(STAGE_QDRANT, f"Qdrant 写入失败: {exc}", summary=summary) from exc
|
||||
|
||||
logger.info(
|
||||
"文档入库完成",
|
||||
doc_id=doc_id,
|
||||
chunks_count=len(chunks),
|
||||
l2_nodes=len(l2_nodes),
|
||||
l3_nodes=len(l3_nodes),
|
||||
category=category.main_category,
|
||||
)
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
summary=summary,
|
||||
category=category.main_category,
|
||||
tags=category.tags,
|
||||
category_confidence=category.confidence,
|
||||
collection="四层集合",
|
||||
chunks_count=len(chunks),
|
||||
)
|
||||
|
||||
def _build_l2_nodes(self, text: str, summary: DocumentSummary) -> list[tuple[str, str]]:
|
||||
"""构建 L2 大纲节点,返回 (text, section_path) 列表
|
||||
|
||||
- 有标题结构(L3 级且标题数 >= 2):每个标题一个节点,
|
||||
text 与 section_path 均为该节点的祖先标题链
|
||||
- 否则若 l2_outline 非空(LLM 生成的大纲):按非空行拆节点,section_path 为空
|
||||
- 2.5 级文档(l2_outline 为 None):无 L2 节点
|
||||
"""
|
||||
headings = parse_headings(text)
|
||||
if summary.level == SummaryLevel.L3 and len(headings) >= 2:
|
||||
return [(path, path) for _, path in _heading_paths(headings)]
|
||||
if summary.l2_outline:
|
||||
return [(line.strip(), "") for line in summary.l2_outline.splitlines() if line.strip()]
|
||||
return []
|
||||
|
||||
def _build_l3_nodes(self, text: str, summary: DocumentSummary) -> list[tuple[str, str]]:
|
||||
"""构建 L3 内容大纲节点,返回 (text, section_path) 列表
|
||||
|
||||
按 "## " 分块(无 "## " 则整块一个节点);
|
||||
section_path 尽力匹配文档标题链(块标题与文档标题文本精确匹配),匹配不到用 ""。
|
||||
"""
|
||||
path_by_title: dict[str, str] = {}
|
||||
for title, path in _heading_paths(parse_headings(text)):
|
||||
path_by_title.setdefault(title, path)
|
||||
return [
|
||||
(block_text, path_by_title.get(block_title, ""))
|
||||
for block_title, block_text in _split_l3_blocks(summary.l3_content_outline)
|
||||
]
|
||||
|
||||
async def _write_qdrant(
|
||||
self,
|
||||
doc_id: str,
|
||||
doc: DocumentInput,
|
||||
summary: DocumentSummary,
|
||||
category: CategoryResult,
|
||||
chunks: list[ChunkModel],
|
||||
l2_nodes: list[tuple[str, str]],
|
||||
l3_nodes: list[tuple[str, str]],
|
||||
l1_vector: list[float],
|
||||
l2_vectors: list[list[float]],
|
||||
l3_vectors: list[list[float]],
|
||||
chunk_vectors: list[list[float]],
|
||||
l1_sparse: SparseVectorTuple | None,
|
||||
chunk_sparses: list[SparseVectorTuple | None],
|
||||
) -> None:
|
||||
"""将 L1/L2/L3/chunks 四层数据写入 Qdrant(任一失败向上抛出)"""
|
||||
await self.qdrant.upsert_l1(
|
||||
doc_id=doc_id,
|
||||
title=doc.title,
|
||||
summary=summary.l1_summary,
|
||||
category=category.main_category,
|
||||
tags=category.tags,
|
||||
dense_vector=l1_vector,
|
||||
sparse_vector=l1_sparse,
|
||||
)
|
||||
|
||||
# L2/L3 大纲节点(为空时跳过对应集合的 upsert)
|
||||
for collection, nodes, vectors in (
|
||||
(COLLECTION_L2, l2_nodes, l2_vectors),
|
||||
(COLLECTION_L3, l3_nodes, l3_vectors),
|
||||
):
|
||||
if not nodes:
|
||||
continue
|
||||
await self.qdrant.upsert_nodes(
|
||||
collection,
|
||||
[
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"section_path": section_path,
|
||||
"text": node_text,
|
||||
"category": category.main_category,
|
||||
"tags": category.tags,
|
||||
"dense_vector": vector,
|
||||
}
|
||||
for (node_text, section_path), vector in zip(nodes, vectors, strict=True)
|
||||
],
|
||||
)
|
||||
|
||||
if chunks:
|
||||
# chunk dict 额外携带 doc_summary(= L1 总结),检索侧直接取用,不用回查 L1
|
||||
chunk_dicts: list[dict[str, Any]] = [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"text": chunk.text,
|
||||
"section_path": chunk.section_path,
|
||||
"title": doc.title,
|
||||
"category": category.main_category,
|
||||
"tags": category.tags,
|
||||
"dense_vector": vector,
|
||||
"sparse_vector": sparse,
|
||||
"doc_summary": summary.l1_summary,
|
||||
}
|
||||
for chunk, vector, sparse in zip(chunks, chunk_vectors, chunk_sparses, strict=True)
|
||||
]
|
||||
await self.qdrant.upsert_chunks(chunk_dicts)
|
||||
@@ -0,0 +1,214 @@
|
||||
"""query 解析与分类路由模块
|
||||
|
||||
分层 RAG 在线侧第一步:用 Ollama 小模型将用户 query 解析为结构化 JSON
|
||||
(命中类目+置信度、rewrite 后 query、关键词),再由纯函数做路由决策:
|
||||
- 高置信且命中类目数 <= 上限 → 按类目过滤检索
|
||||
- 低置信 / 解析失败 / 命中类目过多 → 全库兜底(不丢召回)
|
||||
|
||||
解析(LLM 调用)与决策(纯函数)分离,便于单元测试。
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from hashlib import sha256
|
||||
|
||||
import structlog
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from app.config import settings
|
||||
from app.models.knowledge import UNCATEGORIZED, TaxonomyCategory
|
||||
from app.services.ollama import OllamaClient
|
||||
from app.services.redis import RedisCache, get_cache
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 从 LLM 输出中提取第一个 {...} JSON 块(贪婪匹配到最后的 },兼容嵌套对象)
|
||||
_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
|
||||
|
||||
|
||||
class CategoryHit(BaseModel):
|
||||
"""query 命中的类目及置信度"""
|
||||
|
||||
name: str = Field(description="类目名称,必须来自 taxonomy")
|
||||
confidence: float = Field(ge=0, le=1, description="命中置信度,范围 [0, 1]")
|
||||
|
||||
|
||||
class ParsedQuery(BaseModel):
|
||||
"""query 解析结果"""
|
||||
|
||||
raw_query: str = Field(description="原始 query")
|
||||
rewrite: str = Field(description="rewrite 后的 query")
|
||||
keywords: list[str] = Field(default_factory=list, description="提取的关键词")
|
||||
categories: list[CategoryHit] = Field(default_factory=list, description="命中类目列表")
|
||||
parse_failed: bool = Field(default=False, description="LLM 输出解析是否失败")
|
||||
|
||||
|
||||
class RouteDecision(BaseModel):
|
||||
"""路由决策结果"""
|
||||
|
||||
fallback: bool = Field(description="是否走全库兜底")
|
||||
filter_categories: list[str] | None = Field(default=None, description="过滤类目名列表,None 表示全库不过滤")
|
||||
reason: str = Field(description="决策原因:parse_failed | low_confidence | too_many_categories | routed")
|
||||
parsed: ParsedQuery = Field(description="对应的 query 解析结果")
|
||||
|
||||
|
||||
def _extract_json(raw: str) -> dict | None:
|
||||
"""从 LLM 输出中提取 JSON 对象
|
||||
|
||||
先尝试直接解析;失败则用正则提取第一个 {...} 块再解析。
|
||||
返回 None 表示无法提取出合法的 JSON 对象。
|
||||
"""
|
||||
text = raw.strip()
|
||||
try:
|
||||
data = json.loads(text)
|
||||
return data if isinstance(data, dict) else None
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
match = _JSON_BLOCK_RE.search(text)
|
||||
if not match:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(match.group(0))
|
||||
return data if isinstance(data, dict) else None
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
|
||||
def decide_route(parsed: ParsedQuery, threshold: float, max_categories: int) -> RouteDecision:
|
||||
"""路由决策(纯函数)
|
||||
|
||||
- 解析失败 → 全库兜底(parse_failed)
|
||||
- 无 confidence >= threshold 的类目 → 全库兜底(low_confidence)
|
||||
- 命中类目数 > max_categories → 全库兜底(too_many_categories)
|
||||
- 否则按类目过滤检索,类目按 confidence 降序排列(routed)
|
||||
"""
|
||||
if parsed.parse_failed:
|
||||
return RouteDecision(fallback=True, filter_categories=None, reason="parse_failed", parsed=parsed)
|
||||
|
||||
hits = [c for c in parsed.categories if c.confidence >= threshold]
|
||||
if not hits:
|
||||
return RouteDecision(fallback=True, filter_categories=None, reason="low_confidence", parsed=parsed)
|
||||
|
||||
if len(hits) > max_categories:
|
||||
return RouteDecision(fallback=True, filter_categories=None, reason="too_many_categories", parsed=parsed)
|
||||
|
||||
hits.sort(key=lambda c: c.confidence, reverse=True)
|
||||
return RouteDecision(
|
||||
fallback=False,
|
||||
filter_categories=[c.name for c in hits],
|
||||
reason="routed",
|
||||
parsed=parsed,
|
||||
)
|
||||
|
||||
|
||||
class QueryParser:
|
||||
"""query 解析器:调用 Ollama 小模型将 query 解析为结构化 JSON"""
|
||||
|
||||
def __init__(
|
||||
self, ollama: OllamaClient, taxonomy: list[TaxonomyCategory], cache: RedisCache | None = None
|
||||
) -> None:
|
||||
self.ollama = ollama
|
||||
self.taxonomy = taxonomy
|
||||
# 解析结果缓存,缺省用全局单例;RedisCache 全操作容错,缓存不可用时退化为无缓存行为
|
||||
self.cache = cache if cache is not None else get_cache()
|
||||
# 可作为路由命中类目的名字集合(uncategorized 不可作为路由命中类目)
|
||||
self._routable_names = {c.name for c in taxonomy if c.name != UNCATEGORIZED}
|
||||
|
||||
async def parse(self, query: str) -> ParsedQuery:
|
||||
"""调用 LLM 将 query 解析为结构化结果
|
||||
|
||||
解析失败(输出非 JSON / 必填字段缺失或类型错误)时,
|
||||
返回 parse_failed=True 的兜底结果(rewrite 为原 query)。
|
||||
"""
|
||||
prompt = self._build_prompt(query)
|
||||
raw = await self.ollama.generate(prompt, json_mode=True)
|
||||
|
||||
data = _extract_json(raw)
|
||||
if data is None:
|
||||
logger.warning("query 解析失败:LLM 输出非合法 JSON", query=query, output=raw[:200])
|
||||
return ParsedQuery(raw_query=query, rewrite=query, parse_failed=True)
|
||||
|
||||
rewrite = data.get("rewrite")
|
||||
keywords = data.get("keywords")
|
||||
categories = data.get("categories")
|
||||
if not isinstance(rewrite, str) or not isinstance(keywords, list) or not isinstance(categories, list):
|
||||
logger.warning("query 解析失败:JSON 必填字段缺失或类型错误", query=query, output=raw[:200])
|
||||
return ParsedQuery(raw_query=query, rewrite=query, parse_failed=True)
|
||||
|
||||
return ParsedQuery(
|
||||
raw_query=query,
|
||||
rewrite=rewrite,
|
||||
keywords=[str(k) for k in keywords],
|
||||
categories=self._validate_categories(categories, query),
|
||||
)
|
||||
|
||||
async def parse_and_route(self, query: str) -> RouteDecision:
|
||||
"""解析 query 并做路由决策(threshold / max_categories 取自 settings)
|
||||
|
||||
parse() 的 LLM 解析结果按 query 哈希缓存;路由决策每次现算(纯函数,
|
||||
阈值取最新 settings),缓存数据损坏时回退为重新走 LLM 解析。
|
||||
"""
|
||||
cache_key = f"qparse:{sha256(query.encode()).hexdigest()[:16]}"
|
||||
parsed = await self._get_cached_parsed(cache_key, query)
|
||||
if parsed is None:
|
||||
parsed = await self.parse(query)
|
||||
await self.cache.set_json(cache_key, parsed.model_dump())
|
||||
return decide_route(
|
||||
parsed,
|
||||
threshold=settings.classify_confidence_threshold,
|
||||
max_categories=settings.classify_max_categories,
|
||||
)
|
||||
|
||||
async def _get_cached_parsed(self, cache_key: str, query: str) -> ParsedQuery | None:
|
||||
"""读取缓存的解析结果;未命中或缓存数据无法重建时返回 None"""
|
||||
cached = await self.cache.get_json(cache_key)
|
||||
if cached is None:
|
||||
return None
|
||||
try:
|
||||
parsed = ParsedQuery.model_validate(cached)
|
||||
except ValidationError:
|
||||
logger.warning("query 解析缓存数据损坏,按未命中处理", query=query, cache_key=cache_key)
|
||||
return None
|
||||
logger.info("query 解析缓存命中", query=query, cache_key=cache_key)
|
||||
return parsed
|
||||
|
||||
def _build_prompt(self, query: str) -> str:
|
||||
"""构造解析 prompt:列出 taxonomy 类目名+描述,要求模型只输出 JSON"""
|
||||
category_lines = [f"- {c.name}: {c.description}" for c in self.taxonomy if c.name != UNCATEGORIZED]
|
||||
category_block = "\n".join(category_lines)
|
||||
return (
|
||||
"你是搜索查询分析助手。请分析用户 query,完成三件事:\n"
|
||||
"1. 判断 query 意图命中以下哪些知识类目,并给出每个类目的置信度(0~1 之间的小数);\n"
|
||||
"2. 将 query 改写为更适合检索的形式;\n"
|
||||
"3. 提取 query 的关键词。\n\n"
|
||||
f"可选类目:\n{category_block}\n\n"
|
||||
"要求:\n"
|
||||
"- categories 中的 name 只能从上面的类目名中选择,不要输出其他名称;\n"
|
||||
"- 若没有明显命中的类目,categories 返回空列表;\n"
|
||||
"- 只输出 JSON,不要输出任何其他内容。\n\n"
|
||||
'输出格式:{"categories": [{"name": "类目名", "confidence": 0.0}], '
|
||||
'"rewrite": "改写后的 query", "keywords": ["关键词"]}\n\n'
|
||||
f"用户 query:{query}"
|
||||
)
|
||||
|
||||
def _validate_categories(self, categories: list, query: str) -> list[CategoryHit]:
|
||||
"""校验并过滤 LLM 输出的类目条目
|
||||
|
||||
丢弃:类目名不在 taxonomy 可路由类目中的条目(含 uncategorized)、
|
||||
confidence 缺失或越界(不在 [0, 1])的条目。
|
||||
"""
|
||||
hits: list[CategoryHit] = []
|
||||
for item in categories:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = item.get("name")
|
||||
confidence = item.get("confidence")
|
||||
if not isinstance(name, str) or name not in self._routable_names:
|
||||
logger.warning("丢弃不在 taxonomy 中的类目", query=query, name=name)
|
||||
continue
|
||||
if not isinstance(confidence, (int, float)) or not 0 <= confidence <= 1:
|
||||
logger.warning("丢弃 confidence 越界的类目", query=query, name=name, confidence=confidence)
|
||||
continue
|
||||
hits.append(CategoryHit(name=name, confidence=float(confidence)))
|
||||
return hits
|
||||
@@ -0,0 +1,24 @@
|
||||
"""检索结果重排:RRF 融合与最终截断"""
|
||||
|
||||
from qdrant_client import models
|
||||
|
||||
|
||||
def rrf_fuse(result_lists: list[list[models.ScoredPoint]], k: int = 60) -> list[models.ScoredPoint]:
|
||||
"""标准 RRF(Reciprocal Rank Fusion)融合
|
||||
|
||||
融合分 = Σ 1/(k + rank)(rank 从 1 开始),按 point id 去重合并,
|
||||
返回按融合分降序的列表,score 字段写回融合分。空输入返回 []。
|
||||
"""
|
||||
scores: dict[str | int, float] = {}
|
||||
points: dict[str | int, models.ScoredPoint] = {}
|
||||
for results in result_lists:
|
||||
for rank, point in enumerate(results, start=1):
|
||||
scores[point.id] = scores.get(point.id, 0.0) + 1.0 / (k + rank)
|
||||
points.setdefault(point.id, point)
|
||||
ordered = sorted(points, key=lambda pid: scores[pid], reverse=True)
|
||||
return [points[pid].model_copy(update={"score": scores[pid]}) for pid in ordered]
|
||||
|
||||
|
||||
def finalize(points: list[models.ScoredPoint], final_k: int) -> list[models.ScoredPoint]:
|
||||
"""截断为最终返回的 top final_k"""
|
||||
return points[:final_k]
|
||||
@@ -0,0 +1,169 @@
|
||||
"""分层检索引擎
|
||||
|
||||
L1(文档总结)→ L2(章节大纲)→ L3(小节定位)→ chunks(原文)逐层收窄:
|
||||
- L1 无候选文档:直接全库 chunk 兜底,fallback=True
|
||||
- L2 无命中:全部候选文档回退到 L3 的 doc 级查询(b 路)
|
||||
- 2.5 级文档无 L2 节点,天然落入 L3 b 路(仅按 doc 过滤)
|
||||
- L3 两路(a:L2 命中文档按 section 过滤;b:其余文档仅按 doc 过滤)RRF 融合;
|
||||
L3 无命中时 chunk 层回退为 L1 候选文档级检索
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import structlog
|
||||
from qdrant_client import models
|
||||
|
||||
from app.config import settings
|
||||
from app.core.embeddings import EmbeddingService, create_embedding_service
|
||||
from app.core.query_parser import QueryParser
|
||||
from app.core.ranker import finalize, rrf_fuse
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.models.knowledge import load_taxonomy
|
||||
from app.models.search import SearchHit, SearchRequest, SearchResponse
|
||||
from app.services.ollama import OllamaClient
|
||||
from app.services.qdrant import (
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
SPARSE_COLLECTIONS,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
def _unique(values: Iterable[str | None]) -> list[str]:
|
||||
"""去重保序并丢弃空值"""
|
||||
return list(dict.fromkeys(v for v in values if v))
|
||||
|
||||
|
||||
class Retriever:
|
||||
"""分层检索引擎,依赖均可注入(默认自建,便于测试替换)"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
qdrant: QdrantService | None = None,
|
||||
query_parser: QueryParser | None = None,
|
||||
embedding: EmbeddingService | None = None,
|
||||
sparse_encoder: SparseEncoder | None = None,
|
||||
) -> None:
|
||||
self.qdrant = qdrant or QdrantService()
|
||||
self.query_parser = query_parser or QueryParser(
|
||||
ollama=OllamaClient(),
|
||||
taxonomy=load_taxonomy(settings.taxonomy_path),
|
||||
)
|
||||
self.embedding = embedding or create_embedding_service()
|
||||
self.sparse_encoder = sparse_encoder or SparseEncoder()
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
"""分层检索主流程"""
|
||||
route = await self.query_parser.parse_and_route(request.query)
|
||||
query_text = route.parsed.rewrite or request.query
|
||||
dense = (await self.embedding.embed([query_text]))[0]
|
||||
sparse = self.sparse_encoder.encode(query_text) if settings.sparse_enabled else None
|
||||
# 路由兜底时不做类目过滤,避免丢召回
|
||||
categories = None if route.fallback else route.filter_categories
|
||||
|
||||
# L1:文档级检索,产出候选文档
|
||||
l1_filter = QdrantService.build_filter(categories=categories)
|
||||
l1_hits = await self._search_collection(COLLECTION_L1, dense, sparse, settings.l1_doc_top_n, l1_filter)
|
||||
logger.info("L1 检索完成", hits=len(l1_hits), categories=categories)
|
||||
|
||||
if not l1_hits:
|
||||
# L1 无候选文档 → 全库 chunk 兜底
|
||||
chunk_hits = await self._search_collection(COLLECTION_CHUNKS, dense, sparse, settings.retrieval_top_k, None)
|
||||
logger.info("L1 无命中,全库 chunk 兜底", hits=len(chunk_hits))
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
hits=self._to_hits(finalize(chunk_hits, self._final_k(request))),
|
||||
routed_categories=route.filter_categories or [],
|
||||
fallback=True,
|
||||
)
|
||||
|
||||
doc_ids = _unique((p.payload or {}).get("doc_id") for p in l1_hits)
|
||||
|
||||
# L2:候选文档内检索章节大纲
|
||||
l2_filter = QdrantService.build_filter(categories=categories, doc_ids=doc_ids)
|
||||
l2_hits = await self._search_collection(
|
||||
COLLECTION_L2, dense, sparse, settings.l2_section_top_n * len(doc_ids), l2_filter
|
||||
)
|
||||
logger.info("L2 检索完成", hits=len(l2_hits))
|
||||
l2_doc_ids = _unique((p.payload or {}).get("doc_id") for p in l2_hits)
|
||||
l2_section_paths = _unique((p.payload or {}).get("section_path") for p in l2_hits)
|
||||
# 无 L2 命中的文档(2.5 级文档或该 doc 的 L2 未命中)走 L3 b 路
|
||||
remaining_doc_ids = [d for d in doc_ids if d not in l2_doc_ids]
|
||||
|
||||
# L3:两路查询后 RRF 融合
|
||||
l3_lists: list[list[models.ScoredPoint]] = []
|
||||
if l2_doc_ids:
|
||||
l3_filter_a = QdrantService.build_filter(
|
||||
categories=categories, doc_ids=l2_doc_ids, section_paths=l2_section_paths
|
||||
)
|
||||
l3_lists.append(await self._search_collection(COLLECTION_L3, dense, sparse, settings.l3_top_n, l3_filter_a))
|
||||
if remaining_doc_ids:
|
||||
l3_filter_b = QdrantService.build_filter(categories=categories, doc_ids=remaining_doc_ids)
|
||||
l3_lists.append(await self._search_collection(COLLECTION_L3, dense, sparse, settings.l3_top_n, l3_filter_b))
|
||||
l3_hits = rrf_fuse(l3_lists)
|
||||
logger.info("L3 检索完成", hits=len(l3_hits))
|
||||
|
||||
# chunk 层:L3 有命中按 section 收窄;无命中回退为 L1 候选文档级检索
|
||||
if l3_hits:
|
||||
l3_doc_ids = _unique((p.payload or {}).get("doc_id") for p in l3_hits)
|
||||
# L3 命中 section_path 全为空(如 2.5 级文档)时仅按 doc 过滤
|
||||
l3_section_paths = _unique((p.payload or {}).get("section_path") for p in l3_hits)
|
||||
chunk_filter = QdrantService.build_filter(
|
||||
categories=categories, doc_ids=l3_doc_ids, section_paths=l3_section_paths
|
||||
)
|
||||
else:
|
||||
logger.info("L3 无命中,回退到 L1 候选文档级 chunk 检索")
|
||||
chunk_filter = QdrantService.build_filter(categories=categories, doc_ids=doc_ids)
|
||||
|
||||
chunk_hits = await self._search_collection(
|
||||
COLLECTION_CHUNKS, dense, sparse, settings.retrieval_top_k, chunk_filter
|
||||
)
|
||||
logger.info("chunk 检索完成", hits=len(chunk_hits))
|
||||
|
||||
final_points = finalize(rrf_fuse([chunk_hits]), self._final_k(request))
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
hits=self._to_hits(final_points),
|
||||
routed_categories=route.filter_categories or [],
|
||||
fallback=route.fallback,
|
||||
)
|
||||
|
||||
async def _search_collection(
|
||||
self,
|
||||
collection: str,
|
||||
dense: list[float],
|
||||
sparse: tuple[list[int], list[float]] | None,
|
||||
limit: int,
|
||||
query_filter: models.Filter | None,
|
||||
) -> list[models.ScoredPoint]:
|
||||
"""按集合是否支持 sparse 选择 hybrid 或 dense 检索"""
|
||||
if sparse is not None and collection in SPARSE_COLLECTIONS:
|
||||
return await self.qdrant.search_hybrid(collection, dense, sparse, limit, query_filter=query_filter)
|
||||
return await self.qdrant.search_dense(collection, dense, limit, query_filter=query_filter)
|
||||
|
||||
@staticmethod
|
||||
def _final_k(request: SearchRequest) -> int:
|
||||
"""最终返回数:请求指定优先,否则用配置默认值"""
|
||||
return request.top_k or settings.retrieval_final_k
|
||||
|
||||
@staticmethod
|
||||
def _to_hits(points: list[models.ScoredPoint]) -> list[SearchHit]:
|
||||
"""chunk 点组装为 SearchHit,字段取自 chunk payload(doc_summary 仅上下文标注)"""
|
||||
hits: list[SearchHit] = []
|
||||
for p in points:
|
||||
payload = p.payload or {}
|
||||
hits.append(
|
||||
SearchHit(
|
||||
text=payload.get("text", ""),
|
||||
doc_id=payload.get("doc_id", ""),
|
||||
title=payload.get("title", ""),
|
||||
section_path=payload.get("section_path") or "",
|
||||
score=p.score,
|
||||
doc_summary=payload.get("doc_summary") or "",
|
||||
)
|
||||
)
|
||||
return hits
|
||||
@@ -0,0 +1,68 @@
|
||||
"""BM25 轻量近似稀疏向量编码器
|
||||
|
||||
零第三方依赖的稀疏编码实现:无全局 IDF 的 BM25 近似(仅用词频 tf 加权),
|
||||
输出 Qdrant SparseVector 所需的 indices/values 格式。
|
||||
后续可替换为 SPLADE / BM42 等更强的稀疏编码器。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import math
|
||||
import re
|
||||
|
||||
# 稀疏向量维度(哈希空间大小)
|
||||
SPARSE_DIM = 2**18
|
||||
|
||||
# 连续 CJK 字符段(一-鿿)
|
||||
_CJK_RE = re.compile(r"[一-鿿]+")
|
||||
# CJK 段或英文/数字连续段
|
||||
_TOKEN_RE = re.compile(r"[一-鿿]+|[a-zA-Z0-9]+")
|
||||
|
||||
|
||||
def _tokenize(text: str) -> list[str]:
|
||||
"""分词(纯正则实现)
|
||||
|
||||
- 连续 CJK 字符段:长度 1 时保留 unigram,长度 >= 2 时生成字符 bigram
|
||||
- 英文/数字连续段:小写化后作为整词
|
||||
"""
|
||||
tokens: list[str] = []
|
||||
for match in _TOKEN_RE.finditer(text):
|
||||
seg = match.group()
|
||||
if _CJK_RE.fullmatch(seg):
|
||||
if len(seg) == 1:
|
||||
tokens.append(seg)
|
||||
else:
|
||||
tokens.extend(seg[i : i + 2] for i in range(len(seg) - 1))
|
||||
else:
|
||||
tokens.append(seg.lower())
|
||||
return tokens
|
||||
|
||||
|
||||
def _hash(token: str) -> int:
|
||||
"""将词哈希到 [0, SPARSE_DIM) 的索引空间"""
|
||||
digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
|
||||
return int.from_bytes(digest, "big") % SPARSE_DIM
|
||||
|
||||
|
||||
class SparseEncoder:
|
||||
"""稀疏向量编码器,输出 Qdrant SparseVector 的 indices/values"""
|
||||
|
||||
def encode(self, text: str) -> tuple[list[int], list[float]]:
|
||||
"""编码单条文本
|
||||
|
||||
权重 = 1 + log(tf),词经哈希到 [0, SPARSE_DIM),哈希冲突时权重累加。
|
||||
|
||||
Returns:
|
||||
(indices, values):indices 升序且无重复,与 values 等长
|
||||
"""
|
||||
# 按哈希索引累计词频,天然处理哈希冲突
|
||||
tf: dict[int, float] = {}
|
||||
for token in _tokenize(text):
|
||||
idx = _hash(token)
|
||||
tf[idx] = tf.get(idx, 0.0) + 1.0
|
||||
indices = sorted(tf)
|
||||
values = [1.0 + math.log(tf[idx]) for idx in indices]
|
||||
return indices, values
|
||||
|
||||
def encode_batch(self, texts: list[str]) -> list[tuple[list[int], list[float]]]:
|
||||
"""批量编码,与逐条 encode 结果一致"""
|
||||
return [self.encode(text) for text in texts]
|
||||
@@ -0,0 +1,138 @@
|
||||
"""文档三级总结模块
|
||||
|
||||
通过 Ollama 本地小模型对文档进行分级总结:
|
||||
- L1: 总结(一句话高度概括)
|
||||
- L2: 大纲(主要章节和关键主题)
|
||||
- L3: 内容大纲(每个章节的详细内容摘要)
|
||||
|
||||
当文档内容不足以支撑三级总结时,自动降级为 2.5 级(L1 + L2.5 内容大纲)。
|
||||
"""
|
||||
|
||||
import structlog
|
||||
|
||||
from app.core.headings import parse_headings, render_outline
|
||||
from app.models.document import DocumentSummary, SummaryLevel
|
||||
from app.services.ollama import OllamaClient
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 文本长度阈值(字符数),低于此值触发 2.5 级回退
|
||||
MIN_TEXT_LENGTH_FOR_L3 = 500
|
||||
|
||||
|
||||
class Summarizer:
|
||||
"""文档三级总结器"""
|
||||
|
||||
def __init__(self, ollama: OllamaClient | None = None) -> None:
|
||||
self.ollama = ollama or OllamaClient()
|
||||
|
||||
async def summarize(self, text: str, *, title: str = "") -> DocumentSummary:
|
||||
"""对文档文本进行三级总结
|
||||
|
||||
Args:
|
||||
text: 文档纯文本内容
|
||||
title: 文档标题(可选,辅助总结)
|
||||
|
||||
Returns:
|
||||
DocumentSummary: 包含各级总结的结果
|
||||
"""
|
||||
text_length = len(text.strip())
|
||||
|
||||
# 生成 L1 总结
|
||||
l1_summary = await self._generate_l1(text, title)
|
||||
logger.info("L1 总结完成", title=title, l1=l1_summary[:100])
|
||||
|
||||
# 判断是否需要 2.5 级回退
|
||||
use_fallback = text_length < MIN_TEXT_LENGTH_FOR_L3
|
||||
|
||||
if use_fallback:
|
||||
logger.info("文档内容不足,使用 2.5 级回退", text_length=text_length)
|
||||
# L2.5: 直接生成内容大纲(跳过大纲层)
|
||||
l2_half = await self._generate_l2_half(text, title, l1_summary)
|
||||
return DocumentSummary(
|
||||
l1_summary=l1_summary,
|
||||
l2_outline=None,
|
||||
l3_content_outline=l2_half,
|
||||
level=SummaryLevel.L2_HALF,
|
||||
)
|
||||
|
||||
# 生成 L2 大纲:优先使用文档原生标题树(结构导航,不调 LLM),
|
||||
# 无结构文本(标题数 < 2)回退 LLM 生成
|
||||
headings = parse_headings(text)
|
||||
if len(headings) >= 2:
|
||||
l2_outline = render_outline(headings)
|
||||
logger.info("L2 大纲完成", title=title, outline_source="headings", heading_count=len(headings))
|
||||
else:
|
||||
l2_outline = await self._generate_l2(text, title, l1_summary)
|
||||
logger.info("L2 大纲完成", title=title, outline_source="llm")
|
||||
|
||||
# 生成 L3 内容大纲
|
||||
l3_content_outline = await self._generate_l3(text, title, l1_summary, l2_outline)
|
||||
logger.info("L3 内容大纲完成", title=title)
|
||||
|
||||
return DocumentSummary(
|
||||
l1_summary=l1_summary,
|
||||
l2_outline=l2_outline,
|
||||
l3_content_outline=l3_content_outline,
|
||||
level=SummaryLevel.L3,
|
||||
)
|
||||
|
||||
async def _generate_l1(self, text: str, title: str) -> str:
|
||||
"""生成 L1 总结:一句话高度概括"""
|
||||
prompt = (
|
||||
"请用一句话对以下文档内容进行高度概括,要求简洁精炼,"
|
||||
"突出文档的核心主题和关键信息。\n\n"
|
||||
)
|
||||
if title:
|
||||
prompt += f"文档标题:{title}\n\n"
|
||||
prompt += f"文档内容:\n{text}"
|
||||
|
||||
return await self.ollama.generate(prompt)
|
||||
|
||||
async def _generate_l2(self, text: str, title: str, l1_summary: str) -> str:
|
||||
"""生成 L2 大纲:主要章节和关键主题"""
|
||||
prompt = (
|
||||
"请提取以下文档的主要章节结构和关键主题,"
|
||||
"以大纲形式呈现。每个主题用一行表示,"
|
||||
"格式为:序号. 主题名称\n\n"
|
||||
f"文档总结:{l1_summary}\n\n"
|
||||
)
|
||||
if title:
|
||||
prompt += f"文档标题:{title}\n\n"
|
||||
prompt += f"文档内容:\n{text}"
|
||||
|
||||
return await self.ollama.generate(prompt)
|
||||
|
||||
async def _generate_l3(
|
||||
self, text: str, title: str, l1_summary: str, l2_outline: str
|
||||
) -> str:
|
||||
"""生成 L3 内容大纲:每个章节的详细内容摘要"""
|
||||
prompt = (
|
||||
"请对以下文档的每个章节/主题进行详细的内容摘要,"
|
||||
"格式为:\n"
|
||||
"## 章节名称\n"
|
||||
"详细摘要内容...\n\n"
|
||||
f"文档总结:{l1_summary}\n\n"
|
||||
f"文档大纲:\n{l2_outline}\n\n"
|
||||
)
|
||||
if title:
|
||||
prompt += f"文档标题:{title}\n\n"
|
||||
prompt += f"文档内容:\n{text}"
|
||||
|
||||
return await self.ollama.generate(prompt)
|
||||
|
||||
async def _generate_l2_half(self, text: str, title: str, l1_summary: str) -> str:
|
||||
"""生成 L2.5 内容大纲(2.5 级回退)
|
||||
|
||||
跳过大纲层,直接对短文本生成详细摘要。
|
||||
"""
|
||||
prompt = (
|
||||
"请对以下文档内容进行详细摘要,"
|
||||
"涵盖所有关键信息点。以要点形式呈现。\n\n"
|
||||
f"文档总结:{l1_summary}\n\n"
|
||||
)
|
||||
if title:
|
||||
prompt += f"文档标题:{title}\n\n"
|
||||
prompt += f"文档内容:\n{text}"
|
||||
|
||||
return await self.ollama.generate(prompt)
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import structlog
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import FileResponse, JSONResponse
|
||||
|
||||
from app.api.response import ApiError, error
|
||||
from app.api.v1.document import router as document_router
|
||||
from app.api.v1.knowledge import router as knowledge_router
|
||||
from app.api.v1.search import router as search_router
|
||||
from app.config import settings
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
"""应用生命周期:启动时初始化 Qdrant 集合
|
||||
|
||||
初始化失败仅记录日志、不阻止启动(本地开发可能无 Qdrant)。
|
||||
"""
|
||||
try:
|
||||
await QdrantService().ensure_collections()
|
||||
logger.info("Qdrant 集合初始化完成")
|
||||
except Exception:
|
||||
logger.error("Qdrant 集合初始化失败,跳过初始化继续启动")
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.app_name,
|
||||
description="AI Agent 分层信息检索服务",
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
app.include_router(search_router)
|
||||
app.include_router(document_router)
|
||||
app.include_router(knowledge_router)
|
||||
|
||||
|
||||
@app.exception_handler(ApiError)
|
||||
async def api_error_handler(request: Request, exc: ApiError) -> JSONResponse:
|
||||
"""业务异常 → 统一错误响应(code 取自异常)"""
|
||||
return JSONResponse(content=error(exc.code, exc.message))
|
||||
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def validation_error_handler(request: Request, exc: RequestValidationError) -> JSONResponse:
|
||||
"""请求参数校验失败 → code 1001"""
|
||||
errors = exc.errors()
|
||||
message = f"请求参数校验失败: {errors[0]['msg']}" if errors else "请求参数校验失败"
|
||||
return JSONResponse(content=error(1001, message))
|
||||
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
async def unhandled_error_handler(request: Request, exc: Exception) -> JSONResponse:
|
||||
"""未捕获异常 → code 2000"""
|
||||
logger.exception("未捕获异常", path=request.url.path)
|
||||
return JSONResponse(content=error(2000, "服务器内部错误"))
|
||||
|
||||
|
||||
@app.get("/api/v1/health")
|
||||
async def health() -> dict[str, str]:
|
||||
"""健康检查"""
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
_ADMIN_HTML = Path(__file__).resolve().parent / "static" / "admin.html"
|
||||
|
||||
|
||||
@app.get("/admin", include_in_schema=False)
|
||||
async def admin_page() -> FileResponse:
|
||||
"""管理后台单页(单文件静态 HTML,零外部依赖)"""
|
||||
return FileResponse(_ADMIN_HTML, media_type="text/html")
|
||||
@@ -0,0 +1,51 @@
|
||||
"""文档和总结相关的数据模型"""
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SummaryLevel(StrEnum):
|
||||
"""总结层级"""
|
||||
|
||||
L3 = "L3" # 完整三级:总结 → 大纲 → 内容大纲
|
||||
L2_HALF = "L2.5" # 2.5 级回退:总结 → 内容大纲(跳过大纲)
|
||||
|
||||
|
||||
class DocumentSummary(BaseModel):
|
||||
"""文档三级总结结果"""
|
||||
|
||||
l1_summary: str = Field(description="L1 总结:一句话高度概括")
|
||||
l2_outline: str | None = Field(default=None, description="L2 大纲:主要章节和关键主题")
|
||||
l3_content_outline: str = Field(description="L3/L2.5 内容大纲:详细内容摘要")
|
||||
level: SummaryLevel = Field(description="实际使用的总结层级")
|
||||
|
||||
|
||||
class DocumentInput(BaseModel):
|
||||
"""文档入库输入"""
|
||||
|
||||
text: str = Field(description="文档纯文本内容")
|
||||
title: str = Field(default="", description="文档标题")
|
||||
source: str = Field(default="", description="来源标识(文件路径/URL等)")
|
||||
metadata: dict[str, str] = Field(default_factory=dict, description="附加元数据")
|
||||
|
||||
|
||||
class ChunkModel(BaseModel):
|
||||
"""文档分块结果"""
|
||||
|
||||
doc_id: str = Field(description="所属文档 ID")
|
||||
chunk_index: int = Field(description="chunk 在文档内的序号")
|
||||
text: str = Field(description="chunk 文本内容")
|
||||
section_path: str = Field(default="", description="chunk 所在的章节路径")
|
||||
|
||||
|
||||
class IngestionResult(BaseModel):
|
||||
"""文档入库结果"""
|
||||
|
||||
document_id: str = Field(description="写入后的文档 ID")
|
||||
summary: DocumentSummary = Field(description="三级总结结果")
|
||||
category: str = Field(description="分类标签(主类目)")
|
||||
collection: str = Field(description="写入的 Qdrant 集合名")
|
||||
chunks_count: int = Field(default=0, description="写入的 chunk 数量")
|
||||
tags: list[str] = Field(default_factory=list, description="附加分类标签")
|
||||
category_confidence: float = Field(default=0.0, description="主类目分类置信度")
|
||||
@@ -0,0 +1,68 @@
|
||||
"""知识分类(taxonomy)相关的数据模型与加载逻辑"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import structlog
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 未分类常量:分类置信度不足或无法归类时使用
|
||||
UNCATEGORIZED = "uncategorized"
|
||||
|
||||
|
||||
class TaxonomyCategory(BaseModel):
|
||||
"""知识分类类目定义"""
|
||||
|
||||
name: str = Field(description="类目名称,全局唯一")
|
||||
description: str = Field(default="", description="类目描述,用于辅助分类判断")
|
||||
|
||||
|
||||
class CategoryResult(BaseModel):
|
||||
"""文档/查询的分类结果"""
|
||||
|
||||
main_category: str = Field(description="主类目名称")
|
||||
tags: list[str] = Field(default_factory=list, description="附加标签列表")
|
||||
confidence: float = Field(ge=0, le=1, description="分类置信度,范围 [0, 1]")
|
||||
|
||||
|
||||
def _default_taxonomy() -> list[TaxonomyCategory]:
|
||||
"""内置默认类目集(通用企业知识库场景)"""
|
||||
return [
|
||||
TaxonomyCategory(name="技术文档", description="架构设计、API 文档、开发规范、运维手册等技术资料"),
|
||||
TaxonomyCategory(name="产品手册", description="产品功能介绍、使用说明、版本发布说明"),
|
||||
TaxonomyCategory(name="运营规范", description="运营流程、活动方案、内容规范、客服话术"),
|
||||
TaxonomyCategory(name="财务行政", description="财务制度、报销流程、行政通知、办公管理"),
|
||||
TaxonomyCategory(name="市场资料", description="市场分析、竞品调研、营销素材、品牌规范"),
|
||||
TaxonomyCategory(name="人事制度", description="招聘、考勤、绩效、培训、员工手册等 HR 制度"),
|
||||
TaxonomyCategory(name="法律法规", description="合同模板、合规要求、法律条文、知识产权"),
|
||||
TaxonomyCategory(name=UNCATEGORIZED, description="无法归入其他类目的文档"),
|
||||
]
|
||||
|
||||
|
||||
def load_taxonomy(path: str = "") -> list[TaxonomyCategory]:
|
||||
"""加载 taxonomy 类目集
|
||||
|
||||
path 为空时使用内置默认类目集;非空时从 JSON 文件加载,
|
||||
文件格式为 [{"name": ..., "description": ...}]。
|
||||
校验类目 name 唯一;若缺少 uncategorized 类目则自动追加。
|
||||
"""
|
||||
if not path:
|
||||
return _default_taxonomy()
|
||||
|
||||
raw = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
categories = [TaxonomyCategory.model_validate(item) for item in raw]
|
||||
|
||||
# 校验 name 唯一
|
||||
names = [c.name for c in categories]
|
||||
if len(names) != len(set(names)):
|
||||
duplicates = sorted({n for n in names if names.count(n) > 1})
|
||||
raise ValueError(f"taxonomy 类目 name 重复: {duplicates}")
|
||||
|
||||
# 必含 uncategorized,缺失则自动追加
|
||||
if UNCATEGORIZED not in names:
|
||||
logger.warning("taxonomy 缺少 uncategorized 类目,已自动追加", path=path)
|
||||
categories.append(TaxonomyCategory(name=UNCATEGORIZED, description="无法归入其他类目的文档"))
|
||||
|
||||
return categories
|
||||
@@ -0,0 +1,30 @@
|
||||
"""检索请求与响应的数据模型"""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""检索请求"""
|
||||
|
||||
query: str = Field(description="查询文本")
|
||||
top_k: int | None = Field(default=None, description="返回结果数,为空时使用 settings.retrieval_final_k")
|
||||
|
||||
|
||||
class SearchHit(BaseModel):
|
||||
"""单条检索命中结果"""
|
||||
|
||||
text: str = Field(description="原文 chunk 内容")
|
||||
doc_id: str = Field(description="所属文档 ID")
|
||||
title: str = Field(default="", description="文档标题")
|
||||
section_path: str = Field(default="", description="chunk 所在的章节路径")
|
||||
score: float = Field(description="相关性得分")
|
||||
doc_summary: str = Field(default="", description="L1 文档总结,仅用于上下文标注")
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""检索响应"""
|
||||
|
||||
query: str = Field(description="原始查询文本")
|
||||
hits: list[SearchHit] = Field(default_factory=list, description="命中结果列表")
|
||||
routed_categories: list[str] = Field(default_factory=list, description="query 路由命中的类目")
|
||||
fallback: bool = Field(default=False, description="是否走了全库兜底路径")
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Ollama 客户端
|
||||
|
||||
通过 Ollama HTTP API 调用本地模型进行文本生成。
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import structlog
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
class OllamaClient:
|
||||
"""Ollama HTTP API 客户端"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str | None = None,
|
||||
model: str | None = None,
|
||||
timeout: float = 120.0,
|
||||
) -> None:
|
||||
self.base_url = (base_url or settings.ollama_base_url).rstrip("/")
|
||||
self.model = model or settings.ollama_model
|
||||
self.timeout = timeout
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
"""调用 Ollama 生成文本
|
||||
|
||||
Args:
|
||||
prompt: 输入提示词
|
||||
json_mode: 为 True 时启用 Ollama 原生 JSON 约束输出(payload 加 format: json)
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
"""
|
||||
url = f"{self.base_url}/api/generate"
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
}
|
||||
if json_mode:
|
||||
payload["format"] = "json"
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
resp = await client.post(url, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
result = data.get("response", "")
|
||||
logger.debug("Ollama 生成完成", model=self.model, output_length=len(result))
|
||||
return result
|
||||
|
||||
async def is_available(self) -> bool:
|
||||
"""检查 Ollama 服务是否可用"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.get(f"{self.base_url}/api/tags")
|
||||
return resp.status_code == 200
|
||||
except httpx.HTTPError:
|
||||
return False
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Qdrant 向量数据库服务封装
|
||||
|
||||
分层摘要索引 RAG 的存储层:
|
||||
- doc_l1:文档级总结(dense + sparse)
|
||||
- doc_l2 / doc_l3:大纲节点(dense)
|
||||
- chunks:原文分块(dense + sparse)
|
||||
|
||||
提供集合初始化(幂等)、按层 upsert、dense / hybrid(RRF 融合)检索接口。
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
from qdrant_client import AsyncQdrantClient, models
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 集合名
|
||||
COLLECTION_L1 = "doc_l1"
|
||||
COLLECTION_L2 = "doc_l2"
|
||||
COLLECTION_L3 = "doc_l3"
|
||||
COLLECTION_CHUNKS = "chunks"
|
||||
|
||||
ALL_COLLECTIONS = (COLLECTION_L1, COLLECTION_L2, COLLECTION_L3, COLLECTION_CHUNKS)
|
||||
|
||||
# 命名向量
|
||||
VECTOR_DENSE = "dense"
|
||||
VECTOR_SPARSE = "sparse"
|
||||
|
||||
# 需要 sparse 向量的集合
|
||||
SPARSE_COLLECTIONS = (COLLECTION_L1, COLLECTION_CHUNKS)
|
||||
|
||||
# 需要建立 KEYWORD payload 索引的字段
|
||||
PAYLOAD_INDEX_FIELDS = ("doc_id", "category", "tags", "section_path")
|
||||
|
||||
# upsert_nodes 集合名 -> point id 层级前缀
|
||||
_NODE_ID_PREFIX = {COLLECTION_L2: "l2", COLLECTION_L3: "l3"}
|
||||
|
||||
# 稀疏向量统一表示:(indices, values)
|
||||
SparseVectorTuple = tuple[list[int], list[float]]
|
||||
|
||||
|
||||
def _point_id(key: str) -> str:
|
||||
"""由确定性 key 生成 UUID5 point id,保证重复入库幂等覆盖"""
|
||||
return str(uuid.uuid5(uuid.NAMESPACE_URL, key))
|
||||
|
||||
|
||||
class QdrantService:
|
||||
"""Qdrant 异步客户端封装"""
|
||||
|
||||
def __init__(self, client: AsyncQdrantClient | None = None) -> None:
|
||||
# 允许注入自定义 client(测试可用 location=":memory:" 的本地模式)
|
||||
self._client = client or AsyncQdrantClient(host=settings.qdrant_host, port=settings.qdrant_port)
|
||||
|
||||
@property
|
||||
def client(self) -> AsyncQdrantClient:
|
||||
return self._client
|
||||
|
||||
async def ensure_collections(self) -> None:
|
||||
"""初始化 4 个集合与 payload 索引(幂等,已存在则跳过)"""
|
||||
existing = await self._client.get_collections()
|
||||
existing_names = {c.name for c in existing.collections}
|
||||
|
||||
for collection in ALL_COLLECTIONS:
|
||||
if collection in existing_names:
|
||||
continue
|
||||
sparse_config = None
|
||||
if settings.sparse_enabled and collection in SPARSE_COLLECTIONS:
|
||||
sparse_config = {VECTOR_SPARSE: models.SparseVectorParams(modifier=models.Modifier.IDF)}
|
||||
dense_params = models.VectorParams(size=settings.embedding_dimension, distance=models.Distance.COSINE)
|
||||
await self._client.create_collection(
|
||||
collection_name=collection,
|
||||
vectors_config={VECTOR_DENSE: dense_params},
|
||||
sparse_vectors_config=sparse_config,
|
||||
)
|
||||
logger.info("创建 Qdrant 集合", collection=collection, sparse=sparse_config is not None)
|
||||
|
||||
# payload 索引:先查已有索引,缺失才创建,保证幂等
|
||||
for collection in ALL_COLLECTIONS:
|
||||
info = await self._client.get_collection(collection)
|
||||
indexed = set(info.payload_schema.keys()) if info.payload_schema else set()
|
||||
for field in PAYLOAD_INDEX_FIELDS:
|
||||
if field in indexed:
|
||||
continue
|
||||
await self._client.create_payload_index(
|
||||
collection_name=collection,
|
||||
field_name=field,
|
||||
field_schema=models.PayloadSchemaType.KEYWORD,
|
||||
)
|
||||
logger.debug("payload 索引就绪", collection=collection)
|
||||
|
||||
# ---------- upsert ----------
|
||||
|
||||
async def upsert_l1(
|
||||
self,
|
||||
doc_id: str,
|
||||
title: str,
|
||||
summary: str,
|
||||
category: str,
|
||||
tags: list[str],
|
||||
dense_vector: list[float],
|
||||
sparse_vector: SparseVectorTuple | None = None,
|
||||
) -> None:
|
||||
"""写入 L1 文档总结,payload 含 doc_id/title/category/tags/text(=summary)"""
|
||||
vector: dict[str, Any] = {VECTOR_DENSE: dense_vector}
|
||||
if sparse_vector is not None:
|
||||
vector[VECTOR_SPARSE] = models.SparseVector(indices=sparse_vector[0], values=sparse_vector[1])
|
||||
point = models.PointStruct(
|
||||
id=_point_id(f"{doc_id}:l1"),
|
||||
vector=vector,
|
||||
payload={"doc_id": doc_id, "title": title, "category": category, "tags": tags, "text": summary},
|
||||
)
|
||||
await self._client.upsert(collection_name=COLLECTION_L1, points=[point])
|
||||
|
||||
async def upsert_nodes(self, collection: str, nodes: list[dict[str, Any]]) -> None:
|
||||
"""批量写入 L2/L3 大纲节点
|
||||
|
||||
每个 node 含 doc_id/section_path/text/category/tags/dense_vector。
|
||||
"""
|
||||
if collection not in _NODE_ID_PREFIX:
|
||||
raise ValueError(f"upsert_nodes 仅支持 {COLLECTION_L2}/{COLLECTION_L3},收到: {collection}")
|
||||
prefix = _NODE_ID_PREFIX[collection]
|
||||
points = [
|
||||
models.PointStruct(
|
||||
id=_point_id(f"{node['doc_id']}:{prefix}:{i}"),
|
||||
vector={VECTOR_DENSE: node["dense_vector"]},
|
||||
payload={
|
||||
"doc_id": node["doc_id"],
|
||||
"section_path": node["section_path"],
|
||||
"text": node["text"],
|
||||
"category": node["category"],
|
||||
"tags": node["tags"],
|
||||
},
|
||||
)
|
||||
for i, node in enumerate(nodes)
|
||||
]
|
||||
if points:
|
||||
await self._client.upsert(collection_name=collection, points=points)
|
||||
|
||||
async def upsert_chunks(self, chunks: list[dict[str, Any]]) -> None:
|
||||
"""批量写入原文 chunk
|
||||
|
||||
每个 chunk 含 doc_id/chunk_index/text/section_path/title/category/tags/dense_vector,
|
||||
sparse_vector 可选((indices, values) 形式)。
|
||||
"""
|
||||
points = []
|
||||
for chunk in chunks:
|
||||
vector: dict[str, Any] = {VECTOR_DENSE: chunk["dense_vector"]}
|
||||
sparse = chunk.get("sparse_vector")
|
||||
if sparse is not None:
|
||||
vector[VECTOR_SPARSE] = models.SparseVector(indices=sparse[0], values=sparse[1])
|
||||
points.append(
|
||||
models.PointStruct(
|
||||
id=_point_id(f"{chunk['doc_id']}:chunk:{chunk['chunk_index']}"),
|
||||
vector=vector,
|
||||
payload={
|
||||
"doc_id": chunk["doc_id"],
|
||||
"chunk_index": chunk["chunk_index"],
|
||||
"text": chunk["text"],
|
||||
"section_path": chunk["section_path"],
|
||||
"title": chunk["title"],
|
||||
"category": chunk["category"],
|
||||
"tags": chunk["tags"],
|
||||
# 所属文档 L1 总结,仅作检索结果上下文标注
|
||||
"doc_summary": chunk.get("doc_summary", ""),
|
||||
},
|
||||
)
|
||||
)
|
||||
if points:
|
||||
await self._client.upsert(collection_name=COLLECTION_CHUNKS, points=points)
|
||||
|
||||
# ---------- 查询 ----------
|
||||
|
||||
async def search_dense(
|
||||
self,
|
||||
collection: str,
|
||||
vector: list[float],
|
||||
limit: int,
|
||||
query_filter: models.Filter | None = None,
|
||||
) -> list[models.ScoredPoint]:
|
||||
"""dense 命名向量检索"""
|
||||
resp = await self._client.query_points(
|
||||
collection_name=collection,
|
||||
query=vector,
|
||||
using=VECTOR_DENSE,
|
||||
limit=limit,
|
||||
query_filter=query_filter,
|
||||
)
|
||||
return resp.points
|
||||
|
||||
async def search_hybrid(
|
||||
self,
|
||||
collection: str,
|
||||
dense_vector: list[float],
|
||||
sparse: SparseVectorTuple,
|
||||
limit: int,
|
||||
query_filter: models.Filter | None = None,
|
||||
) -> list[models.ScoredPoint]:
|
||||
"""dense + sparse 两路 prefetch,服务端 RRF 融合"""
|
||||
resp = await self._client.query_points(
|
||||
collection_name=collection,
|
||||
prefetch=[
|
||||
models.Prefetch(query=dense_vector, using=VECTOR_DENSE, limit=limit, filter=query_filter),
|
||||
models.Prefetch(
|
||||
query=models.SparseVector(indices=sparse[0], values=sparse[1]),
|
||||
using=VECTOR_SPARSE,
|
||||
limit=limit,
|
||||
filter=query_filter,
|
||||
),
|
||||
],
|
||||
query=models.FusionQuery(fusion=models.Fusion.RRF),
|
||||
limit=limit,
|
||||
query_filter=query_filter,
|
||||
)
|
||||
return resp.points
|
||||
|
||||
# ---------- 管理操作 ----------
|
||||
|
||||
async def count(self, collection: str) -> int:
|
||||
"""精确统计集合中点的总数"""
|
||||
result = await self._client.count(collection_name=collection, exact=True)
|
||||
return result.count
|
||||
|
||||
async def scroll_l1(self, limit: int = 20, offset: str | None = None) -> tuple[list[dict[str, Any]], str | None]:
|
||||
"""分页浏览 L1 文档列表(不取向量)
|
||||
|
||||
返回 (items, next_offset):item 含 doc_id/title/category/tags/summary(=payload text);
|
||||
next_offset 为下一页游标,无更多数据时为 None。
|
||||
"""
|
||||
records, next_offset = await self._client.scroll(
|
||||
collection_name=COLLECTION_L1,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
with_payload=["doc_id", "title", "category", "tags", "text"],
|
||||
with_vectors=False,
|
||||
)
|
||||
items = []
|
||||
for record in records:
|
||||
payload = record.payload or {}
|
||||
items.append(
|
||||
{
|
||||
"doc_id": payload.get("doc_id", ""),
|
||||
"title": payload.get("title", ""),
|
||||
"category": payload.get("category", ""),
|
||||
"tags": payload.get("tags", []),
|
||||
"summary": payload.get("text", ""),
|
||||
}
|
||||
)
|
||||
return items, str(next_offset) if next_offset is not None else None
|
||||
|
||||
async def get_doc_detail(self, doc_id: str) -> dict[str, Any] | None:
|
||||
"""获取文档详情:L1 记录 + L2/L3 全部节点 + chunks 数量,文档不存在返回 None"""
|
||||
doc_filter = self.build_filter(doc_ids=[doc_id])
|
||||
l1_records, _ = await self._client.scroll(
|
||||
collection_name=COLLECTION_L1,
|
||||
scroll_filter=doc_filter,
|
||||
limit=1,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
)
|
||||
if not l1_records:
|
||||
return None
|
||||
|
||||
l2_nodes = await self._scroll_payloads(COLLECTION_L2, doc_filter)
|
||||
l3_nodes = await self._scroll_payloads(COLLECTION_L3, doc_filter)
|
||||
chunks_count = await self._client.count(
|
||||
collection_name=COLLECTION_CHUNKS,
|
||||
count_filter=doc_filter,
|
||||
exact=True,
|
||||
)
|
||||
return {
|
||||
"l1": l1_records[0].payload or {},
|
||||
"l2_nodes": l2_nodes,
|
||||
"l3_nodes": l3_nodes,
|
||||
"chunks_count": chunks_count.count,
|
||||
}
|
||||
|
||||
async def delete_by_doc_id(self, doc_id: str) -> dict[str, int]:
|
||||
"""删除四层集合中该 doc_id 的所有点,返回各集合删除数量(不存在的 doc_id 全 0,幂等)"""
|
||||
doc_filter = self.build_filter(doc_ids=[doc_id])
|
||||
deleted: dict[str, int] = {}
|
||||
for collection in ALL_COLLECTIONS:
|
||||
# 先记录删除前数量,再按过滤器删除
|
||||
before = await self._client.count(collection_name=collection, count_filter=doc_filter, exact=True)
|
||||
await self._client.delete(
|
||||
collection_name=collection,
|
||||
points_selector=models.FilterSelector(filter=doc_filter),
|
||||
)
|
||||
deleted[collection] = before.count
|
||||
logger.info("按 doc_id 删除文档数据", doc_id=doc_id, deleted=deleted)
|
||||
return deleted
|
||||
|
||||
async def _scroll_payloads(
|
||||
self,
|
||||
collection: str,
|
||||
scroll_filter: models.Filter | None,
|
||||
page_size: int = 256,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""循环 scroll 取出所有匹配点的 payload(含 section_path/text 等全字段)"""
|
||||
payloads: list[dict[str, Any]] = []
|
||||
offset = None
|
||||
while True:
|
||||
records, next_offset = await self._client.scroll(
|
||||
collection_name=collection,
|
||||
scroll_filter=scroll_filter,
|
||||
limit=page_size,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
with_vectors=False,
|
||||
)
|
||||
payloads.extend(record.payload or {} for record in records)
|
||||
if next_offset is None:
|
||||
return payloads
|
||||
offset = next_offset
|
||||
|
||||
# ---------- 过滤构造 ----------
|
||||
|
||||
@staticmethod
|
||||
def build_filter(
|
||||
categories: list[str] | None = None,
|
||||
doc_ids: list[str] | None = None,
|
||||
section_paths: list[str] | None = None,
|
||||
) -> models.Filter | None:
|
||||
"""构造 payload 过滤器
|
||||
|
||||
- categories:主类硬过滤 + 多标签软召回(category 或 tags 命中其一即召回,min_should=1)
|
||||
- doc_ids / section_paths:must 条件 MatchAny
|
||||
- 全空返回 None
|
||||
"""
|
||||
must: list[models.FieldCondition] = []
|
||||
if doc_ids:
|
||||
must.append(models.FieldCondition(key="doc_id", match=models.MatchAny(any=doc_ids)))
|
||||
if section_paths:
|
||||
must.append(models.FieldCondition(key="section_path", match=models.MatchAny(any=section_paths)))
|
||||
|
||||
min_should = None
|
||||
if categories:
|
||||
min_should = models.MinShould(
|
||||
conditions=[
|
||||
models.FieldCondition(key="category", match=models.MatchAny(any=categories)),
|
||||
models.FieldCondition(key="tags", match=models.MatchAny(any=categories)),
|
||||
],
|
||||
min_count=1,
|
||||
)
|
||||
|
||||
if not must and min_should is None:
|
||||
return None
|
||||
return models.Filter(must=must or None, min_should=min_should)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Redis 缓存服务
|
||||
|
||||
为检索结果与 query 解析结果提供 JSON 缓存。所有操作容错:
|
||||
Redis 不可用或缓存数据异常时降级为未命中 / 写入失败,绝不影响主流程。
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
from redis import asyncio as redis_async
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
class RedisCache:
|
||||
"""Redis JSON 缓存客户端(懒连接,全操作容错)"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client: redis_async.Redis | None = None
|
||||
|
||||
def _get_client(self) -> redis_async.Redis:
|
||||
"""懒创建 Redis 客户端(TCP 连接在首次执行命令时才建立)"""
|
||||
if self._client is None:
|
||||
self._client = redis_async.from_url(settings.redis_url, decode_responses=True)
|
||||
return self._client
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
"""读取缓存并反序列化为 dict
|
||||
|
||||
任何异常(连接失败 / JSON 解析失败)或值非 dict 时都降级为未命中,返回 None。
|
||||
"""
|
||||
try:
|
||||
raw = await self._get_client().get(key)
|
||||
if raw is None:
|
||||
return None
|
||||
data = json.loads(raw)
|
||||
return data if isinstance(data, dict) else None
|
||||
except Exception:
|
||||
logger.warning("Redis 读取缓存失败,降级为未命中", key=key, exc_info=True)
|
||||
return None
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
"""写入缓存(SETEX),ttl 缺省取 settings.cache_ttl
|
||||
|
||||
异常时记录日志并返回 False,不影响主流程。
|
||||
"""
|
||||
try:
|
||||
await self._get_client().setex(
|
||||
key, ttl if ttl is not None else settings.cache_ttl, json.dumps(value, ensure_ascii=False)
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
logger.warning("Redis 写入缓存失败", key=key, exc_info=True)
|
||||
return False
|
||||
|
||||
async def close(self) -> None:
|
||||
"""关闭底层连接"""
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
# 模块级懒加载单例,避免每请求重建客户端
|
||||
_cache: RedisCache | None = None
|
||||
|
||||
|
||||
def get_cache() -> RedisCache:
|
||||
"""获取全局 RedisCache 单例"""
|
||||
global _cache
|
||||
if _cache is None:
|
||||
_cache = RedisCache()
|
||||
return _cache
|
||||
@@ -0,0 +1,692 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>知识库管理后台</title>
|
||||
<style>
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: -apple-system, "PingFang SC", "Microsoft YaHei", sans-serif;
|
||||
background: #f5f6f8;
|
||||
color: #24292f;
|
||||
}
|
||||
header {
|
||||
background: #1f2937;
|
||||
color: #fff;
|
||||
padding: 12px 24px;
|
||||
}
|
||||
header h1 { font-size: 18px; margin: 0 0 10px; }
|
||||
nav button {
|
||||
background: transparent;
|
||||
border: 1px solid #4b5563;
|
||||
color: #d1d5db;
|
||||
padding: 6px 14px;
|
||||
margin-right: 8px;
|
||||
border-radius: 4px;
|
||||
cursor: pointer;
|
||||
font-size: 14px;
|
||||
}
|
||||
nav button.active { background: #3b82f6; border-color: #3b82f6; color: #fff; }
|
||||
main { padding: 20px 24px; max-width: 1080px; margin: 0 auto; }
|
||||
section { background: #fff; border: 1px solid #e5e7eb; border-radius: 6px; padding: 16px 20px; }
|
||||
section h2 { font-size: 16px; margin: 0 0 14px; border-bottom: 1px solid #e5e7eb; padding-bottom: 8px; }
|
||||
.hidden { display: none; }
|
||||
.error-bar {
|
||||
background: #fef2f2;
|
||||
border: 1px solid #fca5a5;
|
||||
color: #b91c1c;
|
||||
border-radius: 4px;
|
||||
padding: 8px 12px;
|
||||
margin-bottom: 12px;
|
||||
font-size: 13px;
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
.cards { display: flex; flex-wrap: wrap; gap: 12px; margin-bottom: 16px; }
|
||||
.card {
|
||||
flex: 1 1 140px;
|
||||
border: 1px solid #e5e7eb;
|
||||
border-radius: 6px;
|
||||
padding: 12px;
|
||||
background: #fafafa;
|
||||
}
|
||||
.card .num { font-size: 24px; font-weight: 600; }
|
||||
.card .label { font-size: 12px; color: #6b7280; margin-top: 4px; }
|
||||
.bar-row { display: flex; align-items: center; margin-bottom: 6px; font-size: 13px; }
|
||||
.bar-row .name { width: 160px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.bar-row .bar-track { flex: 1; background: #f3f4f6; border-radius: 3px; height: 16px; margin: 0 8px; }
|
||||
.bar-row .bar-fill { background: #3b82f6; height: 100%; border-radius: 3px; min-width: 2px; }
|
||||
.bar-row .count { width: 48px; text-align: right; color: #374151; }
|
||||
table { width: 100%; border-collapse: collapse; font-size: 13px; }
|
||||
th, td { border-bottom: 1px solid #e5e7eb; padding: 8px; text-align: left; vertical-align: top; }
|
||||
th { background: #f9fafb; color: #374151; }
|
||||
button.action {
|
||||
border: 1px solid #d1d5db;
|
||||
background: #fff;
|
||||
border-radius: 4px;
|
||||
padding: 4px 10px;
|
||||
margin-right: 6px;
|
||||
cursor: pointer;
|
||||
font-size: 12px;
|
||||
}
|
||||
button.action:hover { background: #f3f4f6; }
|
||||
button.danger { color: #b91c1c; border-color: #fca5a5; }
|
||||
button.primary {
|
||||
background: #3b82f6; color: #fff; border: none;
|
||||
border-radius: 4px; padding: 8px 18px; cursor: pointer; font-size: 14px;
|
||||
}
|
||||
button.primary:disabled { background: #9ca3af; cursor: not-allowed; }
|
||||
form .field { margin-bottom: 12px; }
|
||||
form label { display: block; font-size: 13px; margin-bottom: 4px; color: #374151; }
|
||||
input[type="text"], input[type="number"], textarea {
|
||||
width: 100%; border: 1px solid #d1d5db; border-radius: 4px;
|
||||
padding: 8px; font-size: 13px; font-family: inherit;
|
||||
}
|
||||
textarea { min-height: 160px; resize: vertical; }
|
||||
.result-box {
|
||||
margin-top: 14px; border: 1px solid #e5e7eb; border-radius: 6px;
|
||||
background: #fafafa; padding: 12px; font-size: 13px;
|
||||
}
|
||||
.result-box .kv { margin-bottom: 6px; }
|
||||
.result-box .k { color: #6b7280; display: inline-block; min-width: 110px; }
|
||||
.detail-panel {
|
||||
margin-top: 14px; border: 1px solid #d1d5db; border-radius: 6px; padding: 12px;
|
||||
}
|
||||
.detail-panel h3 { margin: 0 0 8px; font-size: 14px; }
|
||||
.detail-panel pre {
|
||||
background: #f9fafb; border: 1px solid #e5e7eb; border-radius: 4px;
|
||||
padding: 8px; white-space: pre-wrap; word-break: break-word;
|
||||
font-size: 12px; max-height: 260px; overflow: auto;
|
||||
}
|
||||
.node-item { border-left: 3px solid #93c5fd; padding: 4px 8px; margin-bottom: 6px; background: #f8fafc; }
|
||||
.node-item .path { font-size: 12px; color: #6b7280; }
|
||||
.hit-item { border: 1px solid #e5e7eb; border-radius: 6px; padding: 10px; margin-bottom: 10px; }
|
||||
.hit-item .meta { font-size: 12px; color: #6b7280; margin-bottom: 6px; }
|
||||
.hit-item .snippet { font-size: 13px; white-space: pre-wrap; word-break: break-word; }
|
||||
.tag { display: inline-block; background: #eff6ff; color: #1d4ed8; border-radius: 3px; padding: 1px 6px; font-size: 12px; margin-right: 4px; }
|
||||
.fallback-flag { color: #b45309; background: #fffbeb; border: 1px solid #fcd34d; border-radius: 4px; padding: 6px 10px; font-size: 13px; margin-bottom: 10px; }
|
||||
.toolbar { margin: 12px 0; }
|
||||
.muted { color: #6b7280; font-size: 12px; }
|
||||
.status-badge {
|
||||
display: inline-block; border-radius: 3px; padding: 1px 8px;
|
||||
font-size: 12px; margin-left: 6px; border: 1px solid transparent;
|
||||
}
|
||||
.status-running { background: #eff6ff; color: #1d4ed8; border-color: #93c5fd; }
|
||||
.status-done { background: #f0fdf4; color: #15803d; border-color: #86efac; }
|
||||
.status-failed { background: #fef2f2; color: #b91c1c; border-color: #fca5a5; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<h1>知识库管理后台</h1>
|
||||
<nav id="nav">
|
||||
<button type="button" data-target="section-overview" class="active">概览</button>
|
||||
<button type="button" data-target="section-docs">文档管理</button>
|
||||
<button type="button" data-target="section-ingest">文档入库</button>
|
||||
<button type="button" data-target="section-search">检索测试台</button>
|
||||
<button type="button" data-target="section-categories">类目列表</button>
|
||||
</nav>
|
||||
</header>
|
||||
<main>
|
||||
<section id="section-overview">
|
||||
<h2>概览</h2>
|
||||
<div class="error-bar hidden" id="error-overview"></div>
|
||||
<div class="toolbar">
|
||||
<button type="button" class="action" id="btn-refresh-overview">刷新</button>
|
||||
</div>
|
||||
<div class="cards" id="overview-cards"></div>
|
||||
<h3 style="font-size:14px;">类目分布</h3>
|
||||
<div id="overview-categories"></div>
|
||||
</section>
|
||||
|
||||
<section id="section-docs" class="hidden">
|
||||
<h2>文档管理</h2>
|
||||
<div class="error-bar hidden" id="error-docs"></div>
|
||||
<table>
|
||||
<thead>
|
||||
<tr><th>标题</th><th>类目</th><th>标签</th><th>L1 摘要</th><th>操作</th></tr>
|
||||
</thead>
|
||||
<tbody id="docs-tbody"></tbody>
|
||||
</table>
|
||||
<div class="toolbar">
|
||||
<button type="button" class="action" id="btn-load-more">加载更多</button>
|
||||
<span class="muted" id="docs-status"></span>
|
||||
</div>
|
||||
<div class="detail-panel hidden" id="doc-detail"></div>
|
||||
</section>
|
||||
|
||||
<section id="section-ingest" class="hidden">
|
||||
<h2>文档入库</h2>
|
||||
<div class="error-bar hidden" id="error-ingest"></div>
|
||||
<form id="ingest-form">
|
||||
<div class="field">
|
||||
<label for="ingest-title">标题</label>
|
||||
<input type="text" id="ingest-title" name="title" required>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="ingest-source">来源</label>
|
||||
<input type="text" id="ingest-source" name="source" placeholder="例如:manual / web / file">
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="ingest-text">正文</label>
|
||||
<textarea id="ingest-text" name="text" required></textarea>
|
||||
</div>
|
||||
<button type="submit" class="primary" id="btn-ingest-submit">提交入库</button>
|
||||
</form>
|
||||
<div class="result-box hidden" id="ingest-result"></div>
|
||||
</section>
|
||||
|
||||
<section id="section-search" class="hidden">
|
||||
<h2>检索测试台</h2>
|
||||
<div class="error-bar hidden" id="error-search"></div>
|
||||
<form id="search-form">
|
||||
<div class="field">
|
||||
<label for="search-query">查询语句</label>
|
||||
<input type="text" id="search-query" name="query" required>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="search-topk">top_k</label>
|
||||
<input type="number" id="search-topk" name="top_k" value="5" min="1" max="50">
|
||||
</div>
|
||||
<button type="submit" class="primary" id="btn-search-submit">检索</button>
|
||||
</form>
|
||||
<div class="result-box hidden" id="search-result"></div>
|
||||
</section>
|
||||
|
||||
<section id="section-categories" class="hidden">
|
||||
<h2>类目列表</h2>
|
||||
<div class="error-bar hidden" id="error-categories"></div>
|
||||
<table>
|
||||
<thead>
|
||||
<tr><th>名称</th><th>描述</th></tr>
|
||||
</thead>
|
||||
<tbody id="categories-tbody"></tbody>
|
||||
</table>
|
||||
</section>
|
||||
</main>
|
||||
|
||||
<script>
|
||||
"use strict";
|
||||
|
||||
/* ---------- 通用工具 ---------- */
|
||||
|
||||
function el(tag, text, className) {
|
||||
var node = document.createElement(tag);
|
||||
if (className) { node.className = className; }
|
||||
if (text !== undefined && text !== null) { node.textContent = String(text); }
|
||||
return node;
|
||||
}
|
||||
|
||||
function clearChildren(node) {
|
||||
while (node.firstChild) { node.removeChild(node.firstChild); }
|
||||
}
|
||||
|
||||
function showError(boxId, err) {
|
||||
var box = document.getElementById(boxId);
|
||||
clearChildren(box);
|
||||
var code = (err && err.code !== undefined) ? err.code : "网络错误";
|
||||
var message = (err && err.message) ? err.message : String(err);
|
||||
box.textContent = "错误 [" + code + "] " + message;
|
||||
box.classList.remove("hidden");
|
||||
}
|
||||
|
||||
function hideError(boxId) {
|
||||
document.getElementById(boxId).classList.add("hidden");
|
||||
}
|
||||
|
||||
/* 统一 API 封装:code !== 0 抛错,网络错误同样捕获 */
|
||||
function api(path, options) {
|
||||
return fetch(path, options).then(function (resp) {
|
||||
return resp.json().catch(function () {
|
||||
throw { code: "HTTP " + resp.status, message: "响应解析失败" };
|
||||
});
|
||||
}).then(function (body) {
|
||||
if (body.code !== 0) {
|
||||
throw { code: body.code, message: body.message };
|
||||
}
|
||||
return body.data;
|
||||
}).catch(function (err) {
|
||||
if (err && err.code !== undefined) { throw err; }
|
||||
throw { code: "NETWORK", message: "网络请求失败: " + (err && err.message ? err.message : err) };
|
||||
});
|
||||
}
|
||||
|
||||
function truncate(text, maxLen) {
|
||||
if (!text) { return ""; }
|
||||
var s = String(text);
|
||||
return s.length > maxLen ? s.slice(0, maxLen) + "…" : s;
|
||||
}
|
||||
|
||||
/* ---------- 导航切换 ---------- */
|
||||
|
||||
var loadedOnce = {};
|
||||
|
||||
function activateSection(targetId) {
|
||||
var sections = document.querySelectorAll("main section");
|
||||
sections.forEach(function (sec) {
|
||||
sec.classList.toggle("hidden", sec.id !== targetId);
|
||||
});
|
||||
var buttons = document.querySelectorAll("#nav button");
|
||||
buttons.forEach(function (btn) {
|
||||
btn.classList.toggle("active", btn.getAttribute("data-target") === targetId);
|
||||
});
|
||||
if (!loadedOnce[targetId]) {
|
||||
loadedOnce[targetId] = true;
|
||||
if (targetId === "section-overview") { loadOverview(); }
|
||||
if (targetId === "section-docs") { loadDocuments(true); }
|
||||
if (targetId === "section-categories") { loadCategories(); }
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById("nav").addEventListener("click", function (event) {
|
||||
var btn = event.target.closest("button[data-target]");
|
||||
if (btn) { activateSection(btn.getAttribute("data-target")); }
|
||||
});
|
||||
|
||||
/* ---------- 1. 概览 ---------- */
|
||||
|
||||
function loadOverview() {
|
||||
hideError("error-overview");
|
||||
api("/api/v1/knowledge/stats").then(function (data) {
|
||||
renderOverviewCards(data);
|
||||
renderOverviewCategories(data.categories || {});
|
||||
}).catch(function (err) { showError("error-overview", err); });
|
||||
}
|
||||
|
||||
function renderOverviewCards(data) {
|
||||
var container = document.getElementById("overview-cards");
|
||||
clearChildren(container);
|
||||
var collections = data.collections || {};
|
||||
var cards = [
|
||||
["L1 文档总结", collections.doc_l1],
|
||||
["L2 大纲节点", collections.doc_l2],
|
||||
["L3 内容大纲", collections.doc_l3],
|
||||
["Chunks", collections.chunks],
|
||||
["文档总数", data.documents_total],
|
||||
["未分类文档", data.uncategorized_count]
|
||||
];
|
||||
cards.forEach(function (pair) {
|
||||
var card = el("div", null, "card");
|
||||
card.appendChild(el("div", pair[1] === undefined ? 0 : pair[1], "num"));
|
||||
card.appendChild(el("div", pair[0], "label"));
|
||||
container.appendChild(card);
|
||||
});
|
||||
}
|
||||
|
||||
function renderOverviewCategories(categories) {
|
||||
var container = document.getElementById("overview-categories");
|
||||
clearChildren(container);
|
||||
var names = Object.keys(categories);
|
||||
if (names.length === 0) {
|
||||
container.appendChild(el("div", "暂无数据", "muted"));
|
||||
return;
|
||||
}
|
||||
var max = 1;
|
||||
names.forEach(function (name) { max = Math.max(max, categories[name]); });
|
||||
names.sort(function (a, b) { return categories[b] - categories[a]; });
|
||||
names.forEach(function (name) {
|
||||
var count = categories[name];
|
||||
var row = el("div", null, "bar-row");
|
||||
row.appendChild(el("span", name, "name"));
|
||||
var track = el("div", null, "bar-track");
|
||||
var fill = el("div", null, "bar-fill");
|
||||
fill.style.width = Math.round((count / max) * 100) + "%";
|
||||
track.appendChild(fill);
|
||||
row.appendChild(track);
|
||||
row.appendChild(el("span", count, "count"));
|
||||
container.appendChild(row);
|
||||
});
|
||||
}
|
||||
|
||||
document.getElementById("btn-refresh-overview").addEventListener("click", loadOverview);
|
||||
|
||||
/* ---------- 2. 文档管理 ---------- */
|
||||
|
||||
var docsOffset = null;
|
||||
|
||||
function loadDocuments(reset) {
|
||||
hideError("error-docs");
|
||||
if (reset) {
|
||||
docsOffset = null;
|
||||
clearChildren(document.getElementById("docs-tbody"));
|
||||
}
|
||||
var path = "/api/v1/documents?limit=20";
|
||||
if (docsOffset !== null) { path += "&offset=" + encodeURIComponent(docsOffset); }
|
||||
api(path).then(function (data) {
|
||||
renderDocuments(data.items || []);
|
||||
docsOffset = data.next_offset;
|
||||
document.getElementById("btn-load-more").disabled = docsOffset === null;
|
||||
document.getElementById("docs-status").textContent =
|
||||
docsOffset === null ? "已加载全部" : "还有更多";
|
||||
}).catch(function (err) { showError("error-docs", err); });
|
||||
}
|
||||
|
||||
function renderDocuments(items) {
|
||||
var tbody = document.getElementById("docs-tbody");
|
||||
items.forEach(function (item) {
|
||||
var tr = el("tr");
|
||||
tr.appendChild(el("td", item.title || "(无标题)"));
|
||||
tr.appendChild(el("td", item.category || ""));
|
||||
|
||||
var tagsTd = el("td");
|
||||
(item.tags || []).forEach(function (tag) { tagsTd.appendChild(el("span", tag, "tag")); });
|
||||
tr.appendChild(tagsTd);
|
||||
|
||||
tr.appendChild(el("td", truncate(item.summary, 80)));
|
||||
|
||||
var opsTd = el("td");
|
||||
var detailBtn = el("button", "详情", "action");
|
||||
detailBtn.type = "button";
|
||||
detailBtn.addEventListener("click", function () { loadDocDetail(item.doc_id); });
|
||||
var deleteBtn = el("button", "删除", "action danger");
|
||||
deleteBtn.type = "button";
|
||||
deleteBtn.addEventListener("click", function () { deleteDocument(item.doc_id, item.title); });
|
||||
opsTd.appendChild(detailBtn);
|
||||
opsTd.appendChild(deleteBtn);
|
||||
tr.appendChild(opsTd);
|
||||
|
||||
tbody.appendChild(tr);
|
||||
});
|
||||
}
|
||||
|
||||
document.getElementById("btn-load-more").addEventListener("click", function () {
|
||||
loadDocuments(false);
|
||||
});
|
||||
|
||||
function loadDocDetail(docId) {
|
||||
hideError("error-docs");
|
||||
api("/api/v1/documents/" + encodeURIComponent(docId)).then(function (data) {
|
||||
renderDocDetail(data);
|
||||
}).catch(function (err) { showError("error-docs", err); });
|
||||
}
|
||||
|
||||
function renderDocDetail(data) {
|
||||
var panel = document.getElementById("doc-detail");
|
||||
clearChildren(panel);
|
||||
panel.classList.remove("hidden");
|
||||
|
||||
var l1 = data.l1 || {};
|
||||
panel.appendChild(el("h3", "文档详情:" + (l1.title || data.l1 && l1.doc_id || "")));
|
||||
|
||||
var meta = el("div", null, "kv");
|
||||
meta.appendChild(el("span", "doc_id", "k"));
|
||||
meta.appendChild(el("span", l1.doc_id || ""));
|
||||
panel.appendChild(meta);
|
||||
|
||||
var meta2 = el("div", null, "kv");
|
||||
meta2.appendChild(el("span", "类目 / 标签", "k"));
|
||||
meta2.appendChild(el("span", (l1.category || "") + " / " + (l1.tags || []).join(", ")));
|
||||
panel.appendChild(meta2);
|
||||
|
||||
var meta3 = el("div", null, "kv");
|
||||
meta3.appendChild(el("span", "chunks_count", "k"));
|
||||
meta3.appendChild(el("span", data.chunks_count));
|
||||
panel.appendChild(meta3);
|
||||
|
||||
panel.appendChild(el("h3", "L1 全文"));
|
||||
panel.appendChild(el("pre", l1.text || ""));
|
||||
|
||||
panel.appendChild(el("h3", "L2 节点(" + (data.l2_nodes || []).length + ")"));
|
||||
renderNodes(panel, data.l2_nodes || []);
|
||||
|
||||
panel.appendChild(el("h3", "L3 节点(" + (data.l3_nodes || []).length + ")"));
|
||||
renderNodes(panel, data.l3_nodes || []);
|
||||
}
|
||||
|
||||
function renderNodes(panel, nodes) {
|
||||
if (nodes.length === 0) {
|
||||
panel.appendChild(el("div", "无", "muted"));
|
||||
return;
|
||||
}
|
||||
nodes.forEach(function (node) {
|
||||
var item = el("div", null, "node-item");
|
||||
item.appendChild(el("div", node.section_path || "", "path"));
|
||||
item.appendChild(el("div", node.text || ""));
|
||||
panel.appendChild(item);
|
||||
});
|
||||
}
|
||||
|
||||
function deleteDocument(docId, title) {
|
||||
if (!confirm("确定删除文档「" + (title || docId) + "」(" + docId + ") 吗?该操作将删除四层集合中的全部数据,不可恢复。")) {
|
||||
return;
|
||||
}
|
||||
hideError("error-docs");
|
||||
api("/api/v1/documents/" + encodeURIComponent(docId), { method: "DELETE" }).then(function (data) {
|
||||
document.getElementById("doc-detail").classList.add("hidden");
|
||||
loadDocuments(true);
|
||||
loadOverview();
|
||||
var status = document.getElementById("docs-status");
|
||||
status.textContent = "已删除 " + (data.deleted_total || 0) + " 条数据";
|
||||
}).catch(function (err) { showError("error-docs", err); });
|
||||
}
|
||||
|
||||
/* ---------- 3. 文档入库(异步任务轮询) ---------- */
|
||||
|
||||
var INGEST_STATUS_TEXT = {
|
||||
pending: "排队中",
|
||||
summarizing: "总结中",
|
||||
classifying: "分类中",
|
||||
embedding: "向量化中",
|
||||
writing: "写入中",
|
||||
done: "完成",
|
||||
failed: "失败"
|
||||
};
|
||||
|
||||
var INGEST_POLL_INTERVAL_MS = 2000;
|
||||
var INGEST_POLL_MAX_ATTEMPTS = 150; /* 150 次 × 2s = 5 分钟超时 */
|
||||
|
||||
var ingestPollTimer = null;
|
||||
|
||||
function stopIngestPolling() {
|
||||
if (ingestPollTimer !== null) {
|
||||
clearInterval(ingestPollTimer);
|
||||
ingestPollTimer = null;
|
||||
}
|
||||
}
|
||||
|
||||
function restoreIngestButton() {
|
||||
var submitBtn = document.getElementById("btn-ingest-submit");
|
||||
submitBtn.disabled = false;
|
||||
submitBtn.textContent = "提交入库";
|
||||
}
|
||||
|
||||
function makeStatusBadge(status) {
|
||||
var cls = "status-running";
|
||||
if (status === "done") { cls = "status-done"; }
|
||||
if (status === "failed") { cls = "status-failed"; }
|
||||
var badge = el("span", INGEST_STATUS_TEXT[status] || status, "status-badge " + cls);
|
||||
badge.id = "ingest-status-badge";
|
||||
return badge;
|
||||
}
|
||||
|
||||
function updateIngestStatus(status) {
|
||||
var old = document.getElementById("ingest-status-badge");
|
||||
if (old && old.parentNode) {
|
||||
old.parentNode.replaceChild(makeStatusBadge(status), old);
|
||||
}
|
||||
}
|
||||
|
||||
function ingestKv(box, label, value) {
|
||||
var row = el("div", null, "kv");
|
||||
row.appendChild(el("span", label, "k"));
|
||||
row.appendChild(el("span", value));
|
||||
box.appendChild(row);
|
||||
}
|
||||
|
||||
function renderIngestHeader(taskId, status) {
|
||||
var box = document.getElementById("ingest-result");
|
||||
clearChildren(box);
|
||||
box.classList.remove("hidden");
|
||||
var row = el("div", null, "kv");
|
||||
row.appendChild(el("span", "任务已提交:" + taskId));
|
||||
row.appendChild(makeStatusBadge(status));
|
||||
box.appendChild(row);
|
||||
return box;
|
||||
}
|
||||
|
||||
function renderIngestDone(task) {
|
||||
var box = renderIngestHeader(task.task_id, "done");
|
||||
var result = task.result || {};
|
||||
var summary = result.summary || {};
|
||||
ingestKv(box, "document_id", result.document_id);
|
||||
ingestKv(box, "类目", (result.category || "") + "(置信度 " + result.category_confidence + ")");
|
||||
ingestKv(box, "标签", (result.tags || []).join(", "));
|
||||
ingestKv(box, "总结层级 level", summary.level);
|
||||
ingestKv(box, "写入集合", result.collection);
|
||||
ingestKv(box, "chunks_count", result.chunks_count);
|
||||
box.appendChild(el("h3", "L1 总结"));
|
||||
box.appendChild(el("pre", summary.l1_summary || ""));
|
||||
}
|
||||
|
||||
function renderIngestFailed(task) {
|
||||
var box = renderIngestHeader(task.task_id, "failed");
|
||||
var error = task.error || {};
|
||||
ingestKv(box, "失败阶段", error.stage || "");
|
||||
ingestKv(box, "错误信息", error.message || "");
|
||||
var partial = error.partial_summary;
|
||||
if (partial && partial.l1_summary) {
|
||||
box.appendChild(el("div", "已产出总结保留:任务失败前已生成 L1 摘要", "fallback-flag"));
|
||||
box.appendChild(el("h3", "L1 总结"));
|
||||
box.appendChild(el("pre", partial.l1_summary));
|
||||
}
|
||||
}
|
||||
|
||||
function renderIngestTimeout(taskId) {
|
||||
var box = document.getElementById("ingest-result");
|
||||
box.appendChild(el("div", "任务仍在进行,可稍后凭 task_id 查询:" + taskId, "fallback-flag"));
|
||||
}
|
||||
|
||||
function startIngestPolling(taskId) {
|
||||
var attempts = 0;
|
||||
document.getElementById("btn-ingest-submit").textContent = "入库中…";
|
||||
ingestPollTimer = setInterval(function () {
|
||||
attempts += 1;
|
||||
if (attempts > INGEST_POLL_MAX_ATTEMPTS) {
|
||||
stopIngestPolling();
|
||||
renderIngestTimeout(taskId);
|
||||
restoreIngestButton();
|
||||
return;
|
||||
}
|
||||
api("/api/v1/documents/tasks/" + encodeURIComponent(taskId)).then(function (task) {
|
||||
updateIngestStatus(task.status);
|
||||
if (task.status === "done") {
|
||||
stopIngestPolling();
|
||||
renderIngestDone(task);
|
||||
restoreIngestButton();
|
||||
} else if (task.status === "failed") {
|
||||
stopIngestPolling();
|
||||
renderIngestFailed(task);
|
||||
restoreIngestButton();
|
||||
}
|
||||
}).catch(function (err) {
|
||||
stopIngestPolling();
|
||||
showError("error-ingest", err);
|
||||
restoreIngestButton();
|
||||
});
|
||||
}, INGEST_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
document.getElementById("ingest-form").addEventListener("submit", function (event) {
|
||||
event.preventDefault();
|
||||
hideError("error-ingest");
|
||||
stopIngestPolling();
|
||||
var submitBtn = document.getElementById("btn-ingest-submit");
|
||||
submitBtn.disabled = true;
|
||||
submitBtn.textContent = "提交中…";
|
||||
var payload = {
|
||||
title: document.getElementById("ingest-title").value,
|
||||
source: document.getElementById("ingest-source").value,
|
||||
text: document.getElementById("ingest-text").value
|
||||
};
|
||||
api("/api/v1/documents", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload)
|
||||
}).then(function (data) {
|
||||
renderIngestHeader(data.task_id, data.status || "pending");
|
||||
startIngestPolling(data.task_id);
|
||||
}).catch(function (err) {
|
||||
showError("error-ingest", err);
|
||||
restoreIngestButton();
|
||||
});
|
||||
});
|
||||
|
||||
/* ---------- 4. 检索测试台 ---------- */
|
||||
|
||||
document.getElementById("search-form").addEventListener("submit", function (event) {
|
||||
event.preventDefault();
|
||||
hideError("error-search");
|
||||
var submitBtn = document.getElementById("btn-search-submit");
|
||||
submitBtn.disabled = true;
|
||||
var payload = {
|
||||
query: document.getElementById("search-query").value,
|
||||
top_k: parseInt(document.getElementById("search-topk").value, 10) || 5
|
||||
};
|
||||
api("/api/v1/search", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload)
|
||||
}).then(function (data) {
|
||||
renderSearchResult(data);
|
||||
}).catch(function (err) {
|
||||
showError("error-search", err);
|
||||
}).finally(function () {
|
||||
submitBtn.disabled = false;
|
||||
});
|
||||
});
|
||||
|
||||
function renderSearchResult(data) {
|
||||
var box = document.getElementById("search-result");
|
||||
clearChildren(box);
|
||||
box.classList.remove("hidden");
|
||||
|
||||
var routed = el("div", null, "kv");
|
||||
routed.appendChild(el("span", "routed_categories", "k"));
|
||||
routed.appendChild(el("span", (data.routed_categories || []).join(", ") || "(无)"));
|
||||
box.appendChild(routed);
|
||||
|
||||
if (data.fallback) {
|
||||
box.appendChild(el("div", "fallback:本次检索触发了回退策略(路由类目无命中,已降级全局检索)", "fallback-flag"));
|
||||
}
|
||||
|
||||
var hits = data.hits || [];
|
||||
box.appendChild(el("div", "命中 " + hits.length + " 条", "muted"));
|
||||
hits.forEach(function (hit) {
|
||||
var item = el("div", null, "hit-item");
|
||||
var meta = el("div", null, "meta");
|
||||
meta.textContent = "score=" + hit.score + " | doc_id=" + hit.doc_id +
|
||||
" | 标题=" + (hit.title || "") + " | section=" + (hit.section_path || "");
|
||||
item.appendChild(meta);
|
||||
item.appendChild(el("div", hit.text || "", "snippet"));
|
||||
if (hit.doc_summary) {
|
||||
var summaryRow = el("div", null, "meta");
|
||||
summaryRow.textContent = "文档摘要:" + hit.doc_summary;
|
||||
item.appendChild(summaryRow);
|
||||
}
|
||||
box.appendChild(item);
|
||||
});
|
||||
}
|
||||
|
||||
/* ---------- 5. 类目列表 ---------- */
|
||||
|
||||
function loadCategories() {
|
||||
hideError("error-categories");
|
||||
api("/api/v1/knowledge/categories").then(function (data) {
|
||||
var tbody = document.getElementById("categories-tbody");
|
||||
clearChildren(tbody);
|
||||
(data.categories || []).forEach(function (cat) {
|
||||
var tr = el("tr");
|
||||
tr.appendChild(el("td", cat.name));
|
||||
tr.appendChild(el("td", cat.description || ""));
|
||||
tbody.appendChild(tr);
|
||||
});
|
||||
}).catch(function (err) { showError("error-categories", err); });
|
||||
}
|
||||
|
||||
/* ---------- 初始化 ---------- */
|
||||
|
||||
activateSection("section-overview");
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,78 @@
|
||||
services:
|
||||
app:
|
||||
build: .
|
||||
container_name: qmdsearch-app
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${APP_PORT:-8000}:8000"
|
||||
environment:
|
||||
- QDRANT_HOST=qdrant
|
||||
- QDRANT_PORT=6333
|
||||
- REDIS_URL=redis://redis:6379/0
|
||||
- OLLAMA_BASE_URL=http://ollama:11434
|
||||
env_file:
|
||||
- .env
|
||||
depends_on:
|
||||
qdrant:
|
||||
condition: service_healthy
|
||||
redis:
|
||||
condition: service_healthy
|
||||
ollama:
|
||||
condition: service_started
|
||||
volumes:
|
||||
- ${NAS_DATA_DIR:-./data}/logs:/app/logs
|
||||
networks:
|
||||
- qmdsearch
|
||||
|
||||
qdrant:
|
||||
image: qdrant/qdrant:latest
|
||||
container_name: qmdsearch-qdrant
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${QDRANT_PORT:-6333}:6333"
|
||||
- "${QDRANT_DASHBOARD_PORT:-6334}:6334"
|
||||
volumes:
|
||||
- ${NAS_DATA_DIR:-./data}/qdrant:/qdrant/storage
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:6333/healthz"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
networks:
|
||||
- qmdsearch
|
||||
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
container_name: qmdsearch-redis
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${REDIS_PORT:-6379}:6379"
|
||||
volumes:
|
||||
- ${NAS_DATA_DIR:-./data}/redis:/data
|
||||
healthcheck:
|
||||
test: ["CMD", "redis-cli", "ping"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
networks:
|
||||
- qmdsearch
|
||||
|
||||
ollama:
|
||||
image: ollama/ollama:latest
|
||||
container_name: qmdsearch-ollama
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "${OLLAMA_PORT:-11434}:11434"
|
||||
volumes:
|
||||
- ${NAS_DATA_DIR:-./data}/ollama:/root/.ollama
|
||||
# 首次启动后需手动拉取模型:
|
||||
# docker exec qmdsearch-ollama ollama pull qwen2.5:1.5b
|
||||
# 或取消下方 entrypoint 注释以自动拉取(需等待下载完成)
|
||||
# entrypoint: /bin/bash
|
||||
# command: -c "ollama serve & sleep 5 && ollama pull qwen2.5:1.5b && wait"
|
||||
networks:
|
||||
- qmdsearch
|
||||
|
||||
networks:
|
||||
qmdsearch:
|
||||
driver: bridge
|
||||
@@ -0,0 +1,41 @@
|
||||
[project]
|
||||
name = "qmdsearch"
|
||||
version = "0.1.0"
|
||||
description = "AI Agent 分层信息检索服务"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"fastapi>=0.115.0",
|
||||
"uvicorn[standard]>=0.34.0",
|
||||
"pydantic>=2.10.0",
|
||||
"pydantic-settings>=2.7.0",
|
||||
"qdrant-client>=1.12.0",
|
||||
"redis>=5.2.0",
|
||||
"httpx>=0.28.0",
|
||||
"structlog>=24.4.0",
|
||||
"openai>=1.58.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.3.0",
|
||||
"pytest-asyncio>=0.24.0",
|
||||
"ruff>=0.8.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["app"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py312"
|
||||
line-length = 120
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "N", "UP", "B"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
@@ -0,0 +1,93 @@
|
||||
"""LLM-as-judge 评测指标:基于 Ollama 小模型的幻觉率与类目一致性判定
|
||||
|
||||
所有判定均要求模型以 JSON 输出(json_mode),解析失败时保守处理并记录日志:
|
||||
- hallucination_rate 解析失败返回 0.0(不低报问题,避免阻塞评测主流程)
|
||||
- taxonomy_consistency 解析失败返回 True(不冤枉摘要)
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
|
||||
import structlog
|
||||
|
||||
from app.services.ollama import OllamaClient
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 每次判定的断言抽取上限,避免长摘要导致判定过慢
|
||||
_MAX_CLAIMS = 5
|
||||
|
||||
_HALLUCINATION_PROMPT = """你是一个事实核查员。请从下面的【摘要】中抽取最多 {max_claims} 条事实性断言,
|
||||
并逐条判断【原文】是否支持该断言(断言中的数字、名称、条款等关键信息必须与原文一致才算支持)。
|
||||
|
||||
【原文】
|
||||
{source}
|
||||
|
||||
【摘要】
|
||||
{summary}
|
||||
|
||||
只输出 JSON,格式:{{"assertions": [{{"claim": "断言内容", "supported": true 或 false}}]}}"""
|
||||
|
||||
_TAXONOMY_PROMPT = """你是一个文档分类审核员。请判断下面的【摘要】所描述的文档是否适合归入类目【{category}】。
|
||||
|
||||
【摘要】
|
||||
{summary}
|
||||
|
||||
只输出 JSON,格式:{{"consistent": true 或 false}}"""
|
||||
|
||||
|
||||
def _extract_json(text: str) -> dict | None:
|
||||
"""从模型输出提取首个 JSON 对象(容忍前后噪声),失败返回 None"""
|
||||
match = re.search(r"\{.*\}", text, re.DOTALL)
|
||||
if not match:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(match.group())
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
async def hallucination_rate(summary: str, source_text: str, ollama: OllamaClient) -> float:
|
||||
"""幻觉率:摘要中不被原文支持的断言占比
|
||||
|
||||
流程:模型抽取断言(最多 _MAX_CLAIMS 条)→ 逐条判原文是否支持 → 返回不支持比例。
|
||||
调用或解析失败保守返回 0.0 并记录日志。
|
||||
"""
|
||||
prompt = _HALLUCINATION_PROMPT.format(max_claims=_MAX_CLAIMS, source=source_text, summary=summary)
|
||||
try:
|
||||
raw = await ollama.generate(prompt, json_mode=True)
|
||||
except Exception as exc:
|
||||
logger.warning("幻觉率判定调用失败,保守返回 0.0", error=str(exc))
|
||||
return 0.0
|
||||
|
||||
data = _extract_json(raw)
|
||||
if data is None:
|
||||
logger.warning("幻觉率判定输出解析失败,保守返回 0.0", raw=raw[:200])
|
||||
return 0.0
|
||||
|
||||
claims = data.get("assertions")
|
||||
if not isinstance(claims, list):
|
||||
logger.warning("幻觉率判定输出缺少 assertions 字段,保守返回 0.0", raw=raw[:200])
|
||||
return 0.0
|
||||
supported_flags = [bool(c["supported"]) for c in claims[:_MAX_CLAIMS] if isinstance(c, dict) and "supported" in c]
|
||||
if not supported_flags:
|
||||
return 0.0
|
||||
unsupported = sum(1 for supported in supported_flags if not supported)
|
||||
return unsupported / len(supported_flags)
|
||||
|
||||
|
||||
async def taxonomy_consistency(summary: str, category: str, ollama: OllamaClient) -> bool:
|
||||
"""类目一致性:摘要是否支持归入指定类目;调用或解析失败保守返回 True"""
|
||||
prompt = _TAXONOMY_PROMPT.format(category=category, summary=summary)
|
||||
try:
|
||||
raw = await ollama.generate(prompt, json_mode=True)
|
||||
except Exception as exc:
|
||||
logger.warning("类目一致性判定调用失败,保守返回 True", error=str(exc))
|
||||
return True
|
||||
|
||||
data = _extract_json(raw)
|
||||
if data is None or "consistent" not in data:
|
||||
logger.warning("类目一致性判定输出解析失败,保守返回 True", raw=raw[:200])
|
||||
return True
|
||||
return bool(data["consistent"])
|
||||
@@ -0,0 +1,106 @@
|
||||
"""离线评测指标:纯函数实现,不依赖 Qdrant / Ollama,可独立单测
|
||||
|
||||
- entity_recall:实体保留率(摘要质量)
|
||||
- routing_f1:L1 路由层 micro P/R/F1
|
||||
- pruning_loss:上层剪枝损失率
|
||||
- precision_at_k / recall_at_k:检索效用
|
||||
- aggregate:按 key 汇总均值
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
# 数字串(含小数/版本号/百分数),如 1536、3.12、95%
|
||||
_NUMBER_RE = re.compile(r"\d+(?:\.\d+)*%?")
|
||||
# 型号/版本号模式,如 bge-m3、v1.2.0、FD-07
|
||||
_MODEL_RE = re.compile(r"[A-Za-z]+[-_]\w*\d\w*|[A-Za-z]+\d+(?:\.\d+)*")
|
||||
# 条款号,如 第三条、第 5 章、第十条
|
||||
_CLAUSE_RE = re.compile(r"第\s*[一二三四五六七八九十百千万零0-9]+\s*[条款章节项]")
|
||||
# 英文/大小写混合词(含连字符/点号连接),如 Qdrant、OpenAI、text-embedding-3-small
|
||||
_EN_WORD_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*(?:[-_.][A-Za-z0-9]+)*")
|
||||
# 中文书名号内词,如 《部署手册》
|
||||
_BOOK_TITLE_RE = re.compile(r"《([^》]+)》")
|
||||
|
||||
|
||||
def _extract_entities(text: str) -> set[str]:
|
||||
"""从文本抽取关键实体 token(去重集合)"""
|
||||
entities: set[str] = set()
|
||||
entities.update(_NUMBER_RE.findall(text))
|
||||
entities.update(_MODEL_RE.findall(text))
|
||||
entities.update(m.replace(" ", "") for m in _CLAUSE_RE.findall(text))
|
||||
entities.update(_BOOK_TITLE_RE.findall(text))
|
||||
# 英文词至少 2 个字符,避免单字母噪声
|
||||
entities.update(w for w in _EN_WORD_RE.findall(text) if len(w) >= 2)
|
||||
return entities
|
||||
|
||||
|
||||
def entity_recall(source_text: str, summary: str) -> float:
|
||||
"""实体保留率:原文抽取的关键实体在摘要中保留的比例
|
||||
|
||||
原文无可抽取实体时返回 1.0(无实体可丢,视为满分)。
|
||||
"""
|
||||
entities = _extract_entities(source_text)
|
||||
if not entities:
|
||||
return 1.0
|
||||
kept = sum(1 for entity in entities if entity in summary)
|
||||
return kept / len(entities)
|
||||
|
||||
|
||||
def routing_f1(golden_doc_ids_per_query: list[set[str]], routed_doc_ids_per_query: list[set[str]]) -> dict[str, float]:
|
||||
"""L1 路由层 micro precision / recall / f1
|
||||
|
||||
逐 query 累加 tp/fp/fn 后统一计算(micro 平均);
|
||||
两条列表必须等长,通常只统计 positive query(negative query 无 golden doc)。
|
||||
"""
|
||||
tp = fp = fn = 0
|
||||
for golden, routed in zip(golden_doc_ids_per_query, routed_doc_ids_per_query, strict=True):
|
||||
tp += len(golden & routed)
|
||||
fp += len(routed - golden)
|
||||
fn += len(golden - routed)
|
||||
precision = tp / (tp + fp) if tp + fp else 0.0
|
||||
recall = tp / (tp + fn) if tp + fn else 0.0
|
||||
f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
|
||||
return {"precision": precision, "recall": recall, "f1": f1}
|
||||
|
||||
|
||||
def pruning_loss(cases: list[bool]) -> float:
|
||||
"""剪枝损失率:原文含答案但 L1 未命中 golden doc 的 positive query 占比
|
||||
|
||||
cases 中每个元素表示一条 positive query 是否被上层剪掉(True = 被剪)。
|
||||
"""
|
||||
if not cases:
|
||||
return 0.0
|
||||
return sum(cases) / len(cases)
|
||||
|
||||
|
||||
def precision_at_k(hit_doc_ids: list[str], golden: set[str], k: int) -> float:
|
||||
"""Precision@k:前 k 条命中中 doc_id 属于 golden 的比例(按命中条数计,不去重)"""
|
||||
if k <= 0:
|
||||
return 0.0
|
||||
top = hit_doc_ids[:k]
|
||||
if not top:
|
||||
return 0.0
|
||||
hits = sum(1 for doc_id in top if doc_id in golden)
|
||||
return hits / k
|
||||
|
||||
|
||||
def recall_at_k(hit_doc_ids: list[str], golden: set[str], k: int) -> float:
|
||||
"""Recall@k:前 k 条命中覆盖的 golden doc 比例(按 doc 去重);golden 为空返回 0.0"""
|
||||
if not golden:
|
||||
return 0.0
|
||||
found = {doc_id for doc_id in hit_doc_ids[:k] if doc_id in golden}
|
||||
return len(found) / len(golden)
|
||||
|
||||
|
||||
def aggregate(records: list[dict[str, float]]) -> dict[str, float]:
|
||||
"""按 key 汇总均值:对记录列表中每个指标求算术平均
|
||||
|
||||
某 key 只在部分记录中出现时,按出现的记录取均值;空列表返回 {}。
|
||||
"""
|
||||
if not records:
|
||||
return {}
|
||||
keys = list(dict.fromkeys(key for record in records for key in record))
|
||||
result: dict[str, float] = {}
|
||||
for key in keys:
|
||||
values = [record[key] for record in records if key in record]
|
||||
result[key] = sum(values) / len(values) if values else 0.0
|
||||
return result
|
||||
@@ -0,0 +1,73 @@
|
||||
{
|
||||
"documents": [
|
||||
{
|
||||
"id": "doc001",
|
||||
"title": "QMDSearch 向量检索服务部署运维手册",
|
||||
"text": "# 部署概述\nQMDSearch v1.2.0 依赖 Qdrant 1.12、Redis 7.2 与 Ollama 0.5 三个组件,默认端口分别为 6333、6379 与 11434。生产环境推荐使用 docker compose 一键拉起,健康检查路径为 /health,首次启动约需 30 秒完成集合初始化。\n# 关键配置\nembedding_dimension 默认 1536,sparse_enabled 默认开启。按照《运维规范》第三条要求:修改 retrieval_top_k 或 l1_doc_top_n 后必须重启服务,并执行回归冒烟用例。\n# 故障排查\n若 Qdrant 连接超时,先检查容器网络 qmd-net 是否互通;Ollama 拉取模型失败时重试 ollama pull qwen2.5:1.5b;Redis 不可用时会自动降级为直连检索,不影响主流程。",
|
||||
"golden_category": "技术文档"
|
||||
},
|
||||
{
|
||||
"id": "doc002",
|
||||
"title": "Embedding 服务接入指南",
|
||||
"text": "# 接入方式\nEmbeddingService 统一接口支持 openai 与 local 两种 provider。openai provider 默认模型 text-embedding-3-small,走 OpenAI 兼容 API;local provider 调用 Ollama 的 /api/embed 接口,默认模型 bge-m3,超时 60 秒。\n# 批量调用\nembed(texts) 支持批量编码,传入空列表时直接返回空列表。单次批量建议不超过 64 条文本,超长文本建议先按 800 字符切分。\n# 维度校验\n返回向量维度与 settings.embedding_dimension 不一致时仅记录 warning 不抛错,且每次调用最多提示一次,避免日志刷屏。",
|
||||
"golden_category": "技术文档"
|
||||
},
|
||||
{
|
||||
"id": "doc003",
|
||||
"title": "智能搜索助手 Pro 2.0 产品说明书",
|
||||
"text": "# 产品简介\n智能搜索助手 Pro 2.0 面向企业知识库场景,单机支持 200 人并发查询,平均响应时间 350 毫秒,P99 延迟不超过 1.2 秒。\n# 核心功能\n产品支持分层摘要检索、类目路由与平铺全文检索三种模式,可按租户切换。第 5 章介绍高级筛选语法,更多示例见《用户操作手册》。\n# 版本记录\n2.0 版新增 RRF 融合排序与 sparse 稀疏检索;1.8 版的旧版筛选语法兼容至 2026 年 12 月 31 日,之后停止维护。",
|
||||
"golden_category": "产品手册"
|
||||
},
|
||||
{
|
||||
"id": "doc004",
|
||||
"title": "云文档协作平台快速上手指南",
|
||||
"text": "# 快速开始\n注册账号后 3 分钟内即可创建首个知识空间,免费额度包含 5GB 存储与每月 1000 次 API 调用,超出后按 0.01 元每次计费。\n# 协作功能\n平台支持多人实时编辑、行内评论与版本回滚,历史版本默认保留 90 天,企业版可延长至 365 天。详见《协作白皮书》第二章。\n# 移动端\niOS 与 Android 客户端支持离线缓存,单文件上限 200MB,弱网环境下自动启用增量同步。",
|
||||
"golden_category": "产品手册"
|
||||
},
|
||||
{
|
||||
"id": "doc005",
|
||||
"title": "差旅费用报销管理办法",
|
||||
"text": "# 适用范围\n本办法适用于全体正式员工与实习生,自 2026 年 1 月 1 日起施行,原 2024 版办法同时废止。\n# 报销标准\n一线城市住宿限额每晚 500 元,二线城市 400 元,其他城市 300 元;高铁二等座、飞机经济舱据实报销。第十条规定:超标部分需部门总监书面审批。\n# 报销流程\n发票开具后 30 日内须提交 OA 系统,超过 90 天的票据不予受理。常见问题见《财务报销常见问题》第七条。",
|
||||
"golden_category": "财务行政"
|
||||
},
|
||||
{
|
||||
"id": "doc006",
|
||||
"title": "固定资产采购与领用规定",
|
||||
"text": "# 采购审批\n单笔金额超过 5000 元的采购须走 OA 审批流,超过 20000 元需分管副总裁签字,紧急采购可先邮件报备后补流程。\n# 领用登记\n固定资产领用后 3 个工作日内在行政系统登记资产编号,编号规则见《资产管理细则》第四条,笔记本等移动设备须加贴防伪标签。\n# 盘点与报废\n每年 12 月进行年度盘点,报废资产须填写 FD-07 表单并附照片存档,残值率按 5% 计提。",
|
||||
"golden_category": "财务行政"
|
||||
}
|
||||
],
|
||||
"queries": [
|
||||
{"query": "QMDSearch 依赖哪些组件,分别用什么端口?", "type": "positive", "golden_doc_id": "doc001", "golden_section": "部署概述"},
|
||||
{"query": "修改 retrieval_top_k 之后需要做什么?", "type": "positive", "golden_doc_id": "doc001", "golden_section": "关键配置"},
|
||||
{"query": "Ollama 拉取模型失败应该怎么处理?", "type": "positive", "golden_doc_id": "doc001", "golden_section": "故障排查"},
|
||||
{"query": "embedding_dimension 的默认值是多少?", "type": "positive", "golden_doc_id": "doc001", "golden_section": "关键配置"},
|
||||
{"query": "公司组织年度体检的医院是哪家?", "type": "negative", "golden_doc_id": null},
|
||||
|
||||
{"query": "local provider 默认使用哪个嵌入模型?", "type": "positive", "golden_doc_id": "doc002", "golden_section": "接入方式"},
|
||||
{"query": "embed 接口传入空列表会返回什么?", "type": "positive", "golden_doc_id": "doc002", "golden_section": "批量调用"},
|
||||
{"query": "向量维度和配置不一致时会抛异常吗?", "type": "positive", "golden_doc_id": "doc002", "golden_section": "维度校验"},
|
||||
{"query": "如何申请欧洲申根旅游签证?", "type": "negative", "golden_doc_id": null},
|
||||
|
||||
{"query": "智能搜索助手 Pro 2.0 支持多少并发查询?", "type": "positive", "golden_doc_id": "doc003", "golden_section": "产品简介"},
|
||||
{"query": "高级筛选语法在产品说明书的哪一章介绍?", "type": "positive", "golden_doc_id": "doc003", "golden_section": "核心功能"},
|
||||
{"query": "旧版筛选语法兼容到什么时候?", "type": "positive", "golden_doc_id": "doc003", "golden_section": "版本记录"},
|
||||
{"query": "竞品的按年订阅价格是多少?", "type": "negative", "golden_doc_id": null},
|
||||
|
||||
{"query": "云文档平台免费额度包含多少存储空间?", "type": "positive", "golden_doc_id": "doc004", "golden_section": "快速开始"},
|
||||
{"query": "历史版本默认保留多长时间?", "type": "positive", "golden_doc_id": "doc004", "golden_section": "协作功能"},
|
||||
{"query": "移动端单文件上传上限是多少?", "type": "positive", "golden_doc_id": "doc004", "golden_section": "移动端"},
|
||||
{"query": "视频会议最多支持多少人同时开启摄像头?", "type": "negative", "golden_doc_id": null},
|
||||
|
||||
{"query": "一线城市住宿报销限额是多少?", "type": "positive", "golden_doc_id": "doc005", "golden_section": "报销标准"},
|
||||
{"query": "发票超过多少天就不能报销了?", "type": "positive", "golden_doc_id": "doc005", "golden_section": "报销流程"},
|
||||
{"query": "住宿超标部分需要谁审批?", "type": "positive", "golden_doc_id": "doc005", "golden_section": "报销标准"},
|
||||
{"query": "新版差旅报销办法什么时候开始施行?", "type": "positive", "golden_doc_id": "doc005", "golden_section": "适用范围"},
|
||||
{"query": "员工结婚礼金的公司福利标准是多少?", "type": "negative", "golden_doc_id": null},
|
||||
|
||||
{"query": "采购金额超过多少需要副总裁签字?", "type": "positive", "golden_doc_id": "doc006", "golden_section": "采购审批"},
|
||||
{"query": "资产编号规则在哪个文件里规定的?", "type": "positive", "golden_doc_id": "doc006", "golden_section": "领用登记"},
|
||||
{"query": "报废资产需要填写什么表单?", "type": "positive", "golden_doc_id": "doc006", "golden_section": "盘点与报废"},
|
||||
{"query": "办公区绿植养护的排班表在哪里查?", "type": "negative", "golden_doc_id": null}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
#!/usr/bin/env python3
|
||||
"""分层摘要 RAG 离线评测 harness
|
||||
|
||||
流程:回归集 → 独立 _eval 后缀集合入库(Ingester)→ 摘要质量指标(Entity Recall /
|
||||
幻觉率 / 类目一致性)→ 检索效用指标(Routing F1 / Pruning Loss / Precision@5 /
|
||||
Recall@10)→ 平铺 chunks baseline 对比 → Markdown 报告(stdout + report.md)。
|
||||
|
||||
用法:
|
||||
uv run python scripts/eval/run_eval.py [--regression PATH] [--keep-data] [--no-judge]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# 脚本直运行时把项目根加入 sys.path,保证可以 import app 与 scripts.eval
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(_PROJECT_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_PROJECT_ROOT))
|
||||
|
||||
import structlog # noqa: E402
|
||||
|
||||
from app.config import settings # noqa: E402
|
||||
from app.core import ingestion as ingestion_mod # noqa: E402
|
||||
from app.core import retriever as retriever_mod # noqa: E402
|
||||
from app.core.ingestion import Ingester # noqa: E402
|
||||
from app.core.retriever import Retriever # noqa: E402
|
||||
from app.models.document import DocumentInput # noqa: E402
|
||||
from app.models.search import SearchRequest # noqa: E402
|
||||
from app.services import qdrant as qdrant_mod # noqa: E402
|
||||
from app.services.ollama import OllamaClient # noqa: E402
|
||||
from app.services.qdrant import QdrantService # noqa: E402
|
||||
from scripts.eval.judge import hallucination_rate, taxonomy_consistency # noqa: E402
|
||||
from scripts.eval.metrics import ( # noqa: E402
|
||||
aggregate,
|
||||
entity_recall,
|
||||
precision_at_k,
|
||||
pruning_loss,
|
||||
recall_at_k,
|
||||
routing_f1,
|
||||
)
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 评测集合后缀与默认路径
|
||||
EVAL_SUFFIX = "_eval"
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_REGRESSION = _SCRIPT_DIR / "regression_set.json"
|
||||
REPORT_PATH = _SCRIPT_DIR / "report.md"
|
||||
|
||||
# 门槛(Spec 规定):低于/高于门槛在报告中标红
|
||||
THRESHOLD_L1_ER = 0.85 # L1 Entity Recall ≥ 0.85
|
||||
THRESHOLD_L3_ER = 0.9 # L3 Entity Recall ≥ 0.9
|
||||
THRESHOLD_HALLUCINATION = 0.02 # 幻觉率 < 2%
|
||||
THRESHOLD_PRUNING_LOSS = 0.08 # Pruning Loss < 8%
|
||||
|
||||
|
||||
def _switch_to_eval_collections() -> list[str]:
|
||||
"""把 app 内模块级集合名常量整体切换为 _eval 后缀的评测集合
|
||||
|
||||
QdrantService / Retriever / Ingester 均在各自模块命名空间引用了集合名常量,
|
||||
评测脚本统一改写这些模块属性实现集合隔离,不改动 app 源码。
|
||||
返回评测集合名列表(用于评测结束后清理)。
|
||||
"""
|
||||
eval_names = {
|
||||
"COLLECTION_L1": f"{qdrant_mod.COLLECTION_L1}{EVAL_SUFFIX}",
|
||||
"COLLECTION_L2": f"{qdrant_mod.COLLECTION_L2}{EVAL_SUFFIX}",
|
||||
"COLLECTION_L3": f"{qdrant_mod.COLLECTION_L3}{EVAL_SUFFIX}",
|
||||
"COLLECTION_CHUNKS": f"{qdrant_mod.COLLECTION_CHUNKS}{EVAL_SUFFIX}",
|
||||
}
|
||||
for module in (qdrant_mod, retriever_mod, ingestion_mod):
|
||||
for name, value in eval_names.items():
|
||||
if hasattr(module, name):
|
||||
setattr(module, name, value)
|
||||
qdrant_mod.ALL_COLLECTIONS = tuple(eval_names.values())
|
||||
sparse_eval = (eval_names["COLLECTION_L1"], eval_names["COLLECTION_CHUNKS"])
|
||||
qdrant_mod.SPARSE_COLLECTIONS = sparse_eval
|
||||
retriever_mod.SPARSE_COLLECTIONS = sparse_eval
|
||||
# upsert_nodes 按集合名校验层级前缀,需同步替换
|
||||
qdrant_mod._NODE_ID_PREFIX = {eval_names["COLLECTION_L2"]: "l2", eval_names["COLLECTION_L3"]: "l3"}
|
||||
return list(eval_names.values())
|
||||
|
||||
|
||||
async def _check_services(qdrant: QdrantService, ollama: OllamaClient) -> str | None:
|
||||
"""检查 Qdrant / Ollama 连通性,返回错误消息(None 表示正常)"""
|
||||
try:
|
||||
await qdrant.client.get_collections()
|
||||
except Exception as exc:
|
||||
return f"无法连接 Qdrant({settings.qdrant_host}:{settings.qdrant_port}):{exc}"
|
||||
if not await ollama.is_available():
|
||||
return f"无法连接 Ollama({settings.ollama_base_url}),请确认服务已启动(入库与评测均依赖 Ollama)"
|
||||
return None
|
||||
|
||||
|
||||
async def _l1_candidate_doc_ids(retriever: Retriever, query: str) -> set[str]:
|
||||
"""轻量复现 Retriever 的 L1 路由层:embed → L1 集合 top-N → 候选 doc_id 集合
|
||||
|
||||
用于计算 Pruning Loss 与 Routing F1;不套类目过滤,度量纯 L1 向量召回。
|
||||
"""
|
||||
dense = (await retriever.embedding.embed([query]))[0]
|
||||
sparse = retriever.sparse_encoder.encode(query) if settings.sparse_enabled else None
|
||||
# 复用 Retriever 内部检索方法(hybrid/dense 按集合能力自动选择)
|
||||
hits = await retriever._search_collection(retriever_mod.COLLECTION_L1, dense, sparse, settings.l1_doc_top_n, None)
|
||||
return {(p.payload or {}).get("doc_id", "") for p in hits} - {""}
|
||||
|
||||
|
||||
async def _baseline_doc_ids(retriever: Retriever, query: str) -> list[str]:
|
||||
"""平铺 baseline:直接对 chunks 集合做 hybrid top-k(无路由无剪枝),返回命中 doc_id 列表"""
|
||||
dense = (await retriever.embedding.embed([query]))[0]
|
||||
sparse = retriever.sparse_encoder.encode(query) if settings.sparse_enabled else None
|
||||
points = await retriever._search_collection(
|
||||
retriever_mod.COLLECTION_CHUNKS, dense, sparse, settings.retrieval_top_k, None
|
||||
)
|
||||
return [(p.payload or {}).get("doc_id", "") for p in points]
|
||||
|
||||
|
||||
def _mark(ok: bool) -> str:
|
||||
"""门槛判定标记"""
|
||||
return "✅" if ok else "❌"
|
||||
|
||||
|
||||
def _fmt(value: float | None, percent: bool = False) -> str:
|
||||
"""数值格式化;None 表示未评测(judge 跳过)"""
|
||||
if value is None:
|
||||
return "N/A"
|
||||
return f"{value:.1%}" if percent else f"{value:.4f}"
|
||||
|
||||
|
||||
def _build_report(
|
||||
regression_path: Path,
|
||||
doc_records: list[dict],
|
||||
query_summary: dict,
|
||||
hier_metrics: dict[str, float],
|
||||
baseline_metrics: dict[str, float],
|
||||
routing: dict[str, float],
|
||||
prune_loss: float,
|
||||
judge_enabled: bool,
|
||||
kept_data: bool,
|
||||
) -> str:
|
||||
"""组装 Markdown 评测报告"""
|
||||
lines: list[str] = [
|
||||
"# 分层摘要 RAG 评测报告",
|
||||
"",
|
||||
f"- 回归集:`{regression_path}`",
|
||||
f"- 文档数:{len(doc_records)};query 数:{query_summary['total']}"
|
||||
f"(positive {query_summary['positive']} / negative {query_summary['negative']})",
|
||||
f"- LLM judge:{'开启' if judge_enabled else '跳过(--no-judge 或 Ollama 不可用)'}",
|
||||
f"- 评测集合:`*{EVAL_SUFFIX}`({'保留' if kept_data else '已清理'})",
|
||||
"",
|
||||
"## 摘要质量(按文档)",
|
||||
"",
|
||||
"| 文档 | 标题 | golden 类目 | 实际类目 | L1 Entity Recall | L3 Entity Recall | 幻觉率 | 类目一致 |",
|
||||
"| --- | --- | --- | --- | --- | --- | --- | --- |",
|
||||
]
|
||||
for record in doc_records:
|
||||
consistency = record.get("taxonomy_consistency")
|
||||
lines.append(
|
||||
f"| {record['id']} | {record['title']} | {record['golden_category']} | {record['category']} "
|
||||
f"| {_fmt(record['entity_recall_l1'])} | {_fmt(record['entity_recall_l3'])} "
|
||||
f"| {_fmt(record.get('hallucination_rate'), percent=True)} "
|
||||
f"| {('是' if consistency else '否') if consistency is not None else 'N/A'} |"
|
||||
)
|
||||
|
||||
# 汇总与门槛判定
|
||||
er_l1 = aggregate([{"v": r["entity_recall_l1"]} for r in doc_records]).get("v", 0.0)
|
||||
er_l3 = aggregate([{"v": r["entity_recall_l3"]} for r in doc_records]).get("v", 0.0)
|
||||
hall_records = [{"v": r["hallucination_rate"]} for r in doc_records if r.get("hallucination_rate") is not None]
|
||||
hall = aggregate(hall_records).get("v") if hall_records else None
|
||||
lines += [
|
||||
"",
|
||||
"## 指标汇总与门槛判定",
|
||||
"",
|
||||
"| 指标 | 数值 | 门槛 | 判定 |",
|
||||
"| --- | --- | --- | --- |",
|
||||
f"| L1 Entity Recall | {_fmt(er_l1)} | ≥ {THRESHOLD_L1_ER} | {_mark(er_l1 >= THRESHOLD_L1_ER)} |",
|
||||
f"| L3 Entity Recall | {_fmt(er_l3)} | ≥ {THRESHOLD_L3_ER} | {_mark(er_l3 >= THRESHOLD_L3_ER)} |",
|
||||
f"| Hallucination Rate | {_fmt(hall, percent=True)} | < {THRESHOLD_HALLUCINATION:.0%} "
|
||||
f"| {_mark(hall < THRESHOLD_HALLUCINATION) if hall is not None else 'N/A'} |",
|
||||
f"| Pruning Loss | {_fmt(prune_loss, percent=True)} | < {THRESHOLD_PRUNING_LOSS:.0%} "
|
||||
f"| {_mark(prune_loss < THRESHOLD_PRUNING_LOSS)} |",
|
||||
"",
|
||||
"## 检索效用(hierarchical vs 平铺 baseline)",
|
||||
"",
|
||||
"| 指标 | 分层检索 | 平铺 baseline |",
|
||||
"| --- | --- | --- |",
|
||||
f"| Precision@5 | {_fmt(hier_metrics.get('precision@5', 0.0))} "
|
||||
f"| {_fmt(baseline_metrics.get('precision@5', 0.0))} |",
|
||||
f"| Recall@10 | {_fmt(hier_metrics.get('recall@10', 0.0))} | {_fmt(baseline_metrics.get('recall@10', 0.0))} |",
|
||||
"",
|
||||
"### 路由层(L1)",
|
||||
"",
|
||||
f"- Routing Precision:{_fmt(routing['precision'])}",
|
||||
f"- Routing Recall:{_fmt(routing['recall'])}",
|
||||
f"- Routing F1:{_fmt(routing['f1'])}",
|
||||
f"- Pruning Loss:{_fmt(prune_loss, percent=True)}({_mark(prune_loss < THRESHOLD_PRUNING_LOSS)})",
|
||||
"",
|
||||
"### negative query",
|
||||
"",
|
||||
f"- 误中(返回了任意结果):{query_summary['negative_false_alarm']} / {query_summary['negative']}",
|
||||
"",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def run(regression_path: Path, keep_data: bool, judge_enabled: bool) -> int:
|
||||
"""评测主流程,返回进程退出码"""
|
||||
data = json.loads(regression_path.read_text(encoding="utf-8"))
|
||||
documents: list[dict] = data["documents"]
|
||||
queries: list[dict] = data["queries"]
|
||||
|
||||
eval_collections = _switch_to_eval_collections()
|
||||
qdrant = QdrantService()
|
||||
ollama = OllamaClient()
|
||||
|
||||
# 连通性检查:失败给出中文提示并以退出码 2 结束
|
||||
error = await _check_services(qdrant, ollama)
|
||||
if error:
|
||||
print(f"错误:{error}", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
await qdrant.ensure_collections()
|
||||
logger.info("评测集合就绪", collections=eval_collections)
|
||||
|
||||
ingester = Ingester(qdrant=qdrant)
|
||||
retriever = Retriever(qdrant=qdrant)
|
||||
|
||||
id_map: dict[str, str] = {} # 回归集 doc id -> 实际入库 doc_id
|
||||
doc_records: list[dict] = []
|
||||
hier_records: list[dict[str, float]] = []
|
||||
baseline_records: list[dict[str, float]] = []
|
||||
pruned_cases: list[bool] = []
|
||||
golden_sets: list[set[str]] = []
|
||||
routed_sets: list[set[str]] = []
|
||||
negative_total = 0
|
||||
negative_false_alarm = 0
|
||||
|
||||
try:
|
||||
# 1. 逐篇入库并计算摘要质量指标
|
||||
for doc in documents:
|
||||
result = await ingester.ingest(DocumentInput(title=doc["title"], text=doc["text"]))
|
||||
id_map[doc["id"]] = result.document_id
|
||||
logger.info("文档入库完成", id=doc["id"], doc_id=result.document_id, category=result.category)
|
||||
|
||||
record: dict = {
|
||||
"id": doc["id"],
|
||||
"title": doc["title"],
|
||||
"golden_category": doc["golden_category"],
|
||||
"category": result.category,
|
||||
"entity_recall_l1": entity_recall(doc["text"], result.summary.l1_summary),
|
||||
"entity_recall_l3": entity_recall(doc["text"], result.summary.l3_content_outline),
|
||||
"hallucination_rate": None,
|
||||
"taxonomy_consistency": None,
|
||||
}
|
||||
if judge_enabled:
|
||||
record["hallucination_rate"] = await hallucination_rate(result.summary.l1_summary, doc["text"], ollama)
|
||||
record["taxonomy_consistency"] = await taxonomy_consistency(
|
||||
result.summary.l1_summary, result.category, ollama
|
||||
)
|
||||
doc_records.append(record)
|
||||
|
||||
# 2. 逐 query 计算检索效用指标(分层检索 + 轻量 L1 路由层 + 平铺 baseline)
|
||||
for query in queries:
|
||||
query_text = query["query"]
|
||||
golden_ids = {id_map[query["golden_doc_id"]]} if query.get("golden_doc_id") else set()
|
||||
|
||||
response = await retriever.search(SearchRequest(query=query_text))
|
||||
hit_doc_ids = [hit.doc_id for hit in response.hits]
|
||||
l1_candidates = await _l1_candidate_doc_ids(retriever, query_text)
|
||||
|
||||
if query["type"] == "positive" and golden_ids:
|
||||
pruned_cases.append(not golden_ids & l1_candidates)
|
||||
golden_sets.append(golden_ids)
|
||||
routed_sets.append(l1_candidates)
|
||||
hier_records.append(
|
||||
{
|
||||
"precision@5": precision_at_k(hit_doc_ids, golden_ids, 5),
|
||||
"recall@10": recall_at_k(hit_doc_ids, golden_ids, 10),
|
||||
}
|
||||
)
|
||||
baseline_ids = await _baseline_doc_ids(retriever, query_text)
|
||||
baseline_records.append(
|
||||
{
|
||||
"precision@5": precision_at_k(baseline_ids, golden_ids, 5),
|
||||
"recall@10": recall_at_k(baseline_ids, golden_ids, 10),
|
||||
}
|
||||
)
|
||||
else:
|
||||
# negative query:期望无结果,任何返回均计为误中
|
||||
negative_total += 1
|
||||
if hit_doc_ids:
|
||||
negative_false_alarm += 1
|
||||
finally:
|
||||
# 评测集合清理(--keep-data 时保留)
|
||||
if keep_data:
|
||||
logger.info("--keep-data 生效,保留评测集合", collections=eval_collections)
|
||||
else:
|
||||
for collection in eval_collections:
|
||||
try:
|
||||
await qdrant.client.delete_collection(collection)
|
||||
except Exception as exc:
|
||||
logger.warning("评测集合清理失败", collection=collection, error=str(exc))
|
||||
logger.info("评测集合已清理", collections=eval_collections)
|
||||
|
||||
# 3. 汇总并输出报告
|
||||
prune_loss = pruning_loss(pruned_cases)
|
||||
routing = routing_f1(golden_sets, routed_sets)
|
||||
hier_metrics = aggregate(hier_records)
|
||||
baseline_metrics = aggregate(baseline_records)
|
||||
query_summary = {
|
||||
"total": len(queries),
|
||||
"positive": len(queries) - negative_total,
|
||||
"negative": negative_total,
|
||||
"negative_false_alarm": negative_false_alarm,
|
||||
}
|
||||
report = _build_report(
|
||||
regression_path,
|
||||
doc_records,
|
||||
query_summary,
|
||||
hier_metrics,
|
||||
baseline_metrics,
|
||||
routing,
|
||||
prune_loss,
|
||||
judge_enabled,
|
||||
keep_data,
|
||||
)
|
||||
print(report)
|
||||
REPORT_PATH.write_text(report + "\n", encoding="utf-8")
|
||||
logger.info("评测报告已写入", path=str(REPORT_PATH))
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="分层摘要 RAG 离线评测:回归集入库 → 摘要质量/检索效用指标 → 平铺 baseline 对比 → Markdown 报告",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--regression",
|
||||
type=Path,
|
||||
default=DEFAULT_REGRESSION,
|
||||
help=f"回归集 JSON 路径(默认 {DEFAULT_REGRESSION})",
|
||||
)
|
||||
parser.add_argument("--keep-data", action="store_true", help="评测结束后保留 _eval 集合(默认删除)")
|
||||
parser.add_argument("--no-judge", action="store_true", help="跳过 LLM-as-judge 指标(幻觉率 / 类目一致性)")
|
||||
args = parser.parse_args()
|
||||
return asyncio.run(run(args.regression.resolve(), args.keep_data, not args.no_judge))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,165 @@
|
||||
"""真实 Ollama 现场冒烟脚本
|
||||
|
||||
用法:uv run python scripts/smoke_live.py
|
||||
|
||||
链路:内存 Qdrant + 真实本地 Ollama(deepseek-r1:8b)+ 确定性哈希向量(无语义,仅打通流程)。
|
||||
读取回归集第一篇文档走完整 ingest → search 流程,打印各阶段结果、原始 LLM 输出与耗时。
|
||||
|
||||
deepseek-r1 是推理模型,/api/generate 的 response 字段可能带 <think> 思考内容;
|
||||
classifier / query_parser 内建 JSON 容错,若解析失败会走兜底路径——均为合理的现场观察。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
|
||||
# 脚本直接运行时需要把项目根目录加入 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from qdrant_client import AsyncQdrantClient # noqa: E402
|
||||
|
||||
from app.config import settings # noqa: E402
|
||||
from app.core.chunker import Chunker # noqa: E402
|
||||
from app.core.classifier import Classifier # noqa: E402
|
||||
from app.core.ingestion import Ingester # noqa: E402
|
||||
from app.core.query_parser import QueryParser # noqa: E402
|
||||
from app.core.retriever import Retriever # noqa: E402
|
||||
from app.core.sparse import SparseEncoder # noqa: E402
|
||||
from app.core.summarizer import Summarizer # noqa: E402
|
||||
from app.models.document import DocumentInput # noqa: E402
|
||||
from app.models.knowledge import load_taxonomy # noqa: E402
|
||||
from app.models.search import SearchRequest # noqa: E402
|
||||
from app.services.ollama import OllamaClient # noqa: E402
|
||||
from app.services.qdrant import QdrantService # noqa: E402
|
||||
|
||||
OLLAMA_MODEL = "deepseek-r1:8b"
|
||||
OLLAMA_TIMEOUT = 600.0
|
||||
|
||||
|
||||
class DeterministicEmbedding:
|
||||
"""稳定哈希伪向量:同文本恒同向量,维度等于 settings.embedding_dimension(无语义)"""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
vectors: list[list[float]] = []
|
||||
for text in texts:
|
||||
seed = int.from_bytes(sha256(text.encode("utf-8")).digest()[:8], "big")
|
||||
rng = random.Random(seed)
|
||||
vectors.append([rng.random() for _ in range(settings.embedding_dimension)])
|
||||
return vectors
|
||||
|
||||
|
||||
class RecordingOllama:
|
||||
"""包装真实 OllamaClient,记录每次调用的原始输出供现场观察"""
|
||||
|
||||
def __init__(self, inner: OllamaClient) -> None:
|
||||
self.inner = inner
|
||||
self.records: list[tuple[str, str]] = [] # (调用用途, 原始输出)
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
kind = _prompt_kind(prompt)
|
||||
raw = await self.inner.generate(prompt, json_mode=json_mode)
|
||||
self.records.append((kind, raw))
|
||||
return raw
|
||||
|
||||
|
||||
def _prompt_kind(prompt: str) -> str:
|
||||
"""按 prompt 特征串标注调用用途"""
|
||||
if "你是知识库分类助手" in prompt:
|
||||
return "文档分类"
|
||||
if "你是搜索查询分析助手" in prompt:
|
||||
return "query 解析"
|
||||
if "请用一句话对以下文档内容进行高度概括" in prompt:
|
||||
return "L1 总结"
|
||||
if "请提取以下文档的主要章节结构" in prompt:
|
||||
return "L2 大纲"
|
||||
if "请对以下文档的每个章节/主题进行详细的内容摘要" in prompt:
|
||||
return "L3 内容大纲"
|
||||
if "请对以下文档内容进行详细摘要" in prompt:
|
||||
return "L2.5 摘要"
|
||||
return "未知"
|
||||
|
||||
|
||||
def _fmt(seconds: float) -> str:
|
||||
return f"{seconds:.1f}s"
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
started = time.perf_counter()
|
||||
|
||||
# 0. 环境检查
|
||||
ollama_inner = OllamaClient(base_url="http://localhost:11434", model=OLLAMA_MODEL, timeout=OLLAMA_TIMEOUT)
|
||||
if not await ollama_inner.is_available():
|
||||
print("Ollama 不可用(http://localhost:11434),请先启动 Ollama 服务")
|
||||
return 1
|
||||
print(f"[环境] Ollama 可用,模型 {OLLAMA_MODEL};Qdrant 使用 :memory: 本地模式")
|
||||
|
||||
# 1. 读取回归集第一篇文档与一条相关 query
|
||||
regression_path = Path(__file__).resolve().parent / "eval" / "regression_set.json"
|
||||
regression = json.loads(regression_path.read_text(encoding="utf-8"))
|
||||
doc_data = regression["documents"][0]
|
||||
query = next(q["query"] for q in regression["queries"] if q.get("golden_doc_id") == doc_data["id"])
|
||||
doc = DocumentInput(text=doc_data["text"], title=doc_data["title"])
|
||||
print(f"[数据] 文档《{doc.title}》({len(doc.text)} 字符);query:{query}")
|
||||
|
||||
# 2. 组装真实链路(仅 embedding 为确定性伪向量)
|
||||
qdrant = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await qdrant.ensure_collections()
|
||||
ollama = RecordingOllama(ollama_inner)
|
||||
taxonomy = load_taxonomy()
|
||||
embedding = DeterministicEmbedding()
|
||||
ingester = Ingester(
|
||||
summarizer=Summarizer(ollama=ollama), # type: ignore[arg-type]
|
||||
classifier=Classifier(ollama=ollama, taxonomy=taxonomy), # type: ignore[arg-type]
|
||||
chunker=Chunker(),
|
||||
embedding=embedding,
|
||||
sparse=SparseEncoder(),
|
||||
qdrant=qdrant,
|
||||
)
|
||||
retriever = Retriever(
|
||||
qdrant=qdrant,
|
||||
query_parser=QueryParser(ollama=ollama, taxonomy=taxonomy), # type: ignore[arg-type]
|
||||
embedding=embedding,
|
||||
sparse_encoder=SparseEncoder(),
|
||||
)
|
||||
|
||||
# 3. 入库
|
||||
t0 = time.perf_counter()
|
||||
result = await ingester.ingest(doc)
|
||||
t_ingest = time.perf_counter() - t0
|
||||
print(f"\n[入库] 耗时 {_fmt(t_ingest)}")
|
||||
print(f" 总结层级: {result.summary.level.value}")
|
||||
print(f" L1 总结: {result.summary.l1_summary[:100]}")
|
||||
print(f" L2 大纲: {(result.summary.l2_outline or '(None,2.5 级回退)')[:100]}")
|
||||
print(f" 分类: {result.category} (置信度 {result.category_confidence:.2f}),tags={result.tags}")
|
||||
print(f" chunks 数: {result.chunks_count}")
|
||||
|
||||
# 4. 检索
|
||||
t0 = time.perf_counter()
|
||||
response = await retriever.search(SearchRequest(query=query))
|
||||
t_search = time.perf_counter() - t0
|
||||
print(f"\n[检索] 耗时 {_fmt(t_search)}")
|
||||
print(f" routed_categories={response.routed_categories},fallback={response.fallback}")
|
||||
print(f" hits 数: {len(response.hits)}")
|
||||
if response.hits:
|
||||
hit = response.hits[0]
|
||||
print(f" 首条 hit: section_path={hit.section_path!r},score={hit.score:.4f}")
|
||||
print(f" 首条 hit.doc_summary 前 50 字: {hit.doc_summary[:50]!r}")
|
||||
print(f" 首条 hit.text 前 50 字: {hit.text[:50]!r}")
|
||||
|
||||
# 5. r1 模型原始输出观察(thinking / JSON 表现)
|
||||
print("\n[Ollama 原始输出摘录](r1 推理模型的 JSON 表现现场观察)")
|
||||
for kind, raw in ollama.records:
|
||||
excerpt = raw.strip().replace("\n", " ")[:200]
|
||||
has_think = "<think>" in raw
|
||||
print(f" - {kind}: 长度 {len(raw)},含 <think>={has_think},输出摘录: {excerpt!r}")
|
||||
|
||||
print(f"\n[总耗时] {_fmt(time.perf_counter() - started)}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -0,0 +1,329 @@
|
||||
"""管理闭环集成测试(Spec Task 7:集成验证)
|
||||
|
||||
两部分覆盖,均不修改 app/ 与 scripts/ 实现代码:
|
||||
|
||||
1. TestAdminClosedLoop:真实内存 Qdrant(location=":memory:")+ TestClient 走完整管理闭环
|
||||
列表 → 详情 → stats → 删除 → 删除后校验 → 幂等再删。
|
||||
document.py 与 knowledge.py 的 _get_qdrant 单例均 monkeypatch 指向同一内存服务,
|
||||
数据直接经 QdrantService.upsert_* 写入(L1/L2/L3/chunks 四层齐全)。
|
||||
|
||||
2. TestAdminPageContract:读取 app/static/admin.html 文本做静态断言,
|
||||
校验页面 fetch 路径全部落在后端真实路由集合内、页面引用的响应字段名抽样
|
||||
存在于后端模型/路由代码中,防止前后端契约漂移。
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.routing import APIRoute
|
||||
from fastapi.testclient import TestClient
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.api.v1 import knowledge as knowledge_module
|
||||
from app.config import settings
|
||||
from app.main import app
|
||||
from app.services.qdrant import (
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
DIM = settings.embedding_dimension
|
||||
|
||||
# 文档 A:类目「技术文档」,L2 x2 / L3 x3 / chunks x4
|
||||
DOC_A = "doc-a-tech"
|
||||
DOC_A_TITLE = "安装指南"
|
||||
DOC_A_SUMMARY = "安装指南的一句话总结"
|
||||
DOC_A_CATEGORY = "技术文档"
|
||||
DOC_A_TAGS = ["安装", "运维"]
|
||||
DOC_A_L2_PATHS = ["安装指南", "安装指南 / 环境准备"]
|
||||
DOC_A_L3_PATHS = ["安装指南", "安装指南 / 环境准备", "安装指南 / 安装步骤"]
|
||||
DOC_A_CHUNKS = 4
|
||||
|
||||
# 文档 B:类目「uncategorized」,L2 x1 / L3 x2 / chunks x3
|
||||
DOC_B = "doc-b-misc"
|
||||
DOC_B_TITLE = "随手记"
|
||||
DOC_B_SUMMARY = "无法归类的随手记录"
|
||||
DOC_B_CATEGORY = "uncategorized"
|
||||
DOC_B_L2_PATHS = ["随手记"]
|
||||
DOC_B_L3_PATHS = ["随手记", "随手记 / 杂项"]
|
||||
DOC_B_CHUNKS = 3
|
||||
|
||||
# admin.html 与项目根路径
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
_ADMIN_HTML_PATH = _PROJECT_ROOT / "app" / "static" / "admin.html"
|
||||
|
||||
# 页面引用的响应字段抽样 -> 应包含该字段的后端模型/路由文件(相对项目根)
|
||||
_CONTRACT_FIELD_FILES = {
|
||||
"routed_categories": "app/models/search.py",
|
||||
"fallback": "app/models/search.py",
|
||||
"chunks_count": "app/models/document.py",
|
||||
"deleted_total": "app/api/v1/document.py",
|
||||
"next_offset": "app/api/v1/document.py",
|
||||
"uncategorized_count": "app/api/v1/knowledge.py",
|
||||
"documents_total": "app/api/v1/knowledge.py",
|
||||
}
|
||||
|
||||
|
||||
def _backend_paths() -> set[str]:
|
||||
"""展平 app.routes,收集后端真实路由路径
|
||||
|
||||
兼容两种挂载形态:APIRoute 直接平铺,以及 include_router 产生的
|
||||
嵌套代理对象(其路由在 original_router.routes 上)。
|
||||
"""
|
||||
paths: set[str] = set()
|
||||
stack: list[Any] = list(app.routes)
|
||||
while stack:
|
||||
route = stack.pop()
|
||||
if isinstance(route, APIRoute):
|
||||
paths.add(route.path)
|
||||
continue
|
||||
stack.extend(getattr(route, "routes", []) or [])
|
||||
inner = getattr(route, "original_router", None)
|
||||
if inner is not None:
|
||||
stack.extend(getattr(inner, "routes", []))
|
||||
return paths
|
||||
|
||||
|
||||
def _dense(seed: float) -> list[float]:
|
||||
"""构造确定性 dense 向量:前 4 维取特征值,便于区分不同点"""
|
||||
vec = [0.0] * DIM
|
||||
vec[0] = seed
|
||||
vec[1] = 1.0 - seed
|
||||
vec[2] = seed * 0.5
|
||||
vec[3] = 0.1
|
||||
return vec
|
||||
|
||||
|
||||
def _node(doc_id: str, section_path: str, category: str, tags: list[str], seed: float) -> dict[str, Any]:
|
||||
"""构造一个 L2/L3 大纲节点"""
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"section_path": section_path,
|
||||
"text": f"{section_path} 的节点内容",
|
||||
"category": category,
|
||||
"tags": tags,
|
||||
"dense_vector": _dense(seed),
|
||||
}
|
||||
|
||||
|
||||
def _chunk(doc_id: str, index: int, title: str, category: str, tags: list[str]) -> dict[str, Any]:
|
||||
"""构造一个原文 chunk"""
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"chunk_index": index,
|
||||
"text": f"{doc_id} 的第 {index} 个 chunk",
|
||||
"section_path": f"section-{index}",
|
||||
"title": title,
|
||||
"category": category,
|
||||
"tags": tags,
|
||||
"dense_vector": _dense(0.5 + index * 0.05),
|
||||
}
|
||||
|
||||
|
||||
async def _seed_documents(service: QdrantService) -> None:
|
||||
"""写入两篇四层结构完整的文档:A「技术文档」、B「uncategorized」"""
|
||||
await service.upsert_l1(
|
||||
doc_id=DOC_A,
|
||||
title=DOC_A_TITLE,
|
||||
summary=DOC_A_SUMMARY,
|
||||
category=DOC_A_CATEGORY,
|
||||
tags=DOC_A_TAGS,
|
||||
dense_vector=_dense(0.9),
|
||||
)
|
||||
await service.upsert_nodes(
|
||||
COLLECTION_L2, [_node(DOC_A, path, DOC_A_CATEGORY, DOC_A_TAGS, 0.8) for path in DOC_A_L2_PATHS]
|
||||
)
|
||||
await service.upsert_nodes(
|
||||
COLLECTION_L3, [_node(DOC_A, path, DOC_A_CATEGORY, DOC_A_TAGS, 0.7) for path in DOC_A_L3_PATHS]
|
||||
)
|
||||
await service.upsert_chunks(
|
||||
[_chunk(DOC_A, i, DOC_A_TITLE, DOC_A_CATEGORY, DOC_A_TAGS) for i in range(DOC_A_CHUNKS)]
|
||||
)
|
||||
|
||||
await service.upsert_l1(
|
||||
doc_id=DOC_B,
|
||||
title=DOC_B_TITLE,
|
||||
summary=DOC_B_SUMMARY,
|
||||
category=DOC_B_CATEGORY,
|
||||
tags=[],
|
||||
dense_vector=_dense(0.6),
|
||||
)
|
||||
await service.upsert_nodes(
|
||||
COLLECTION_L2, [_node(DOC_B, path, DOC_B_CATEGORY, [], 0.55) for path in DOC_B_L2_PATHS]
|
||||
)
|
||||
await service.upsert_nodes(
|
||||
COLLECTION_L3, [_node(DOC_B, path, DOC_B_CATEGORY, [], 0.5) for path in DOC_B_L3_PATHS]
|
||||
)
|
||||
await service.upsert_chunks([_chunk(DOC_B, i, DOC_B_TITLE, DOC_B_CATEGORY, []) for i in range(DOC_B_CHUNKS)])
|
||||
|
||||
|
||||
class TestAdminClosedLoop:
|
||||
"""文档管理 API + stats 的完整闭环(真实内存 Qdrant + TestClient)"""
|
||||
|
||||
async def test_document_admin_closed_loop(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
service = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await service.ensure_collections()
|
||||
await _seed_documents(service)
|
||||
|
||||
# lifespan 的集合初始化替换为空操作;两个路由模块的 _get_qdrant 指向同一内存服务
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
monkeypatch.setattr(document_module, "_get_qdrant", lambda: service)
|
||||
monkeypatch.setattr(knowledge_module, "_get_qdrant", lambda: service)
|
||||
|
||||
with TestClient(app) as client:
|
||||
# 1. 列表:items 数=2,字段完整
|
||||
resp = client.get("/api/v1/documents")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert set(data.keys()) == {"items", "next_offset"}
|
||||
items = data["items"]
|
||||
assert len(items) == 2
|
||||
assert data["next_offset"] is None
|
||||
by_id = {item["doc_id"]: item for item in items}
|
||||
assert set(by_id) == {DOC_A, DOC_B}
|
||||
item_a = by_id[DOC_A]
|
||||
assert set(item_a.keys()) == {"doc_id", "title", "category", "tags", "summary"}
|
||||
assert item_a["title"] == DOC_A_TITLE
|
||||
assert item_a["category"] == DOC_A_CATEGORY
|
||||
assert item_a["tags"] == DOC_A_TAGS
|
||||
assert item_a["summary"] == DOC_A_SUMMARY
|
||||
item_b = by_id[DOC_B]
|
||||
assert item_b["category"] == DOC_B_CATEGORY
|
||||
assert item_b["tags"] == []
|
||||
|
||||
# 2. 详情:l1/l2_nodes/l3_nodes/chunks_count 与写入一致
|
||||
resp = client.get(f"/api/v1/documents/{DOC_A}")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()["data"]
|
||||
assert set(data.keys()) == {"l1", "l2_nodes", "l3_nodes", "chunks_count"}
|
||||
l1 = data["l1"]
|
||||
assert l1["doc_id"] == DOC_A
|
||||
assert l1["title"] == DOC_A_TITLE
|
||||
assert l1["text"] == DOC_A_SUMMARY
|
||||
assert l1["category"] == DOC_A_CATEGORY
|
||||
assert l1["tags"] == DOC_A_TAGS
|
||||
assert len(data["l2_nodes"]) == len(DOC_A_L2_PATHS)
|
||||
assert {n["section_path"] for n in data["l2_nodes"]} == set(DOC_A_L2_PATHS)
|
||||
assert len(data["l3_nodes"]) == len(DOC_A_L3_PATHS)
|
||||
assert {n["section_path"] for n in data["l3_nodes"]} == set(DOC_A_L3_PATHS)
|
||||
assert data["chunks_count"] == DOC_A_CHUNKS
|
||||
|
||||
# 3. stats:documents_total=2、类目分布与四层点数正确
|
||||
resp = client.get("/api/v1/knowledge/stats")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()["data"]
|
||||
assert set(data.keys()) == {"collections", "categories", "uncategorized_count", "documents_total"}
|
||||
assert data["documents_total"] == 2
|
||||
assert data["categories"] == {DOC_A_CATEGORY: 1, DOC_B_CATEGORY: 1}
|
||||
assert data["uncategorized_count"] == 1
|
||||
assert data["collections"] == {
|
||||
COLLECTION_L1: 2,
|
||||
COLLECTION_L2: len(DOC_A_L2_PATHS) + len(DOC_B_L2_PATHS),
|
||||
COLLECTION_L3: len(DOC_A_L3_PATHS) + len(DOC_B_L3_PATHS),
|
||||
COLLECTION_CHUNKS: DOC_A_CHUNKS + DOC_B_CHUNKS,
|
||||
}
|
||||
|
||||
# 4. 删除文档 A:deleted_total > 0,各集合删除数与写入一致
|
||||
resp = client.delete(f"/api/v1/documents/{DOC_A}")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert set(data.keys()) == {"doc_id", "deleted", "deleted_total"}
|
||||
assert data["doc_id"] == DOC_A
|
||||
assert data["deleted"] == {
|
||||
COLLECTION_L1: 1,
|
||||
COLLECTION_L2: len(DOC_A_L2_PATHS),
|
||||
COLLECTION_L3: len(DOC_A_L3_PATHS),
|
||||
COLLECTION_CHUNKS: DOC_A_CHUNKS,
|
||||
}
|
||||
assert data["deleted_total"] == 1 + len(DOC_A_L2_PATHS) + len(DOC_A_L3_PATHS) + DOC_A_CHUNKS > 0
|
||||
|
||||
# 5. 删除后校验:详情 1004、stats 只剩 1 篇、文档 B 不受影响
|
||||
resp = client.get(f"/api/v1/documents/{DOC_A}")
|
||||
body = resp.json()
|
||||
assert body["code"] == 1004
|
||||
assert body["data"] is None
|
||||
|
||||
resp = client.get("/api/v1/knowledge/stats")
|
||||
data = resp.json()["data"]
|
||||
assert data["documents_total"] == 1
|
||||
assert data["categories"] == {DOC_B_CATEGORY: 1}
|
||||
assert data["uncategorized_count"] == 1
|
||||
assert data["collections"] == {
|
||||
COLLECTION_L1: 1,
|
||||
COLLECTION_L2: len(DOC_B_L2_PATHS),
|
||||
COLLECTION_L3: len(DOC_B_L3_PATHS),
|
||||
COLLECTION_CHUNKS: DOC_B_CHUNKS,
|
||||
}
|
||||
|
||||
resp = client.get(f"/api/v1/documents/{DOC_B}")
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["l1"]["doc_id"] == DOC_B
|
||||
assert len(data["l2_nodes"]) == len(DOC_B_L2_PATHS)
|
||||
assert len(data["l3_nodes"]) == len(DOC_B_L3_PATHS)
|
||||
assert data["chunks_count"] == DOC_B_CHUNKS
|
||||
|
||||
resp = client.get("/api/v1/documents")
|
||||
items = resp.json()["data"]["items"]
|
||||
assert [item["doc_id"] for item in items] == [DOC_B]
|
||||
|
||||
# 6. 幂等再删:code=0 且 deleted_total=0,各集合删除数全 0
|
||||
resp = client.delete(f"/api/v1/documents/{DOC_A}")
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
all_collections = (COLLECTION_L1, COLLECTION_L2, COLLECTION_L3, COLLECTION_CHUNKS)
|
||||
assert data["deleted"] == {name: 0 for name in all_collections}
|
||||
assert data["deleted_total"] == 0
|
||||
|
||||
|
||||
class TestAdminPageContract:
|
||||
"""admin.html 与后端 API 的契约一致性(静态断言,防契约漂移)"""
|
||||
|
||||
def test_fetch_paths_in_backend_routes(self) -> None:
|
||||
"""HTML 中所有 /api/v1/ URL 字面量都在后端真实路由集合内"""
|
||||
html = _ADMIN_HTML_PATH.read_text(encoding="utf-8")
|
||||
# 提取所有以 /api/v1/ 开头的字符串字面量(覆盖 fetch()/api() 调用与路径变量)
|
||||
literals = set(re.findall(r"""["'](/api/v1/[^"']*)["']""", html))
|
||||
|
||||
# 页面至少应覆盖这 5 条契约路径(缺失说明提取失效或页面能力回退)
|
||||
expected = {
|
||||
"/api/v1/search",
|
||||
"/api/v1/documents",
|
||||
"/api/v1/documents/", # 详情/删除:拼接 doc_id 的前缀形式
|
||||
"/api/v1/knowledge/stats",
|
||||
"/api/v1/knowledge/categories",
|
||||
}
|
||||
assert expected <= literals
|
||||
|
||||
backend_paths = _backend_paths()
|
||||
for literal in sorted(literals):
|
||||
path = literal.split("?", 1)[0] # 去掉查询串
|
||||
if path in backend_paths:
|
||||
continue
|
||||
# 以 / 结尾的前缀(如 /api/v1/documents/)应对应带路径参数的路由
|
||||
if path.endswith("/") and any(p.startswith(path) for p in backend_paths):
|
||||
continue
|
||||
pytest.fail(f"页面 fetch 路径不在后端路由集合内: {literal}")
|
||||
|
||||
@pytest.mark.parametrize(("field", "source_file"), sorted(_CONTRACT_FIELD_FILES.items()))
|
||||
def test_response_fields_in_backend_code(self, field: str, source_file: str) -> None:
|
||||
"""页面引用的响应字段名:既在 HTML 中出现,也在后端模型/路由代码中存在"""
|
||||
html = _ADMIN_HTML_PATH.read_text(encoding="utf-8")
|
||||
assert field in html, f"契约字段未在页面中引用: {field}"
|
||||
source = (_PROJECT_ROOT / source_file).read_text(encoding="utf-8")
|
||||
assert field in source, f"契约字段未在后端代码中定义: {field} ({source_file})"
|
||||
@@ -0,0 +1,125 @@
|
||||
"""管理页面挂载测试(GET /admin)
|
||||
|
||||
覆盖:
|
||||
1. 路由存在性:200 + content-type 为 text/html
|
||||
2. 五区块可识别标记(文案与 section id)
|
||||
3. 零外部依赖:无 http(s) 外链资源、无 CDN 引用
|
||||
4. fetch 调用路径与后端 API 契约一致
|
||||
5. 删除操作的 confirm() 二次确认逻辑
|
||||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.main import app
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def admin_html(client: TestClient) -> str:
|
||||
"""请求 /admin 并返回 HTML 文本(前置断言 200)"""
|
||||
resp = client.get("/admin")
|
||||
assert resp.status_code == 200
|
||||
return resp.text
|
||||
|
||||
|
||||
def test_admin_page_ok(client: TestClient) -> None:
|
||||
"""GET /admin → 200,content-type 含 text/html"""
|
||||
resp = client.get("/admin")
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert "text/html" in resp.headers["content-type"]
|
||||
|
||||
|
||||
def test_admin_page_has_five_sections(admin_html: str) -> None:
|
||||
"""HTML 含五个区块的文案与对应 section id"""
|
||||
for section_id in (
|
||||
"section-overview",
|
||||
"section-docs",
|
||||
"section-ingest",
|
||||
"section-search",
|
||||
"section-categories",
|
||||
):
|
||||
assert section_id in admin_html
|
||||
for label in ("概览", "文档管理", "文档入库", "检索测试台", "类目列表"):
|
||||
assert label in admin_html
|
||||
|
||||
|
||||
def test_admin_page_no_external_resources(admin_html: str) -> None:
|
||||
"""零外部依赖:无 src/href http(s) 外链,无 CDN 引用"""
|
||||
assert re.search(r'src=["\']https?://', admin_html) is None
|
||||
assert re.search(r'href=["\']https?://', admin_html) is None
|
||||
assert re.search(r"//cdn", admin_html) is None
|
||||
|
||||
|
||||
def test_admin_page_fetch_paths(admin_html: str) -> None:
|
||||
"""fetch 调用路径与后端 API 契约一致"""
|
||||
for path in (
|
||||
"/api/v1/search",
|
||||
"/api/v1/documents",
|
||||
"/api/v1/knowledge/stats",
|
||||
"/api/v1/knowledge/categories",
|
||||
):
|
||||
assert path in admin_html
|
||||
|
||||
|
||||
def _collect_route_paths(routes: list) -> set[str]:
|
||||
"""递归收集路由路径(FastAPI 0.140+ include_router 包装为 _IncludedRouter)"""
|
||||
paths: set[str] = set()
|
||||
for route in routes:
|
||||
path = getattr(route, "path", None)
|
||||
if path:
|
||||
paths.add(path)
|
||||
sub_router = getattr(route, "original_router", None)
|
||||
if sub_router is not None:
|
||||
paths |= _collect_route_paths(sub_router.routes)
|
||||
return paths
|
||||
|
||||
|
||||
def test_admin_page_fetch_paths_in_real_routes(admin_html: str) -> None:
|
||||
"""页面 api() 调用路径均在真实后端路由集合内(含新增的 tasks 路径)"""
|
||||
route_paths = _collect_route_paths(app.routes)
|
||||
assert "/api/v1/documents/tasks/{task_id}" in route_paths
|
||||
|
||||
prefixes = set(re.findall(r'api\("([^"?]+)', admin_html))
|
||||
assert prefixes, "页面应包含 api() 调用"
|
||||
for prefix in prefixes:
|
||||
assert any(path == prefix or path.startswith(prefix) for path in route_paths), (
|
||||
f"页面 fetch 路径 {prefix} 不在后端路由集合内"
|
||||
)
|
||||
|
||||
|
||||
def test_admin_page_ingest_polling(admin_html: str) -> None:
|
||||
"""入库区块:异步任务轮询逻辑标记"""
|
||||
# 轮询任务状态端点
|
||||
assert "/api/v1/documents/tasks/" in admin_html
|
||||
# 状态中文映射
|
||||
for text in ("排队中", "总结中", "分类中", "向量化中", "写入中", "完成", "失败"):
|
||||
assert text in admin_html
|
||||
# 提交确认与超时提示
|
||||
assert "任务已提交" in admin_html
|
||||
assert "任务仍在进行,可稍后凭 task_id 查询" in admin_html
|
||||
# 5 分钟超时:150 次 × 2s
|
||||
assert "150" in admin_html
|
||||
# 防重复提交:轮询中禁用提交按钮
|
||||
assert "disabled" in admin_html
|
||||
|
||||
|
||||
def test_admin_page_has_confirm(admin_html: str) -> None:
|
||||
"""删除操作包含 confirm() 二次确认逻辑"""
|
||||
assert "confirm(" in admin_html
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Chunker 切分器的单元测试"""
|
||||
|
||||
from app.core.chunker import Chunker
|
||||
|
||||
|
||||
class TestStructuredChunking:
|
||||
"""结构化文档按标题树切分"""
|
||||
|
||||
def test_split_by_sections_and_section_path(self):
|
||||
"""按 section 切分,section_path 为祖先标题链"""
|
||||
text = (
|
||||
"# 安装指南\n这是简介内容,介绍安装流程。\n\n"
|
||||
"## 环境准备\n准备 Python 环境与依赖。\n\n"
|
||||
"## 安装步骤\n执行安装命令完成部署。"
|
||||
)
|
||||
chunks = Chunker(max_chars=30).chunk(text, "doc-1")
|
||||
|
||||
assert [c.chunk_index for c in chunks] == [0, 1, 2]
|
||||
assert [c.section_path for c in chunks] == ["安装指南", "安装指南 / 环境准备", "安装指南 / 安装步骤"]
|
||||
# chunk 文本含所属标题行本身
|
||||
assert chunks[0].text.startswith("# 安装指南")
|
||||
assert chunks[1].text.startswith("## 环境准备")
|
||||
assert chunks[2].text.startswith("## 安装步骤")
|
||||
assert all(c.doc_id == "doc-1" for c in chunks)
|
||||
|
||||
def test_preamble_before_first_heading(self):
|
||||
"""首个标题前的引导正文归入空 section_path 的 chunk"""
|
||||
text = "无标题的引导内容,描述文档主题。\n# 第一章\n章节正文内容。"
|
||||
chunks = Chunker(max_chars=20).chunk(text, "doc-1")
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0].section_path == ""
|
||||
assert chunks[0].text == "无标题的引导内容,描述文档主题。"
|
||||
assert chunks[1].section_path == "第一章"
|
||||
|
||||
def test_sibling_sections_reset_path(self):
|
||||
"""同级标题弹栈,section_path 不含上一分支"""
|
||||
text = "# 甲\n内容甲。\n## 甲一\n内容甲一。\n# 乙\n内容乙。"
|
||||
chunks = Chunker(max_chars=20).chunk(text, "doc-1")
|
||||
|
||||
assert [c.section_path for c in chunks] == ["甲", "甲 / 甲一", "乙"]
|
||||
|
||||
|
||||
class TestLongSectionSplitting:
|
||||
"""超长 section 二次切分"""
|
||||
|
||||
def test_long_section_split_by_paragraphs(self):
|
||||
"""超长 section 按空行段落累加切分,同 section 的 chunk 共享 section_path"""
|
||||
para1 = "第一段内容," * 6 # 36 字符
|
||||
para2 = "第二段内容," * 6
|
||||
text = f"# 大章节\n{para1}\n\n{para2}"
|
||||
chunks = Chunker(max_chars=60).chunk(text, "doc-1")
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert [c.chunk_index for c in chunks] == [0, 1]
|
||||
assert all(c.section_path == "大章节" for c in chunks)
|
||||
assert all(len(c.text) <= 60 for c in chunks)
|
||||
assert chunks[0].text.startswith("# 大章节")
|
||||
assert para1 in chunks[0].text
|
||||
assert para2 in chunks[1].text
|
||||
|
||||
def test_long_paragraph_hard_split(self):
|
||||
"""单段落仍超长时按 max_chars 硬切"""
|
||||
long_para = "长" * 150
|
||||
text = f"# 章节\n{long_para}"
|
||||
chunks = Chunker(max_chars=50).chunk(text, "doc-1")
|
||||
|
||||
# section 文本为 "# 章节\n" + 150 字 = 156 字符,硬切为 4 段
|
||||
assert len(chunks) == 4
|
||||
assert [c.chunk_index for c in chunks] == [0, 1, 2, 3]
|
||||
assert all(len(c.text) <= 50 for c in chunks)
|
||||
assert all(c.section_path == "章节" for c in chunks)
|
||||
assert chunks[0].text.startswith("# 章节")
|
||||
|
||||
|
||||
class TestPlainTextChunking:
|
||||
"""无结构文本按段落切分"""
|
||||
|
||||
def test_split_by_paragraphs(self):
|
||||
"""按空行段落累加切分,section_path 为空"""
|
||||
para = "段落内容," * 5 # 25 字符
|
||||
text = f"{para}\n\n{para}\n\n{para}"
|
||||
chunks = Chunker(max_chars=55).chunk(text, "doc-1")
|
||||
|
||||
# 两段累加 52 字符 ≤ 55,再加一段超限,故切为 2 块
|
||||
assert len(chunks) == 2
|
||||
assert [c.chunk_index for c in chunks] == [0, 1]
|
||||
assert all(c.section_path == "" for c in chunks)
|
||||
assert all(len(c.text) <= 55 for c in chunks)
|
||||
|
||||
def test_long_plain_text_hard_split(self):
|
||||
"""无结构单段落超长时硬切"""
|
||||
text = "字" * 120
|
||||
chunks = Chunker(max_chars=50).chunk(text, "doc-1")
|
||||
|
||||
assert len(chunks) == 3
|
||||
assert all(len(c.text) <= 50 for c in chunks)
|
||||
assert all(c.section_path == "" for c in chunks)
|
||||
|
||||
|
||||
class TestShortAndEmptyText:
|
||||
"""短文本与空文本"""
|
||||
|
||||
def test_short_text_single_chunk(self):
|
||||
"""全文不超过 max_chars 时整篇单 chunk"""
|
||||
chunks = Chunker(max_chars=100).chunk("# 标题\n短文本内容", "doc-1")
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].chunk_index == 0
|
||||
assert chunks[0].text == "# 标题\n短文本内容"
|
||||
assert chunks[0].section_path == ""
|
||||
|
||||
def test_empty_text_returns_empty(self):
|
||||
assert Chunker(max_chars=100).chunk("", "doc-1") == []
|
||||
assert Chunker(max_chars=100).chunk(" \n ", "doc-1") == []
|
||||
|
||||
|
||||
class TestChunkIndexContinuity:
|
||||
"""chunk_index 跨 section 连续递增"""
|
||||
|
||||
def test_indices_continuous_across_sections(self):
|
||||
para = "内容段落," * 6 # 30 字符
|
||||
text = f"# 第一章\n{para}\n\n{para}\n# 第二章\n{para}\n\n{para}"
|
||||
chunks = Chunker(max_chars=50).chunk(text, "doc-1")
|
||||
|
||||
# 每个 section(标题行 + 两段)超 50,各切为 2 块,共 4 块
|
||||
assert len(chunks) == 4
|
||||
assert [c.chunk_index for c in chunks] == [0, 1, 2, 3]
|
||||
assert chunks[0].section_path == chunks[1].section_path == "第一章"
|
||||
assert chunks[2].section_path == chunks[3].section_path == "第二章"
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Classifier 文档分类的单元测试(mock OllamaClient,不真实联网)"""
|
||||
|
||||
import json
|
||||
|
||||
from app.core.classifier import Classifier
|
||||
from app.models.knowledge import UNCATEGORIZED, TaxonomyCategory
|
||||
|
||||
|
||||
def _taxonomy() -> list[TaxonomyCategory]:
|
||||
"""测试用 taxonomy 类目集"""
|
||||
return [
|
||||
TaxonomyCategory(name="技术文档", description="架构设计、API 文档、开发规范等技术资料"),
|
||||
TaxonomyCategory(name="产品手册", description="产品功能介绍、使用说明"),
|
||||
TaxonomyCategory(name="财务行政", description="财务制度、报销流程、行政通知"),
|
||||
TaxonomyCategory(name=UNCATEGORIZED, description="无法归入其他类目的文档"),
|
||||
]
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""返回固定响应的假 OllamaClient,记录调用参数"""
|
||||
|
||||
def __init__(self, response: str) -> None:
|
||||
self.response = response
|
||||
self.calls: list[dict] = []
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
self.calls.append({"prompt": prompt, "json_mode": json_mode})
|
||||
return self.response
|
||||
|
||||
|
||||
def _make_classifier(response: str) -> Classifier:
|
||||
return Classifier(ollama=FakeOllama(response), taxonomy=_taxonomy()) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class TestClassify:
|
||||
"""Classifier.classify 的正常解析与容错"""
|
||||
|
||||
async def test_classify_valid_json(self):
|
||||
"""正常 JSON 响应 → 正确解析主类目/tags/confidence
|
||||
|
||||
tags 清洗规则:剔除主类目名、最多保留 3 个。
|
||||
"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"main_category": "技术文档",
|
||||
# "技术文档" 与主类目重复应被剔除;超出 3 个应被截断
|
||||
"tags": ["架构", "技术文档", "API", "部署", "运维"],
|
||||
"confidence": 0.9,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
classifier = _make_classifier(response)
|
||||
|
||||
result = await classifier.classify("本文介绍系统架构设计。", title="架构文档")
|
||||
|
||||
assert result.main_category == "技术文档"
|
||||
assert result.tags == ["架构", "API", "部署"]
|
||||
assert result.confidence == 0.9
|
||||
|
||||
async def test_classify_uses_json_mode_and_prompt_contains_taxonomy(self):
|
||||
"""以 json_mode 调用 LLM,prompt 包含全部类目名与 uncategorized 用途说明"""
|
||||
classifier = _make_classifier('{"main_category": "财务行政", "tags": [], "confidence": 0.8}')
|
||||
|
||||
result = await classifier.classify("报销流程说明", title="")
|
||||
|
||||
assert result.main_category == "财务行政"
|
||||
ollama = classifier.ollama
|
||||
assert ollama.calls[0]["json_mode"] is True
|
||||
prompt = ollama.calls[0]["prompt"]
|
||||
assert "技术文档" in prompt and "产品手册" in prompt and "财务行政" in prompt
|
||||
assert UNCATEGORIZED in prompt
|
||||
assert "跨多个类目" in prompt
|
||||
assert "报销流程说明" in prompt
|
||||
|
||||
async def test_classify_json_with_surrounding_noise(self):
|
||||
"""响应带前后多余文本 → 正则提取 {...} 块成功"""
|
||||
response = (
|
||||
"好的,分类结果如下:\n"
|
||||
'{"main_category": "产品手册", "tags": ["使用说明"], "confidence": 0.75}\n'
|
||||
"以上是分类结果。"
|
||||
)
|
||||
classifier = _make_classifier(response)
|
||||
|
||||
result = await classifier.classify("产品功能使用说明")
|
||||
|
||||
assert result.main_category == "产品手册"
|
||||
assert result.tags == ["使用说明"]
|
||||
assert result.confidence == 0.75
|
||||
|
||||
async def test_classify_non_json_falls_back_to_uncategorized(self):
|
||||
"""完全非 JSON 响应 → 归 uncategorized,tags 为空,confidence=0.0"""
|
||||
classifier = _make_classifier("抱歉,我无法完成分类。")
|
||||
|
||||
result = await classifier.classify("一段总结")
|
||||
|
||||
assert result.main_category == UNCATEGORIZED
|
||||
assert result.tags == []
|
||||
assert result.confidence == 0.0
|
||||
|
||||
async def test_classify_unknown_category_falls_back(self):
|
||||
"""类目名不在 taxonomy → 归 uncategorized,confidence=0.0"""
|
||||
classifier = _make_classifier('{"main_category": "不存在的类目", "tags": ["x"], "confidence": 0.9}')
|
||||
|
||||
result = await classifier.classify("一段总结")
|
||||
|
||||
assert result.main_category == UNCATEGORIZED
|
||||
assert result.tags == []
|
||||
assert result.confidence == 0.0
|
||||
|
||||
async def test_classify_low_confidence_soft_recall(self):
|
||||
"""置信度低于阈值(默认 0.6)→ 主类目归 uncategorized,候选类目名保留进 tags,confidence 保留原值"""
|
||||
response = json.dumps(
|
||||
{"main_category": "产品手册", "tags": ["手册"], "confidence": 0.4},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
classifier = _make_classifier(response)
|
||||
|
||||
result = await classifier.classify("介于产品和运营之间的内容")
|
||||
|
||||
assert result.main_category == UNCATEGORIZED
|
||||
# 候选类目名插入 tags 首位,供检索侧软召回
|
||||
assert result.tags == ["产品手册", "手册"]
|
||||
assert result.confidence == 0.4
|
||||
|
||||
async def test_classify_low_confidence_uncategorized_not_duplicated_in_tags(self):
|
||||
"""低置信且候选本身为 uncategorized → 不把 uncategorized 塞进 tags"""
|
||||
classifier = _make_classifier('{"main_category": "uncategorized", "tags": [], "confidence": 0.3}')
|
||||
|
||||
result = await classifier.classify("杂项内容")
|
||||
|
||||
assert result.main_category == UNCATEGORIZED
|
||||
assert result.tags == []
|
||||
assert result.confidence == 0.3
|
||||
@@ -0,0 +1,171 @@
|
||||
"""文档管理 API 测试(QdrantService 为 mock,不真实联网)"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.main import app
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
class FakeQdrant:
|
||||
"""假 QdrantService:返回固定数据或抛出固定异常"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scroll_result: tuple[list[dict[str, Any]], str | None] | None = None,
|
||||
detail: dict[str, Any] | None = None,
|
||||
deleted: dict[str, int] | None = None,
|
||||
error: Exception | None = None,
|
||||
) -> None:
|
||||
self.scroll_result = scroll_result if scroll_result is not None else ([], None)
|
||||
self.detail = detail
|
||||
self.deleted = deleted if deleted is not None else {}
|
||||
self.error = error
|
||||
|
||||
async def scroll_l1(self, limit: int = 20, offset: str | None = None) -> tuple[list[dict[str, Any]], str | None]:
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.scroll_result
|
||||
|
||||
async def get_doc_detail(self, doc_id: str) -> dict[str, Any] | None:
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.detail
|
||||
|
||||
async def delete_by_doc_id(self, doc_id: str) -> dict[str, int]:
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.deleted
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
def _install_fake(monkeypatch: pytest.MonkeyPatch, fake: FakeQdrant) -> None:
|
||||
"""将 _get_qdrant 单例替换为假服务"""
|
||||
monkeypatch.setattr(document_module, "_get_qdrant", lambda: fake)
|
||||
|
||||
|
||||
def _make_item(doc_id: str = "doc-1") -> dict[str, Any]:
|
||||
"""构造固定的 L1 列表项"""
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"title": "标题",
|
||||
"category": "技术文档",
|
||||
"tags": ["API"],
|
||||
"summary": "一句话总结",
|
||||
}
|
||||
|
||||
|
||||
def test_list_documents_success(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""列表:正常返回 items 与 next_offset"""
|
||||
fake = FakeQdrant(scroll_result=([_make_item()], "cursor-2"))
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.get("/api/v1/documents", params={"limit": 10, "offset": "cursor-1"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["items"] == [_make_item()]
|
||||
assert data["next_offset"] == "cursor-2"
|
||||
|
||||
|
||||
def test_list_documents_empty(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""列表:空库返回 items=[]、next_offset=None"""
|
||||
fake = FakeQdrant(scroll_result=([], None))
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.get("/api/v1/documents")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"] == {"items": [], "next_offset": None}
|
||||
|
||||
|
||||
def test_get_document_detail(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""详情:存在返回 l1/l2_nodes/l3_nodes/chunks_count"""
|
||||
detail = {
|
||||
"l1": {"doc_id": "doc-1", "title": "标题", "text": "一句话总结"},
|
||||
"l2_nodes": [{"doc_id": "doc-1", "section_path": "1", "text": "大纲节点"}],
|
||||
"l3_nodes": [],
|
||||
"chunks_count": 3,
|
||||
}
|
||||
fake = FakeQdrant(detail=detail)
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.get("/api/v1/documents/doc-1")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"] == detail
|
||||
|
||||
|
||||
def test_get_document_not_found(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""详情:不存在返回 code=1004"""
|
||||
fake = FakeQdrant(detail=None)
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.get("/api/v1/documents/missing")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1004
|
||||
assert body["message"] == "文档不存在"
|
||||
assert body["data"] is None
|
||||
|
||||
|
||||
def test_delete_document_success(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""删除:返回各集合删除数与 deleted_total"""
|
||||
deleted = {"doc_l1": 1, "doc_l2": 2, "doc_l3": 4, "chunks": 6}
|
||||
fake = FakeQdrant(deleted=deleted)
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.delete("/api/v1/documents/doc-1")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["doc_id"] == "doc-1"
|
||||
assert data["deleted"] == deleted
|
||||
assert data["deleted_total"] == 13
|
||||
|
||||
|
||||
def test_delete_document_not_found_idempotent(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""删除:不存在 doc_id 仍 code=0,各集合删除数全 0、deleted_total=0(幂等)"""
|
||||
deleted = {"doc_l1": 0, "doc_l2": 0, "doc_l3": 0, "chunks": 0}
|
||||
fake = FakeQdrant(deleted=deleted)
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.delete("/api/v1/documents/missing")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["deleted"] == deleted
|
||||
assert data["deleted_total"] == 0
|
||||
|
||||
|
||||
def test_list_documents_service_error(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""列表:scroll 抛错返回 code=2000"""
|
||||
fake = FakeQdrant(error=RuntimeError("boom"))
|
||||
_install_fake(monkeypatch, fake)
|
||||
|
||||
resp = client.get("/api/v1/documents")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 2000
|
||||
assert body["data"] is None
|
||||
@@ -0,0 +1,74 @@
|
||||
"""文档入库 API 测试(任务管理器为 FakeManager,不真实联网)"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.main import app
|
||||
from app.models.document import DocumentInput
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
class FakeManager:
|
||||
"""假入库任务管理器:记录 submit 调用并返回固定 task_id"""
|
||||
|
||||
def __init__(self, task_id: str = "task-1") -> None:
|
||||
self.task_id = task_id
|
||||
self.submitted: list[DocumentInput] = []
|
||||
|
||||
async def submit(self, doc: DocumentInput) -> str:
|
||||
self.submitted.append(doc)
|
||||
return self.task_id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
def test_ingest_submit_accepted(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""正常提交:HTTP 202,code=0,data 含 task_id 与 status=pending,文档已登记给管理器"""
|
||||
fake = FakeManager(task_id="task-abc")
|
||||
monkeypatch.setattr(document_module, "_get_task_manager", lambda: fake)
|
||||
|
||||
resp = client.post("/api/v1/documents", json={"text": "正文内容", "title": "标题"})
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["task_id"] == "task-abc"
|
||||
assert data["status"] == "pending"
|
||||
assert [d.title for d in fake.submitted] == ["标题"]
|
||||
|
||||
|
||||
def test_ingest_empty_text(client: TestClient) -> None:
|
||||
"""空 text:路由内校验,返回 code=1001"""
|
||||
resp = client.post("/api/v1/documents", json={"text": ""})
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert body["message"] == "文档内容不能为空"
|
||||
assert body["data"] is None
|
||||
|
||||
|
||||
def test_ingest_no_sync_ingestion_error(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""任务模式下同步路径不再返回 2000:提交即 202,入库失败体现在任务状态中"""
|
||||
fake = FakeManager()
|
||||
monkeypatch.setattr(document_module, "_get_task_manager", lambda: fake)
|
||||
|
||||
resp = client.post("/api/v1/documents", json={"text": "正文内容"})
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"]["status"] == "pending"
|
||||
@@ -0,0 +1,328 @@
|
||||
"""端到端集成测试
|
||||
|
||||
真实 Ingester + Retriever + 内存 Qdrant(location=":memory:")串联分层 RAG 全链路,
|
||||
仅替换两个外部边界:
|
||||
- FakeOllama:按 prompt 内容返回确定性的 L1/L2/L3 总结、分类 JSON、query 解析 JSON
|
||||
- DeterministicEmbedding:稳定哈希生成固定维度向量(无语义,仅保证流程可跑)
|
||||
|
||||
本机无需 Docker / Ollama / Redis,CI 可直接运行。
|
||||
"""
|
||||
|
||||
import json
|
||||
import random
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.api.v1 import search as search_module
|
||||
from app.config import Settings, settings
|
||||
from app.core.chunker import Chunker
|
||||
from app.core.classifier import Classifier
|
||||
from app.core.ingest_tasks import IngestTaskManager, IngestTaskStatus
|
||||
from app.core.ingestion import Ingester
|
||||
from app.core.query_parser import QueryParser
|
||||
from app.core.retriever import Retriever
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.core.summarizer import Summarizer
|
||||
from app.main import app
|
||||
from app.models.document import DocumentInput
|
||||
from app.models.knowledge import load_taxonomy
|
||||
from app.models.search import SearchRequest
|
||||
from app.services.qdrant import (
|
||||
ALL_COLLECTIONS,
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
# FakeOllama 的确定性输出
|
||||
FAKE_L1_SUMMARY = "本文介绍安装指南、环境准备与安装步骤,是一篇集成测试文档。"
|
||||
FAKE_L3_OUTLINE = (
|
||||
"## 安装指南\n安装指南的整体流程说明。\n"
|
||||
"## 环境准备\n环境准备的依赖与注意事项。\n"
|
||||
"## 安装步骤\n安装步骤的命令与验证方法。"
|
||||
)
|
||||
FAKE_L2_HALF = "要点一:短文档的核心通知内容。\n要点二:需要关注的事项。"
|
||||
FAKE_CATEGORY = "技术文档"
|
||||
FAKE_TAGS = ["安装", "运维"]
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""按 prompt 内容返回确定性结果的假 OllamaClient
|
||||
|
||||
覆盖 Summarizer / Classifier / QueryParser 三类调用方;
|
||||
分类与 query 解析的置信度可配置,用于构造路由兜底场景。
|
||||
"""
|
||||
|
||||
def __init__(self, classify_confidence: float = 0.9, query_confidence: float = 0.9) -> None:
|
||||
self.classify_confidence = classify_confidence
|
||||
self.query_confidence = query_confidence
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
if "你是知识库分类助手" in prompt:
|
||||
return json.dumps(
|
||||
{"main_category": FAKE_CATEGORY, "tags": FAKE_TAGS, "confidence": self.classify_confidence},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
if "你是搜索查询分析助手" in prompt:
|
||||
return json.dumps(
|
||||
{
|
||||
"categories": [{"name": FAKE_CATEGORY, "confidence": self.query_confidence}],
|
||||
"rewrite": "安装指南 环境准备 安装步骤",
|
||||
"keywords": ["安装", "环境准备"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
if "请用一句话对以下文档内容进行高度概括" in prompt:
|
||||
return FAKE_L1_SUMMARY
|
||||
if "请提取以下文档的主要章节结构" in prompt:
|
||||
return "1. 主题一\n2. 主题二"
|
||||
if "请对以下文档的每个章节/主题进行详细的内容摘要" in prompt:
|
||||
return FAKE_L3_OUTLINE
|
||||
if "请对以下文档内容进行详细摘要" in prompt:
|
||||
return FAKE_L2_HALF
|
||||
raise AssertionError(f"FakeOllama 收到未识别的 prompt: {prompt[:100]}")
|
||||
|
||||
|
||||
class DeterministicEmbedding:
|
||||
"""稳定哈希伪向量:同文本恒同向量,维度等于 settings.embedding_dimension"""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
vectors: list[list[float]] = []
|
||||
for text in texts:
|
||||
seed = int.from_bytes(sha256(text.encode("utf-8")).digest()[:8], "big")
|
||||
rng = random.Random(seed)
|
||||
vectors.append([rng.random() for _ in range(settings.embedding_dimension)])
|
||||
return vectors
|
||||
|
||||
|
||||
class FakeCache:
|
||||
"""无缓存行为的假 RedisCache(get 恒未命中,set 恒成功)"""
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
return None
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Env:
|
||||
"""一套共享内存 Qdrant 的真实 Ingester + Retriever 环境"""
|
||||
|
||||
qdrant: QdrantService
|
||||
ollama: FakeOllama
|
||||
ingester: Ingester
|
||||
retriever: Retriever
|
||||
|
||||
|
||||
async def _make_env(classify_confidence: float = 0.9, query_confidence: float = 0.9) -> _Env:
|
||||
"""构建集成环境:真实组件 + 内存 Qdrant + FakeOllama + 确定性向量"""
|
||||
qdrant = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await qdrant.ensure_collections()
|
||||
ollama = FakeOllama(classify_confidence=classify_confidence, query_confidence=query_confidence)
|
||||
taxonomy = load_taxonomy()
|
||||
embedding = DeterministicEmbedding()
|
||||
ingester = Ingester(
|
||||
summarizer=Summarizer(ollama=ollama), # type: ignore[arg-type]
|
||||
classifier=Classifier(ollama=ollama, taxonomy=taxonomy), # type: ignore[arg-type]
|
||||
chunker=Chunker(),
|
||||
embedding=embedding,
|
||||
sparse=SparseEncoder(),
|
||||
qdrant=qdrant,
|
||||
)
|
||||
retriever = Retriever(
|
||||
qdrant=qdrant,
|
||||
query_parser=QueryParser(ollama=ollama, taxonomy=taxonomy, cache=FakeCache()), # type: ignore[arg-type]
|
||||
embedding=embedding,
|
||||
sparse_encoder=SparseEncoder(),
|
||||
)
|
||||
return _Env(qdrant=qdrant, ollama=ollama, ingester=ingester, retriever=retriever)
|
||||
|
||||
|
||||
def _structured_doc() -> DocumentInput:
|
||||
"""带 Markdown 标题结构的中文长文档(>500 字符,切出 3 个 section chunk)"""
|
||||
paragraph = "这是章节正文内容,包含足够多的信息量,用于测试切分与向量化流程。" * 20
|
||||
text = f"# 安装指南\n{paragraph}\n\n## 环境准备\n{paragraph}\n\n## 安装步骤\n{paragraph}"
|
||||
return DocumentInput(text=text, title="安装文档")
|
||||
|
||||
|
||||
def _short_doc() -> DocumentInput:
|
||||
"""短文档(< summary_min_text_length),触发 2.5 级回退"""
|
||||
return DocumentInput(text="# 维护通知\n明天凌晨系统维护,请提前保存工作。", title="维护通知")
|
||||
|
||||
|
||||
async def _counts(qdrant: QdrantService) -> dict[str, int]:
|
||||
"""四个集合的点数"""
|
||||
return {name: (await qdrant.client.count(collection_name=name)).count for name in ALL_COLLECTIONS}
|
||||
|
||||
|
||||
class TestIngestIntegration:
|
||||
"""入库链路:四层集合点数与 chunk payload 完整性"""
|
||||
|
||||
async def test_ingest_structured_document(self):
|
||||
"""结构化长文档 → L1=1、L2=标题数、L3/chunks ≥1,chunk payload 字段齐全"""
|
||||
env = await _make_env()
|
||||
doc = _structured_doc()
|
||||
|
||||
result = await env.ingester.ingest(doc)
|
||||
|
||||
# 结果透传
|
||||
assert result.document_id
|
||||
assert result.category == FAKE_CATEGORY
|
||||
assert result.tags == FAKE_TAGS
|
||||
assert result.category_confidence == 0.9
|
||||
assert result.summary.l1_summary == FAKE_L1_SUMMARY
|
||||
assert result.chunks_count >= 1
|
||||
|
||||
counts = await _counts(env.qdrant)
|
||||
assert counts[COLLECTION_L1] == 1
|
||||
assert counts[COLLECTION_L2] == 3 # 标题数:安装指南/环境准备/安装步骤
|
||||
assert counts[COLLECTION_L3] >= 1
|
||||
assert counts[COLLECTION_CHUNKS] == result.chunks_count >= 1
|
||||
|
||||
# chunk payload:doc_summary/category/tags/section_path 齐全且与分类结果一致
|
||||
points, _ = await env.qdrant.client.scroll(
|
||||
collection_name=COLLECTION_CHUNKS, limit=100, with_payload=True
|
||||
)
|
||||
assert len(points) == result.chunks_count
|
||||
expected_paths = {"安装指南", "安装指南 / 环境准备", "安装指南 / 安装步骤"}
|
||||
for point in points:
|
||||
payload = point.payload or {}
|
||||
assert payload["doc_id"] == result.document_id
|
||||
assert payload["doc_summary"] == FAKE_L1_SUMMARY
|
||||
assert payload["category"] == FAKE_CATEGORY
|
||||
assert payload["tags"] == FAKE_TAGS
|
||||
assert payload["section_path"] in expected_paths
|
||||
assert payload["text"] in doc.text
|
||||
|
||||
async def test_ingest_short_document_l2_half(self):
|
||||
"""短文档 → 2.5 级:doc_l2 无该 doc 节点,其余层正常写入"""
|
||||
env = await _make_env()
|
||||
|
||||
result = await env.ingester.ingest(_short_doc())
|
||||
|
||||
assert result.summary.l2_outline is None
|
||||
counts = await _counts(env.qdrant)
|
||||
assert counts[COLLECTION_L1] == 1
|
||||
assert counts[COLLECTION_L2] == 0 # 2.5 级文档无 L2 大纲节点
|
||||
assert counts[COLLECTION_L3] >= 1
|
||||
assert counts[COLLECTION_CHUNKS] >= 1
|
||||
|
||||
async def test_reingest_idempotent(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""同 doc_id 重复入库:幂等覆盖不报错,各层点数不翻倍"""
|
||||
env = await _make_env()
|
||||
doc = _structured_doc()
|
||||
# 固定 doc_id,两次入库写入同一组确定性 point id(uuid5 覆盖)
|
||||
fixed_uuid = uuid.UUID("12345678-1234-5678-1234-567812345678")
|
||||
monkeypatch.setattr("app.core.ingestion.uuid.uuid4", lambda: fixed_uuid)
|
||||
|
||||
first = await env.ingester.ingest(doc)
|
||||
counts_after_first = await _counts(env.qdrant)
|
||||
second = await env.ingester.ingest(doc)
|
||||
counts_after_second = await _counts(env.qdrant)
|
||||
|
||||
assert first.document_id == second.document_id == fixed_uuid.hex
|
||||
assert counts_after_second == counts_after_first
|
||||
assert counts_after_second[COLLECTION_L1] == 1
|
||||
assert counts_after_second[COLLECTION_L2] == 3
|
||||
|
||||
|
||||
class TestSearchIntegration:
|
||||
"""检索链路:正常路由 / 路由兜底 / 空库"""
|
||||
|
||||
async def test_search_normal_query(self):
|
||||
"""正常 query:hits 非空,text 为原文片段,doc_summary 非空,routed_categories 来自分类"""
|
||||
env = await _make_env()
|
||||
doc = _structured_doc()
|
||||
await env.ingester.ingest(doc)
|
||||
|
||||
resp = await env.retriever.search(SearchRequest(query="安装步骤有哪些注意事项?"))
|
||||
|
||||
assert resp.fallback is False
|
||||
assert resp.routed_categories == [FAKE_CATEGORY]
|
||||
assert resp.hits
|
||||
hit = resp.hits[0]
|
||||
assert hit.text in doc.text
|
||||
assert hit.doc_summary == FAKE_L1_SUMMARY
|
||||
assert hit.score > 0
|
||||
|
||||
async def test_search_route_fallback(self):
|
||||
"""query 解析低置信 → 路由兜底:fallback=True 且仍有 hits"""
|
||||
env = await _make_env(query_confidence=0.3)
|
||||
doc = _structured_doc()
|
||||
await env.ingester.ingest(doc)
|
||||
|
||||
resp = await env.retriever.search(SearchRequest(query="随便问点什么"))
|
||||
|
||||
assert resp.fallback is True
|
||||
assert resp.routed_categories == []
|
||||
assert resp.hits
|
||||
assert resp.hits[0].text in doc.text
|
||||
|
||||
async def test_search_empty_db(self):
|
||||
"""空库 search:hits 为空、不报错、fallback=True"""
|
||||
env = await _make_env()
|
||||
|
||||
resp = await env.retriever.search(SearchRequest(query="空库查询"))
|
||||
|
||||
assert resp.hits == []
|
||||
assert resp.fallback is True
|
||||
|
||||
|
||||
class TestApiIntegration:
|
||||
"""API 层冒烟:documents → search → knowledge/categories,统一响应格式 code=0"""
|
||||
|
||||
async def test_api_smoke(self, monkeypatch: pytest.MonkeyPatch):
|
||||
env = await _make_env()
|
||||
doc = _structured_doc()
|
||||
# 入库走异步任务:测试自建纯内存任务管理器(内存 Qdrant + FakeOllama 的 Ingester)
|
||||
manager = IngestTaskManager(env.ingester, None, Settings())
|
||||
|
||||
# lifespan 的 ensure_collections 替换为空操作(真实 Qdrant 不可达时也能启动)
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
# 模块级单例替换为集成环境实例
|
||||
monkeypatch.setattr(document_module, "_task_manager", manager)
|
||||
monkeypatch.setattr(search_module, "_retriever", env.retriever)
|
||||
monkeypatch.setattr(search_module, "get_cache", lambda: FakeCache())
|
||||
|
||||
with TestClient(app) as client:
|
||||
# 入库:202 拿 task_id,等待任务终态后断言入库结果
|
||||
resp_doc = client.post("/api/v1/documents", json={"text": doc.text, "title": doc.title})
|
||||
assert resp_doc.status_code == 202
|
||||
body_doc = resp_doc.json()
|
||||
assert body_doc["code"] == 0
|
||||
assert body_doc["data"]["status"] == "pending"
|
||||
task_id = body_doc["data"]["task_id"]
|
||||
assert task_id
|
||||
|
||||
final = await manager.wait_done(task_id)
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
result = final["result"]
|
||||
assert result["document_id"]
|
||||
assert result["category"] == FAKE_CATEGORY
|
||||
assert result["chunks_count"] >= 1
|
||||
|
||||
# 检索
|
||||
resp_search = client.post("/api/v1/search", json={"query": "安装步骤有哪些注意事项?"})
|
||||
body_search = resp_search.json()
|
||||
assert body_search["code"] == 0
|
||||
assert body_search["data"]["hits"]
|
||||
assert body_search["data"]["routed_categories"] == [FAKE_CATEGORY]
|
||||
|
||||
# 知识分类
|
||||
resp_categories = client.get("/api/v1/knowledge/categories")
|
||||
body_categories = resp_categories.json()
|
||||
assert body_categories["code"] == 0
|
||||
assert body_categories["data"]["count"] > 0
|
||||
@@ -0,0 +1,150 @@
|
||||
"""嵌入服务单元测试(mock httpx/openai,不发起真实网络请求)"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import settings
|
||||
from app.core.embeddings import (
|
||||
EmbeddingService,
|
||||
LocalEmbeddingService,
|
||||
OpenAIEmbeddingService,
|
||||
create_embedding_service,
|
||||
)
|
||||
|
||||
|
||||
def _fake_openai_response(dim: int, n: int) -> MagicMock:
|
||||
"""构造 OpenAI embeddings.create 的假响应"""
|
||||
resp = MagicMock()
|
||||
resp.data = [MagicMock(embedding=[0.1 * (i + 1)] * dim) for i in range(n)]
|
||||
return resp
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
"""httpx.Response 替代品"""
|
||||
|
||||
def __init__(self, data: dict) -> None:
|
||||
self._data = data
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def json(self) -> dict:
|
||||
return self._data
|
||||
|
||||
|
||||
class _FakeAsyncClient:
|
||||
"""httpx.AsyncClient 替代品,记录请求参数"""
|
||||
|
||||
def __init__(self, data: dict, captured: dict, **kwargs) -> None:
|
||||
self._data = data
|
||||
self._captured = captured
|
||||
captured["timeout"] = kwargs.get("timeout")
|
||||
|
||||
async def __aenter__(self) -> "_FakeAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args) -> bool:
|
||||
return False
|
||||
|
||||
async def post(self, url: str, json: dict | None = None) -> _FakeResponse:
|
||||
self._captured["url"] = url
|
||||
self._captured["json"] = json
|
||||
return _FakeResponse(self._data)
|
||||
|
||||
|
||||
class TestOpenAIEmbeddingService:
|
||||
"""OpenAI provider 行为"""
|
||||
|
||||
async def test_embed_batch(self):
|
||||
"""批量嵌入:模型与输入透传,返回与输入等长的向量"""
|
||||
with patch("app.core.embeddings.AsyncOpenAI") as mock_cls:
|
||||
client = mock_cls.return_value
|
||||
client.embeddings.create = AsyncMock(return_value=_fake_openai_response(dim=1536, n=2))
|
||||
service = OpenAIEmbeddingService(api_key="k", base_url="http://x/v1", model="m")
|
||||
vectors = await service.embed(["你好", "world"])
|
||||
|
||||
assert len(vectors) == 2
|
||||
assert all(len(v) == 1536 for v in vectors)
|
||||
client.embeddings.create.assert_awaited_once_with(model="m", input=["你好", "world"])
|
||||
|
||||
async def test_embed_empty(self):
|
||||
"""空列表输入直接返回空列表,不调用 API"""
|
||||
with patch("app.core.embeddings.AsyncOpenAI") as mock_cls:
|
||||
client = mock_cls.return_value
|
||||
client.embeddings.create = AsyncMock()
|
||||
service = OpenAIEmbeddingService(api_key="k", base_url="http://x/v1", model="m")
|
||||
assert await service.embed([]) == []
|
||||
client.embeddings.create.assert_not_called()
|
||||
|
||||
async def test_dimension_mismatch_warns(self):
|
||||
"""返回维度与配置不一致时记录 warning,不抛错"""
|
||||
with (
|
||||
patch("app.core.embeddings.AsyncOpenAI") as mock_cls,
|
||||
patch("app.core.embeddings.logger") as mock_logger,
|
||||
):
|
||||
client = mock_cls.return_value
|
||||
client.embeddings.create = AsyncMock(return_value=_fake_openai_response(dim=8, n=1))
|
||||
service = OpenAIEmbeddingService(api_key="k", base_url="http://x/v1", model="m")
|
||||
vectors = await service.embed(["a"])
|
||||
|
||||
assert len(vectors[0]) == 8
|
||||
mock_logger.warning.assert_called_once()
|
||||
|
||||
|
||||
class TestLocalEmbeddingService:
|
||||
"""Ollama local provider 行为"""
|
||||
|
||||
async def test_embed_batch(self, monkeypatch):
|
||||
"""批量嵌入:POST /api/embed,payload 与 URL 正确,timeout 60s"""
|
||||
captured: dict = {}
|
||||
data = {"embeddings": [[0.1, 0.2], [0.3, 0.4]]}
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: _FakeAsyncClient(data, captured, **kw))
|
||||
|
||||
service = LocalEmbeddingService(base_url="http://localhost:11434/", model="bge-m3")
|
||||
vectors = await service.embed(["你好", "world"])
|
||||
|
||||
assert vectors == [[0.1, 0.2], [0.3, 0.4]]
|
||||
assert captured["url"] == "http://localhost:11434/api/embed"
|
||||
assert captured["json"] == {"model": "bge-m3", "input": ["你好", "world"]}
|
||||
assert captured["timeout"] == 60.0
|
||||
|
||||
async def test_embed_empty(self, monkeypatch):
|
||||
"""空列表输入直接返回空列表,不发起请求"""
|
||||
called = []
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: called.append(kw))
|
||||
|
||||
service = LocalEmbeddingService(base_url="http://localhost:11434", model="bge-m3")
|
||||
assert await service.embed([]) == []
|
||||
assert called == []
|
||||
|
||||
async def test_dimension_mismatch_warns(self, monkeypatch):
|
||||
"""返回维度与配置不一致时记录 warning,不抛错"""
|
||||
data = {"embeddings": [[0.1] * 8]}
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: _FakeAsyncClient(data, {}, **kw))
|
||||
|
||||
with patch("app.core.embeddings.logger") as mock_logger:
|
||||
service = LocalEmbeddingService(base_url="http://localhost:11434", model="bge-m3")
|
||||
vectors = await service.embed(["a"])
|
||||
|
||||
assert len(vectors[0]) == 8
|
||||
mock_logger.warning.assert_called_once()
|
||||
|
||||
|
||||
class TestCreateEmbeddingService:
|
||||
"""工厂函数选择逻辑"""
|
||||
|
||||
def test_provider_openai(self, monkeypatch):
|
||||
monkeypatch.setattr(settings, "embedding_provider", "openai")
|
||||
monkeypatch.setattr(settings, "openai_api_key", "test-key")
|
||||
service = create_embedding_service()
|
||||
assert isinstance(service, OpenAIEmbeddingService)
|
||||
assert isinstance(service, EmbeddingService)
|
||||
|
||||
def test_provider_local(self, monkeypatch):
|
||||
monkeypatch.setattr(settings, "embedding_provider", "local")
|
||||
monkeypatch.setattr(settings, "ollama_embedding_model", "bge-m3")
|
||||
service = create_embedding_service()
|
||||
assert isinstance(service, LocalEmbeddingService)
|
||||
assert service.model == "bge-m3"
|
||||
assert isinstance(service, EmbeddingService)
|
||||
@@ -0,0 +1,94 @@
|
||||
"""评测指标纯函数单元测试(不依赖 Qdrant / Ollama)"""
|
||||
|
||||
import pytest
|
||||
|
||||
from scripts.eval.metrics import (
|
||||
aggregate,
|
||||
entity_recall,
|
||||
precision_at_k,
|
||||
pruning_loss,
|
||||
recall_at_k,
|
||||
routing_f1,
|
||||
)
|
||||
|
||||
|
||||
class TestEntityRecall:
|
||||
def test_all_entities_kept(self) -> None:
|
||||
source = "使用 Qdrant 存储 1536 维向量,模型 bge-m3,见《部署手册》第三条。"
|
||||
summary = "基于 Qdrant 与 bge-m3 的 1536 维向量检索,详见《部署手册》第三条。"
|
||||
assert entity_recall(source, summary) == pytest.approx(1.0)
|
||||
|
||||
def test_partial_entities_kept(self) -> None:
|
||||
source = "支持 text-embedding-3-small 和 bge-m3 两个模型。"
|
||||
summary = "支持 bge-m3 模型。"
|
||||
score = entity_recall(source, summary)
|
||||
assert 0.0 < score < 1.0
|
||||
|
||||
def test_no_entities_returns_full_score(self) -> None:
|
||||
assert entity_recall("这是一段没有任何实体的中文普通句子。", "很短") == 1.0
|
||||
|
||||
def test_empty_summary_drops_all(self) -> None:
|
||||
source = "版本 v1.2.0 修复了 3 个问题。"
|
||||
assert entity_recall(source, "") == pytest.approx(0.0)
|
||||
|
||||
|
||||
class TestRoutingF1:
|
||||
def test_perfect_routing(self) -> None:
|
||||
golden = [{"d1"}, {"d2"}]
|
||||
routed = [{"d1"}, {"d2"}]
|
||||
result = routing_f1(golden, routed)
|
||||
assert result == {"precision": 1.0, "recall": 1.0, "f1": 1.0}
|
||||
|
||||
def test_micro_average(self) -> None:
|
||||
# query1: tp=1 fp=1;query2: tp=0 fn=1 → micro p=1/2 r=1/2 f1=1/2
|
||||
golden = [{"d1"}, {"d2"}]
|
||||
routed = [{"d1", "d3"}, set()]
|
||||
result = routing_f1(golden, routed)
|
||||
assert result["precision"] == pytest.approx(0.5)
|
||||
assert result["recall"] == pytest.approx(0.5)
|
||||
assert result["f1"] == pytest.approx(0.5)
|
||||
|
||||
def test_empty_routed(self) -> None:
|
||||
result = routing_f1([{"d1"}], [set()])
|
||||
assert result["precision"] == 0.0
|
||||
assert result["recall"] == 0.0
|
||||
|
||||
|
||||
class TestPruningLoss:
|
||||
def test_mixed_cases(self) -> None:
|
||||
assert pruning_loss([False, True, False, False]) == pytest.approx(0.25)
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert pruning_loss([]) == 0.0
|
||||
|
||||
|
||||
class TestPrecisionRecallAtK:
|
||||
def test_precision_at_k(self) -> None:
|
||||
hits = ["d1", "d2", "d3", "d4", "d5"]
|
||||
assert precision_at_k(hits, {"d1", "d3"}, 5) == pytest.approx(0.4)
|
||||
|
||||
def test_precision_at_k_empty_hits(self) -> None:
|
||||
assert precision_at_k([], {"d1"}, 5) == 0.0
|
||||
|
||||
def test_recall_at_k_dedup(self) -> None:
|
||||
hits = ["d1", "d1", "d2"]
|
||||
assert recall_at_k(hits, {"d1", "d2", "d3"}, 3) == pytest.approx(2 / 3)
|
||||
|
||||
def test_recall_at_k_empty_golden(self) -> None:
|
||||
assert recall_at_k(["d1"], set(), 5) == 0.0
|
||||
|
||||
|
||||
class TestAggregate:
|
||||
def test_mean_per_key(self) -> None:
|
||||
records = [{"a": 1.0, "b": 0.5}, {"a": 0.0, "b": 1.0}]
|
||||
result = aggregate(records)
|
||||
assert result["a"] == pytest.approx(0.5)
|
||||
assert result["b"] == pytest.approx(0.75)
|
||||
|
||||
def test_partial_keys(self) -> None:
|
||||
result = aggregate([{"a": 1.0}, {"a": 0.0, "b": 1.0}])
|
||||
assert result["a"] == pytest.approx(0.5)
|
||||
assert result["b"] == pytest.approx(1.0)
|
||||
|
||||
def test_empty(self) -> None:
|
||||
assert aggregate([]) == {}
|
||||
@@ -0,0 +1,104 @@
|
||||
"""标题树解析器的单元测试"""
|
||||
|
||||
from app.core.headings import parse_headings, render_outline
|
||||
|
||||
|
||||
class TestMarkdownHeadings:
|
||||
"""Markdown ATX 标题解析"""
|
||||
|
||||
def test_multi_level(self):
|
||||
"""多级 Markdown 标题,层级为 # 数量"""
|
||||
text = "# 一级标题\n正文内容\n## 二级标题\n更多内容\n### 三级标题\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level, h.line_index) for h in headings] == [
|
||||
("一级标题", 1, 0),
|
||||
("二级标题", 2, 2),
|
||||
("三级标题", 3, 4),
|
||||
]
|
||||
|
||||
def test_hash_without_space_not_heading(self):
|
||||
"""# 后无空白不是标题"""
|
||||
assert parse_headings("#不是标题") == []
|
||||
|
||||
def test_up_to_six_levels(self):
|
||||
"""最多支持 6 级标题"""
|
||||
text = "###### 六级标题\n####### 七级不算\n"
|
||||
headings = parse_headings(text)
|
||||
assert len(headings) == 1
|
||||
assert headings[0].level == 6
|
||||
|
||||
|
||||
class TestChineseChapterHeadings:
|
||||
"""中文章/节/篇编号标题解析"""
|
||||
|
||||
def test_chapter_and_section(self):
|
||||
"""章=1 级,节=2 级"""
|
||||
text = "第一章 总则\n正文\n第一节 一般规定\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level) for h in headings] == [("第一章 总则", 1), ("第一节 一般规定", 2)]
|
||||
|
||||
def test_pian_is_level_one(self):
|
||||
"""篇=1 级,支持阿拉伯数字编号"""
|
||||
text = "第一篇 概述\n第2章 背景\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level) for h in headings] == [("第一篇 概述", 1), ("第2章 背景", 1)]
|
||||
|
||||
def test_chinese_enum_is_level_one(self):
|
||||
"""一、二、…… 固定 1 级"""
|
||||
text = "一、项目背景\n二、建设目标\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level) for h in headings] == [("一、项目背景", 1), ("二、建设目标", 1)]
|
||||
|
||||
def test_overlong_line_not_heading(self):
|
||||
"""编号行超过 60 字符视为正文"""
|
||||
text = "第一章 " + "很长的标题" * 12 + "\n"
|
||||
assert parse_headings(text) == []
|
||||
|
||||
|
||||
class TestNumericHeadings:
|
||||
"""数字编号标题解析"""
|
||||
|
||||
def test_dot_segments_determine_level(self):
|
||||
"""按点分段数定层级:1.=1、1.1=2、1.1.1=3"""
|
||||
text = "1. 概述\n1.1 背景\n1.1.1 细节\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level) for h in headings] == [
|
||||
("1. 概述", 1),
|
||||
("1.1 背景", 2),
|
||||
("1.1.1 细节", 3),
|
||||
]
|
||||
|
||||
def test_dunhao_separator(self):
|
||||
"""顿号分隔的数字编号为 1 级"""
|
||||
headings = parse_headings("1、基本要求\n2、总体架构\n")
|
||||
assert [(h.title, h.level) for h in headings] == [("1、基本要求", 1), ("2、总体架构", 1)]
|
||||
|
||||
|
||||
class TestMixedAndPlain:
|
||||
"""混合模式与无结构文本"""
|
||||
|
||||
def test_no_headings_returns_empty(self):
|
||||
"""无标题文本返回空列表"""
|
||||
text = "第一段正文内容。\n\n第二段正文内容。\n"
|
||||
assert parse_headings(text) == []
|
||||
|
||||
def test_empty_text(self):
|
||||
assert parse_headings("") == []
|
||||
|
||||
def test_mixed_patterns(self):
|
||||
"""Markdown 与中文编号混合时均能识别"""
|
||||
text = "# Markdown 标题\n正文\n第一章 中文标题\n1.1 数字标题\n"
|
||||
headings = parse_headings(text)
|
||||
assert [(h.title, h.level) for h in headings] == [
|
||||
("Markdown 标题", 1),
|
||||
("第一章 中文标题", 1),
|
||||
("1.1 数字标题", 2),
|
||||
]
|
||||
|
||||
|
||||
class TestRenderOutline:
|
||||
"""标题树渲染为大纲文本"""
|
||||
|
||||
def test_indent_by_level(self):
|
||||
headings = parse_headings("# 安装指南\n## 环境准备\n## 安装步骤\n### 验证\n")
|
||||
assert render_outline(headings) == "- 安装指南\n - 环境准备\n - 安装步骤\n - 验证"
|
||||
@@ -0,0 +1,262 @@
|
||||
"""异步入库集成验证
|
||||
|
||||
在真实内存 Qdrant + FakeOllama 环境下验证异步入库全链路闭环:
|
||||
POST 202 拿 task_id → wait_done 终态 → 任务查询 → 文档详情 → 检索命中 → 删除;
|
||||
并覆盖两条边界路径:
|
||||
- 失败路径:分类阶段抛错 → failed + error.stage/partial_summary,Redis 失败镜像 TTL=7 天
|
||||
- Redis 降级路径:无 Redis / Redis 全异常时任务照常完成、查询正常
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.api.v1 import search as search_module
|
||||
from app.config import Settings
|
||||
from app.core.chunker import Chunker
|
||||
from app.core.classifier import Classifier
|
||||
from app.core.ingest_tasks import IngestTaskManager, IngestTaskStatus
|
||||
from app.core.ingestion import Ingester
|
||||
from app.core.query_parser import QueryParser
|
||||
from app.core.retriever import Retriever
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.core.summarizer import Summarizer
|
||||
from app.main import app
|
||||
from app.models.document import DocumentInput, IngestionResult
|
||||
from app.models.knowledge import load_taxonomy
|
||||
from app.services.qdrant import QdrantService
|
||||
from tests.test_e2e_integration import (
|
||||
FAKE_CATEGORY,
|
||||
FAKE_L1_SUMMARY,
|
||||
DeterministicEmbedding,
|
||||
FakeCache,
|
||||
FakeOllama,
|
||||
)
|
||||
|
||||
|
||||
class ClassifyFailOllama(FakeOllama):
|
||||
"""分类阶段抛错的 FakeOllama(总结等其余行为与正常版一致)"""
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
if "你是知识库分类助手" in prompt:
|
||||
raise RuntimeError("模拟分类模型不可用")
|
||||
return await super().generate(prompt, json_mode=json_mode)
|
||||
|
||||
|
||||
class RecordingRedis:
|
||||
"""记录 set_json 调用的假 Redis(get_json 恒未命中)"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.set_calls: list[tuple[str, dict[str, Any], int | None]] = []
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
return None
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
self.set_calls.append((key, value, ttl))
|
||||
return True
|
||||
|
||||
|
||||
class BrokenRedis:
|
||||
"""全部方法抛异常的假 Redis,验证任务管理器的容错降级"""
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
raise ConnectionError("模拟 Redis 不可用")
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
raise ConnectionError("模拟 Redis 不可用")
|
||||
|
||||
|
||||
class RecordingIngester:
|
||||
"""包装真实 Ingester,记录 progress_cb 回调的阶段序列(接口与 Ingester 一致)"""
|
||||
|
||||
def __init__(self, ingester: Ingester) -> None:
|
||||
self._ingester = ingester
|
||||
self.stages: list[str] = []
|
||||
|
||||
async def ingest(
|
||||
self, doc: DocumentInput, progress_cb: Callable[[str], None] | None = None
|
||||
) -> IngestionResult:
|
||||
def _cb(stage: str) -> None:
|
||||
self.stages.append(stage)
|
||||
if progress_cb is not None:
|
||||
progress_cb(stage)
|
||||
|
||||
return await self._ingester.ingest(doc, progress_cb=_cb)
|
||||
|
||||
|
||||
async def _make_env(ollama: FakeOllama) -> tuple[QdrantService, Ingester, Retriever]:
|
||||
"""构建集成环境:真实组件 + 内存 Qdrant + 指定 FakeOllama + 确定性向量"""
|
||||
qdrant = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await qdrant.ensure_collections()
|
||||
taxonomy = load_taxonomy()
|
||||
embedding = DeterministicEmbedding()
|
||||
ingester = Ingester(
|
||||
summarizer=Summarizer(ollama=ollama), # type: ignore[arg-type]
|
||||
classifier=Classifier(ollama=ollama, taxonomy=taxonomy), # type: ignore[arg-type]
|
||||
chunker=Chunker(),
|
||||
embedding=embedding, # type: ignore[arg-type]
|
||||
sparse=SparseEncoder(),
|
||||
qdrant=qdrant,
|
||||
)
|
||||
retriever = Retriever(
|
||||
qdrant=qdrant,
|
||||
query_parser=QueryParser(ollama=ollama, taxonomy=taxonomy, cache=FakeCache()), # type: ignore[arg-type]
|
||||
embedding=embedding, # type: ignore[arg-type]
|
||||
sparse_encoder=SparseEncoder(),
|
||||
)
|
||||
return qdrant, ingester, retriever
|
||||
|
||||
|
||||
def _structured_doc() -> DocumentInput:
|
||||
"""带 Markdown 标题结构的中文长文档(>500 字符,走完整三级总结)"""
|
||||
paragraph = "这是章节正文内容,包含足够多的信息量,用于测试切分与向量化流程。" * 20
|
||||
text = f"# 安装指南\n{paragraph}\n\n## 环境准备\n{paragraph}\n\n## 安装步骤\n{paragraph}"
|
||||
return DocumentInput(text=text, title="安装文档")
|
||||
|
||||
|
||||
def _patch_app(
|
||||
monkeypatch: pytest.MonkeyPatch, manager: IngestTaskManager, qdrant: QdrantService, retriever: Retriever
|
||||
) -> None:
|
||||
"""替换模块级单例:任务管理器 / Qdrant / Retriever,lifespan 建集合改空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
monkeypatch.setattr(document_module, "_task_manager", manager)
|
||||
monkeypatch.setattr(document_module, "_qdrant", qdrant)
|
||||
monkeypatch.setattr(search_module, "_retriever", retriever)
|
||||
monkeypatch.setattr(search_module, "get_cache", lambda: FakeCache())
|
||||
|
||||
|
||||
class TestIngestAsyncFullLoop:
|
||||
"""全链路闭环:异步入库 → 任务查询 → 文档管理 → 检索 → 删除"""
|
||||
|
||||
async def test_full_loop(self, monkeypatch: pytest.MonkeyPatch):
|
||||
qdrant, ingester, retriever = await _make_env(FakeOllama())
|
||||
recording = RecordingIngester(ingester)
|
||||
manager = IngestTaskManager(recording, None, Settings()) # type: ignore[arg-type]
|
||||
_patch_app(monkeypatch, manager, qdrant, retriever)
|
||||
doc = _structured_doc()
|
||||
|
||||
with TestClient(app) as client:
|
||||
# 1. 提交入库:202 + task_id
|
||||
resp_post = client.post("/api/v1/documents", json={"text": doc.text, "title": doc.title})
|
||||
assert resp_post.status_code == 202
|
||||
body_post = resp_post.json()
|
||||
assert body_post["code"] == 0
|
||||
assert body_post["data"]["status"] == "pending"
|
||||
task_id = body_post["data"]["task_id"]
|
||||
assert task_id
|
||||
|
||||
# 2. 等待终态:done + 结果完整 + 阶段序列完整
|
||||
final = await manager.wait_done(task_id)
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
result = final["result"]
|
||||
assert result["document_id"]
|
||||
assert result["category"] == FAKE_CATEGORY
|
||||
assert result["chunks_count"] >= 1
|
||||
assert recording.stages == ["summarizing", "classifying", "embedding", "writing"]
|
||||
document_id = result["document_id"]
|
||||
|
||||
# 3. 任务状态查询:code=0、done、含 result
|
||||
resp_task = client.get(f"/api/v1/documents/tasks/{task_id}")
|
||||
body_task = resp_task.json()
|
||||
assert body_task["code"] == 0
|
||||
assert body_task["data"]["status"] == "done"
|
||||
assert body_task["data"]["result"]["document_id"] == document_id
|
||||
|
||||
# 4. 文档详情:入库完成后即可管理
|
||||
resp_detail = client.get(f"/api/v1/documents/{document_id}")
|
||||
body_detail = resp_detail.json()
|
||||
assert body_detail["code"] == 0
|
||||
|
||||
# 5. 检索:相关 query 命中该文档
|
||||
resp_search = client.post("/api/v1/search", json={"query": "安装步骤有哪些注意事项?"})
|
||||
body_search = resp_search.json()
|
||||
assert body_search["code"] == 0
|
||||
hits = body_search["data"]["hits"]
|
||||
assert hits
|
||||
assert any(hit["doc_id"] == document_id for hit in hits)
|
||||
|
||||
# 6. 删除:四层集合中该文档全部清除
|
||||
resp_delete = client.delete(f"/api/v1/documents/{document_id}")
|
||||
body_delete = resp_delete.json()
|
||||
assert body_delete["code"] == 0
|
||||
assert body_delete["data"]["deleted_total"] > 0
|
||||
|
||||
|
||||
class TestIngestAsyncFailure:
|
||||
"""失败路径:分类阶段抛错 → failed 状态、错误透传与 Redis 失败镜像 TTL"""
|
||||
|
||||
async def test_classify_failure(self, monkeypatch: pytest.MonkeyPatch):
|
||||
qdrant, ingester, retriever = await _make_env(ClassifyFailOllama())
|
||||
redis = RecordingRedis()
|
||||
manager = IngestTaskManager(ingester, redis, Settings()) # type: ignore[arg-type]
|
||||
_patch_app(monkeypatch, manager, qdrant, retriever)
|
||||
|
||||
# 直接经 manager 提交(与 API 提交同路径),在测试事件循环内等待终态与镜像写完
|
||||
task_id = await manager.submit(_structured_doc())
|
||||
final = await manager.wait_done(task_id)
|
||||
|
||||
# 终态 failed:阶段与已产出总结透传
|
||||
assert final["status"] == IngestTaskStatus.FAILED
|
||||
error = final["error"]
|
||||
assert error["stage"] == "classify"
|
||||
assert error["partial_summary"]
|
||||
assert error["partial_summary"]["l1_summary"] == FAKE_L1_SUMMARY
|
||||
|
||||
# Redis 镜像:failed 快照使用失败 TTL(7 天),其余快照用 done TTL(24h)
|
||||
failed_mirrors = [c for c in redis.set_calls if c[1]["status"] == IngestTaskStatus.FAILED]
|
||||
assert failed_mirrors
|
||||
assert all(ttl == 604800 for _, _, ttl in failed_mirrors)
|
||||
non_failed = [c for c in redis.set_calls if c[1]["status"] != IngestTaskStatus.FAILED]
|
||||
assert non_failed
|
||||
assert all(ttl == 86400 for _, _, ttl in non_failed)
|
||||
|
||||
# API 查询:GET tasks/{task_id} 返回同样的失败信息
|
||||
with TestClient(app) as client:
|
||||
resp_task = client.get(f"/api/v1/documents/tasks/{task_id}")
|
||||
body_task = resp_task.json()
|
||||
assert body_task["code"] == 0
|
||||
assert body_task["data"]["status"] == "failed"
|
||||
assert body_task["data"]["error"]["stage"] == "classify"
|
||||
assert body_task["data"]["error"]["partial_summary"]
|
||||
assert body_task["data"]["error"]["partial_summary"]["l1_summary"] == FAKE_L1_SUMMARY
|
||||
|
||||
|
||||
class TestIngestAsyncRedisDegraded:
|
||||
"""Redis 降级路径:无 Redis 或 Redis 全异常时任务照常完成"""
|
||||
|
||||
async def test_redis_none(self):
|
||||
"""redis=None:纯内存模式,提交→done→查询全流程正常"""
|
||||
_, ingester, _ = await _make_env(FakeOllama())
|
||||
manager = IngestTaskManager(ingester, None, Settings())
|
||||
|
||||
task_id = await manager.submit(_structured_doc())
|
||||
final = await manager.wait_done(task_id)
|
||||
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
assert final["result"]["document_id"]
|
||||
record = await manager.get(task_id)
|
||||
assert record is not None
|
||||
assert record["status"] == IngestTaskStatus.DONE
|
||||
|
||||
async def test_redis_broken(self):
|
||||
"""Redis 读写全抛异常:镜像/读取容错降级,任务流程与查询不受影响"""
|
||||
_, ingester, _ = await _make_env(FakeOllama())
|
||||
manager = IngestTaskManager(ingester, BrokenRedis(), Settings()) # type: ignore[arg-type]
|
||||
|
||||
task_id = await manager.submit(_structured_doc())
|
||||
final = await manager.wait_done(task_id)
|
||||
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
assert final["result"]["document_id"]
|
||||
record = await manager.get(task_id)
|
||||
assert record is not None
|
||||
assert record["status"] == IngestTaskStatus.DONE
|
||||
@@ -0,0 +1,160 @@
|
||||
"""入库异步任务 API 测试(TestClient + FakeManager,不真实联网)"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.main import app
|
||||
from app.models.document import DocumentInput
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
class FakeManager:
|
||||
"""假入库任务管理器:记录 submit 调用,按 task_id 返回预置任务"""
|
||||
|
||||
def __init__(self, tasks: dict[str, dict[str, Any]] | None = None, task_id: str = "task-1") -> None:
|
||||
self.tasks = tasks or {}
|
||||
self.task_id = task_id
|
||||
self.submitted: list[DocumentInput] = []
|
||||
|
||||
async def submit(self, doc: DocumentInput) -> str:
|
||||
self.submitted.append(doc)
|
||||
return self.task_id
|
||||
|
||||
async def get(self, task_id: str) -> dict[str, Any] | None:
|
||||
return self.tasks.get(task_id)
|
||||
|
||||
|
||||
def _task(task_id: str, status: str, **extra: Any) -> dict[str, Any]:
|
||||
"""构造一条任务记录(时间字段为固定 ISO8601 字符串)"""
|
||||
record: dict[str, Any] = {
|
||||
"task_id": task_id,
|
||||
"status": status,
|
||||
"created_at": "2026-07-29T08:00:00+00:00",
|
||||
"updated_at": "2026-07-29T08:00:01+00:00",
|
||||
"result": None,
|
||||
"error": None,
|
||||
}
|
||||
record.update(extra)
|
||||
return record
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
def _inject_manager(monkeypatch: pytest.MonkeyPatch, manager: FakeManager) -> None:
|
||||
"""将 FakeManager 注入路由的 _get_task_manager"""
|
||||
monkeypatch.setattr(document_module, "_get_task_manager", lambda: manager)
|
||||
|
||||
|
||||
def test_post_documents_returns_202(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""POST /documents:HTTP 202,data.task_id 非空且 status=pending,文档已提交给管理器"""
|
||||
manager = FakeManager(task_id="abc123")
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post("/api/v1/documents", json={"text": "正文内容", "title": "标题"})
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"]["task_id"] == "abc123"
|
||||
assert body["data"]["status"] == "pending"
|
||||
assert len(manager.submitted) == 1
|
||||
|
||||
|
||||
def test_get_task_pending(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""GET tasks/{id}:进行中任务 → code=0,含 status/created_at/updated_at"""
|
||||
manager = FakeManager(tasks={"t-pending": _task("t-pending", "summarizing")})
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.get("/api/v1/documents/tasks/t-pending")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["task_id"] == "t-pending"
|
||||
assert data["status"] == "summarizing"
|
||||
assert data["created_at"]
|
||||
assert data["updated_at"]
|
||||
assert data["result"] is None
|
||||
assert data["error"] is None
|
||||
|
||||
|
||||
def test_get_task_done(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""GET tasks/{id}:done 任务 → data.result 含 document_id"""
|
||||
result = {
|
||||
"document_id": "doc-1",
|
||||
"summary": {"l1_summary": "一句话", "l2_outline": None, "l3_content_outline": "大纲", "level": "L3"},
|
||||
"category": "技术文档",
|
||||
"collection": "四层集合",
|
||||
"chunks_count": 3,
|
||||
"tags": ["API"],
|
||||
"category_confidence": 0.9,
|
||||
}
|
||||
manager = FakeManager(tasks={"t-done": _task("t-done", "done", result=result)})
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.get("/api/v1/documents/tasks/t-done")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["status"] == "done"
|
||||
assert data["result"]["document_id"] == "doc-1"
|
||||
assert data["result"]["chunks_count"] == 3
|
||||
assert data["error"] is None
|
||||
|
||||
|
||||
def test_get_task_failed(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""GET tasks/{id}:failed 任务 → data.error 含 stage/message/partial_summary"""
|
||||
partial = {"l1_summary": "L1", "l2_outline": None, "l3_content_outline": "L3", "level": "L3"}
|
||||
error = {"stage": "embed", "message": "向量化失败: boom", "partial_summary": partial}
|
||||
manager = FakeManager(tasks={"t-failed": _task("t-failed", "failed", error=error)})
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.get("/api/v1/documents/tasks/t-failed")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["status"] == "failed"
|
||||
assert data["error"]["stage"] == "embed"
|
||||
assert "boom" in data["error"]["message"]
|
||||
assert data["error"]["partial_summary"] == partial
|
||||
assert data["result"] is None
|
||||
|
||||
|
||||
def test_get_task_not_found(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""GET tasks/{id}:不存在的 task_id → code=1004"""
|
||||
_inject_manager(monkeypatch, FakeManager())
|
||||
|
||||
resp = client.get("/api/v1/documents/tasks/不存在")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1004
|
||||
assert body["message"] == "任务不存在"
|
||||
assert body["data"] is None
|
||||
|
||||
|
||||
def test_post_empty_text_creates_no_task(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""POST 空文本:code=1001,且不产生任务(submit 未被调用)"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post("/api/v1/documents", json={"text": " "})
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert manager.submitted == []
|
||||
@@ -0,0 +1,200 @@
|
||||
"""IngestTaskManager 单元测试(FakeIngester/FakeRedis,不真实联网)"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from app.config import Settings
|
||||
from app.core.ingest_tasks import REDIS_KEY_PREFIX, IngestTaskManager, IngestTaskStatus
|
||||
from app.core.ingestion import IngestionError
|
||||
from app.models.document import DocumentInput, DocumentSummary, IngestionResult, SummaryLevel
|
||||
|
||||
|
||||
def make_summary() -> DocumentSummary:
|
||||
"""构造固定的三级总结"""
|
||||
return DocumentSummary(
|
||||
l1_summary="一句话总结", l2_outline=None, l3_content_outline="内容大纲", level=SummaryLevel.L3
|
||||
)
|
||||
|
||||
|
||||
def make_result(doc_id: str = "doc-1") -> IngestionResult:
|
||||
"""构造固定的入库结果"""
|
||||
return IngestionResult(
|
||||
document_id=doc_id,
|
||||
summary=make_summary(),
|
||||
category="tech",
|
||||
collection="四层集合",
|
||||
chunks_count=2,
|
||||
tags=["t"],
|
||||
category_confidence=0.9,
|
||||
)
|
||||
|
||||
|
||||
class FakeIngester:
|
||||
"""假入库器:上报固定阶段序列;可配置抛错或用闸门阻塞以验证并发限流"""
|
||||
|
||||
def __init__(self, error: Exception | None = None) -> None:
|
||||
self.error = error
|
||||
self.stages: list[str] = []
|
||||
self.calls: list[DocumentInput] = []
|
||||
self.gate: asyncio.Event | None = None
|
||||
self.started = asyncio.Event()
|
||||
|
||||
async def ingest(self, doc: DocumentInput, progress_cb: Callable[[str], None] | None = None) -> IngestionResult:
|
||||
self.calls.append(doc)
|
||||
self.started.set()
|
||||
if progress_cb is not None:
|
||||
for stage in ("summarizing", "classifying", "embedding", "writing"):
|
||||
progress_cb(stage)
|
||||
self.stages.append(stage)
|
||||
if self.gate is not None:
|
||||
await self.gate.wait()
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return make_result()
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
"""内存版 Redis:记录全部写入(含 TTL),可配置写入抛错模拟故障降级"""
|
||||
|
||||
def __init__(self, fail_writes: bool = False) -> None:
|
||||
self.fail_writes = fail_writes
|
||||
self.writes: list[tuple[str, dict[str, Any], int | None]] = []
|
||||
self.store: dict[str, dict[str, Any]] = {}
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
self.writes.append((key, value, ttl))
|
||||
if self.fail_writes:
|
||||
raise RuntimeError("redis down")
|
||||
self.store[key] = value
|
||||
return True
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
return self.store.get(key)
|
||||
|
||||
|
||||
async def test_submit_returns_immediately_and_completes() -> None:
|
||||
"""submit 立即返回;任务后台跑完为 done,结果完整,阶段序列齐全,Redis 镜像同步"""
|
||||
ingester = FakeIngester()
|
||||
redis = FakeRedis()
|
||||
manager = IngestTaskManager(ingester, redis, Settings())
|
||||
|
||||
task_id = await manager.submit(DocumentInput(text="正文", title="标题"))
|
||||
assert isinstance(task_id, str) and len(task_id) == 32
|
||||
|
||||
# submit 后立即可查:状态在合法集合内(pending 或已进入某阶段)
|
||||
record = await manager.get(task_id)
|
||||
assert record is not None
|
||||
assert record["status"] in set(IngestTaskStatus)
|
||||
|
||||
final = await manager.wait_done(task_id, timeout=5)
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
assert final["result"] == make_result().model_dump(mode="json")
|
||||
assert final["error"] is None
|
||||
assert ingester.stages == ["summarizing", "classifying", "embedding", "writing"]
|
||||
# 时间字段为 ISO8601 字符串
|
||||
datetime.fromisoformat(final["created_at"])
|
||||
datetime.fromisoformat(final["updated_at"])
|
||||
# Redis 镜像已写入终态
|
||||
mirrored = await redis.get_json(f"{REDIS_KEY_PREFIX}{task_id}")
|
||||
assert mirrored is not None
|
||||
assert mirrored["status"] == IngestTaskStatus.DONE
|
||||
|
||||
|
||||
async def test_ingestion_error_marks_failed_with_stage_and_partial_summary() -> None:
|
||||
"""IngestionError:任务 failed,error.stage 透传,partial_summary 保留"""
|
||||
summary = DocumentSummary(l1_summary="L1", l2_outline=None, l3_content_outline="L3", level=SummaryLevel.L3)
|
||||
ingester = FakeIngester(error=IngestionError("classify", "分类判定失败: boom", summary=summary))
|
||||
manager = IngestTaskManager(ingester, FakeRedis(), Settings())
|
||||
|
||||
task_id = await manager.submit(DocumentInput(text="正文"))
|
||||
final = await manager.wait_done(task_id, timeout=5)
|
||||
assert final["status"] == IngestTaskStatus.FAILED
|
||||
assert final["result"] is None
|
||||
assert final["error"]["stage"] == "classify"
|
||||
assert "boom" in final["error"]["message"]
|
||||
assert final["error"]["partial_summary"] == summary.model_dump(mode="json")
|
||||
|
||||
|
||||
async def test_unexpected_error_marks_failed_with_unknown_stage() -> None:
|
||||
"""普通异常:任务 failed,error.stage 为 unknown"""
|
||||
ingester = FakeIngester(error=RuntimeError("炸了"))
|
||||
manager = IngestTaskManager(ingester, None, Settings())
|
||||
|
||||
task_id = await manager.submit(DocumentInput(text="正文"))
|
||||
final = await manager.wait_done(task_id, timeout=5)
|
||||
assert final["status"] == IngestTaskStatus.FAILED
|
||||
assert final["error"]["stage"] == "unknown"
|
||||
assert "炸了" in final["error"]["message"]
|
||||
|
||||
|
||||
async def test_concurrency_limited_by_semaphore() -> None:
|
||||
"""并发上限 1 时两个任务串行:第一个放行前第二个不得进入执行"""
|
||||
ingester = FakeIngester()
|
||||
ingester.gate = asyncio.Event() # 首次调用阻塞,验证第二个任务在排队
|
||||
manager = IngestTaskManager(ingester, None, Settings(ingest_max_concurrency=1))
|
||||
|
||||
task1 = await manager.submit(DocumentInput(text="a", title="t1"))
|
||||
task2 = await manager.submit(DocumentInput(text="b", title="t2"))
|
||||
await asyncio.wait_for(ingester.started.wait(), timeout=1)
|
||||
await asyncio.sleep(0.05) # 给第二个任务调度机会
|
||||
assert len(ingester.calls) == 1
|
||||
queued = await manager.get(task2)
|
||||
assert queued is not None
|
||||
assert queued["status"] == IngestTaskStatus.PENDING
|
||||
|
||||
ingester.gate.set()
|
||||
final1 = await manager.wait_done(task1, timeout=5)
|
||||
final2 = await manager.wait_done(task2, timeout=5)
|
||||
assert final1["status"] == IngestTaskStatus.DONE
|
||||
assert final2["status"] == IngestTaskStatus.DONE
|
||||
assert [doc.title for doc in ingester.calls] == ["t1", "t2"]
|
||||
|
||||
|
||||
async def test_redis_write_failure_does_not_break_task() -> None:
|
||||
"""Redis 写失败仅告警:任务仍正常跑完为 done"""
|
||||
redis = FakeRedis(fail_writes=True)
|
||||
manager = IngestTaskManager(FakeIngester(), redis, Settings())
|
||||
|
||||
task_id = await manager.submit(DocumentInput(text="正文"))
|
||||
final = await manager.wait_done(task_id, timeout=5)
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
assert final["result"] is not None
|
||||
assert redis.writes # 确实尝试过写 Redis
|
||||
|
||||
|
||||
async def test_redis_mirror_ttl_done_and_failed() -> None:
|
||||
"""Redis 镜像 TTL:进行中与 done 用 ttl_done,failed 用 ttl_failed"""
|
||||
settings = Settings(ingest_task_ttl_done=86400, ingest_task_ttl_failed=604800)
|
||||
|
||||
redis = FakeRedis()
|
||||
manager = IngestTaskManager(FakeIngester(), redis, settings)
|
||||
done_id = await manager.submit(DocumentInput(text="x"))
|
||||
await manager.wait_done(done_id, timeout=5)
|
||||
done_writes = [(v, ttl) for key, v, ttl in redis.writes if key.endswith(done_id)]
|
||||
assert done_writes
|
||||
assert all(ttl == 86400 for _, ttl in done_writes)
|
||||
assert done_writes[-1][0]["status"] == IngestTaskStatus.DONE
|
||||
|
||||
redis2 = FakeRedis()
|
||||
failing_manager = IngestTaskManager(FakeIngester(error=RuntimeError("boom")), redis2, settings)
|
||||
failed_id = await failing_manager.submit(DocumentInput(text="y"))
|
||||
await failing_manager.wait_done(failed_id, timeout=5)
|
||||
failed_writes = [(v, ttl) for key, v, ttl in redis2.writes if key.endswith(failed_id)]
|
||||
assert failed_writes
|
||||
assert failed_writes[-1][0]["status"] == IngestTaskStatus.FAILED
|
||||
assert failed_writes[-1][1] == 604800
|
||||
|
||||
|
||||
async def test_get_falls_back_to_redis_then_none() -> None:
|
||||
"""get:内存 miss 时回查 Redis;两者都没有返回 None"""
|
||||
redis = FakeRedis()
|
||||
manager = IngestTaskManager(FakeIngester(), redis, Settings())
|
||||
|
||||
assert await manager.get("不存在") is None
|
||||
|
||||
redis.store[f"{REDIS_KEY_PREFIX}abc"] = {"task_id": "abc", "status": "done"}
|
||||
record = await manager.get("abc")
|
||||
assert record is not None
|
||||
assert record["status"] == "done"
|
||||
@@ -0,0 +1,273 @@
|
||||
"""Ingester 全链路单元测试(Summarizer/Classifier/Embedding/Qdrant 均为假实现,不真实联网)"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.chunker import Chunker
|
||||
from app.core.ingestion import Ingester, IngestionError
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.models.document import DocumentInput, DocumentSummary, SummaryLevel
|
||||
from app.models.knowledge import CategoryResult
|
||||
from app.services.qdrant import COLLECTION_L2, COLLECTION_L3
|
||||
|
||||
|
||||
class FakeSummarizer:
|
||||
"""返回固定总结结果的假 Summarizer"""
|
||||
|
||||
def __init__(self, summary: DocumentSummary) -> None:
|
||||
self.summary = summary
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
async def summarize(self, text: str, *, title: str = "") -> DocumentSummary:
|
||||
self.calls.append((text, title))
|
||||
return self.summary
|
||||
|
||||
|
||||
class FakeClassifier:
|
||||
"""返回固定分类结果的假 Classifier"""
|
||||
|
||||
def __init__(self, result: CategoryResult) -> None:
|
||||
self.result = result
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
async def classify(self, l1_summary: str, title: str = "") -> CategoryResult:
|
||||
self.calls.append((l1_summary, title))
|
||||
return self.result
|
||||
|
||||
|
||||
class FakeEmbedding:
|
||||
"""按输入数量返回伪向量的假 EmbeddingService,记录每次调用的文本"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[list[str]] = []
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
self.calls.append(list(texts))
|
||||
return [[float(i), 1.0] for i in range(len(texts))]
|
||||
|
||||
|
||||
class FakeQdrant:
|
||||
"""内存版 QdrantService,记录各层 upsert 调用;可配置在某一层抛错"""
|
||||
|
||||
def __init__(self, fail_on: str = "") -> None:
|
||||
self.fail_on = fail_on
|
||||
self.l1_calls: list[dict[str, Any]] = []
|
||||
self.nodes_calls: list[tuple[str, list[dict[str, Any]]]] = []
|
||||
self.chunks_calls: list[list[dict[str, Any]]] = []
|
||||
|
||||
async def upsert_l1(
|
||||
self,
|
||||
doc_id: str,
|
||||
title: str,
|
||||
summary: str,
|
||||
category: str,
|
||||
tags: list[str],
|
||||
dense_vector: list[float],
|
||||
sparse_vector: Any = None,
|
||||
) -> None:
|
||||
if self.fail_on == "l1":
|
||||
raise RuntimeError("qdrant down")
|
||||
self.l1_calls.append(
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"title": title,
|
||||
"summary": summary,
|
||||
"category": category,
|
||||
"tags": tags,
|
||||
"dense_vector": dense_vector,
|
||||
"sparse_vector": sparse_vector,
|
||||
}
|
||||
)
|
||||
|
||||
async def upsert_nodes(self, collection: str, nodes: list[dict[str, Any]]) -> None:
|
||||
if self.fail_on == collection:
|
||||
raise RuntimeError("qdrant down")
|
||||
self.nodes_calls.append((collection, nodes))
|
||||
|
||||
async def upsert_chunks(self, chunks: list[dict[str, Any]]) -> None:
|
||||
if self.fail_on == "chunks":
|
||||
raise RuntimeError("qdrant down")
|
||||
self.chunks_calls.append(chunks)
|
||||
|
||||
|
||||
def _make_ingester(
|
||||
summary: DocumentSummary,
|
||||
category: CategoryResult,
|
||||
qdrant: FakeQdrant,
|
||||
embedding: FakeEmbedding | None = None,
|
||||
) -> Ingester:
|
||||
return Ingester(
|
||||
summarizer=FakeSummarizer(summary), # type: ignore[arg-type]
|
||||
classifier=FakeClassifier(category), # type: ignore[arg-type]
|
||||
chunker=Chunker(),
|
||||
embedding=embedding or FakeEmbedding(), # type: ignore[arg-type]
|
||||
sparse=SparseEncoder(),
|
||||
qdrant=qdrant, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def _structured_doc() -> DocumentInput:
|
||||
"""带标题结构的长文档(切出多个 chunk,L2 走标题树)"""
|
||||
paragraph = "这是章节正文内容,包含足够多的信息量,用于测试切分与向量化流程。" * 10
|
||||
text = f"# 安装指南\n{paragraph}\n\n## 环境准备\n{paragraph}\n\n## 安装步骤\n{paragraph}"
|
||||
return DocumentInput(text=text, title="安装文档")
|
||||
|
||||
|
||||
def _structured_summary() -> DocumentSummary:
|
||||
return DocumentSummary(
|
||||
l1_summary="本文介绍软件的安装流程。",
|
||||
l2_outline="- 安装指南\n - 环境准备\n - 安装步骤",
|
||||
l3_content_outline=(
|
||||
"## 安装指南\n整体安装流程说明。\n## 环境准备\n准备依赖环境。\n## 安装步骤\n执行安装命令。"
|
||||
),
|
||||
level=SummaryLevel.L3,
|
||||
)
|
||||
|
||||
|
||||
def _category() -> CategoryResult:
|
||||
return CategoryResult(main_category="技术文档", tags=["安装", "运维"], confidence=0.9)
|
||||
|
||||
|
||||
class TestStructuredDocument:
|
||||
"""结构化长文档:四层集合 upsert 均被调用,结果字段透传"""
|
||||
|
||||
async def test_full_pipeline(self):
|
||||
doc = _structured_doc()
|
||||
summary = _structured_summary()
|
||||
qdrant = FakeQdrant()
|
||||
embedding = FakeEmbedding()
|
||||
ingester = _make_ingester(summary, _category(), qdrant, embedding)
|
||||
|
||||
result = await ingester.ingest(doc)
|
||||
|
||||
# 结果字段透传
|
||||
assert result.document_id
|
||||
assert result.summary is summary
|
||||
assert result.category == "技术文档"
|
||||
assert result.tags == ["安装", "运维"]
|
||||
assert result.category_confidence == 0.9
|
||||
|
||||
# chunk 数与 Chunker 直出一致
|
||||
expected_chunks = Chunker().chunk(doc.text, "expected")
|
||||
assert result.chunks_count == len(expected_chunks) > 1
|
||||
|
||||
# 批量 embedding 恰好一次:L1 + L2 节点 + L3 节点 + chunks
|
||||
assert len(embedding.calls) == 1
|
||||
l2_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L2]
|
||||
l3_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L3]
|
||||
assert len(embedding.calls[0]) == 1 + len(l2_calls[0]) + len(l3_calls[0]) + result.chunks_count
|
||||
|
||||
# L1:category/tags 透传,sparse 已启用
|
||||
assert len(qdrant.l1_calls) == 1
|
||||
l1 = qdrant.l1_calls[0]
|
||||
assert l1["doc_id"] == result.document_id
|
||||
assert l1["title"] == "安装文档"
|
||||
assert l1["summary"] == summary.l1_summary
|
||||
assert l1["category"] == "技术文档"
|
||||
assert l1["tags"] == ["安装", "运维"]
|
||||
assert l1["sparse_vector"] is not None
|
||||
|
||||
# L2:每个标题一个节点,text 与 section_path 均为祖先标题链
|
||||
assert len(l2_calls) == 1
|
||||
l2_nodes = l2_calls[0]
|
||||
assert len(l2_nodes) == 3
|
||||
assert [n["section_path"] for n in l2_nodes] == [
|
||||
"安装指南",
|
||||
"安装指南 / 环境准备",
|
||||
"安装指南 / 安装步骤",
|
||||
]
|
||||
assert all(n["text"] == n["section_path"] for n in l2_nodes)
|
||||
assert all(n["category"] == "技术文档" and n["tags"] == ["安装", "运维"] for n in l2_nodes)
|
||||
|
||||
# L3:按 "## " 分块,section_path 精确匹配到标题链
|
||||
assert len(l3_calls) == 1
|
||||
l3_nodes = l3_calls[0]
|
||||
assert len(l3_nodes) == 3
|
||||
assert [n["section_path"] for n in l3_nodes] == [
|
||||
"安装指南",
|
||||
"安装指南 / 环境准备",
|
||||
"安装指南 / 安装步骤",
|
||||
]
|
||||
assert l3_nodes[0]["text"].startswith("## 安装指南")
|
||||
|
||||
# chunks:携带 doc_summary 与 sparse 向量
|
||||
assert len(qdrant.chunks_calls) == 1
|
||||
chunk_dicts = qdrant.chunks_calls[0]
|
||||
assert len(chunk_dicts) == result.chunks_count
|
||||
assert all(c["doc_summary"] == summary.l1_summary for c in chunk_dicts)
|
||||
assert all(c["sparse_vector"] is not None for c in chunk_dicts)
|
||||
assert all(c["category"] == "技术文档" for c in chunk_dicts)
|
||||
assert [c["chunk_index"] for c in chunk_dicts] == list(range(result.chunks_count))
|
||||
|
||||
|
||||
class TestFallbackDocuments:
|
||||
"""2.5 级文档与无结构文档的 L2/L3 节点构建"""
|
||||
|
||||
async def test_l2_half_document_skips_l2_upsert(self):
|
||||
"""2.5 级文档(l2_outline=None)→ 不写 L2 集合,L3 整块一个节点"""
|
||||
summary = DocumentSummary(
|
||||
l1_summary="一条简短通知。",
|
||||
l2_outline=None,
|
||||
l3_content_outline="要点一:明天放假。\n要点二:注意安全。",
|
||||
level=SummaryLevel.L2_HALF,
|
||||
)
|
||||
doc = DocumentInput(text="简短通知正文,无标题结构。", title="通知")
|
||||
qdrant = FakeQdrant()
|
||||
ingester = _make_ingester(summary, _category(), qdrant)
|
||||
|
||||
result = await ingester.ingest(doc)
|
||||
|
||||
collections = [c for c, _ in qdrant.nodes_calls]
|
||||
assert COLLECTION_L2 not in collections
|
||||
# L3 内容大纲无 "## " → 整块一个节点,section_path 为空
|
||||
l3_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L3]
|
||||
assert len(l3_calls) == 1
|
||||
assert len(l3_calls[0]) == 1
|
||||
assert l3_calls[0][0]["section_path"] == ""
|
||||
assert l3_calls[0][0]["text"] == summary.l3_content_outline
|
||||
# 其余层级正常写入
|
||||
assert len(qdrant.l1_calls) == 1
|
||||
assert result.chunks_count == len(qdrant.chunks_calls[0])
|
||||
|
||||
async def test_l2_from_llm_outline_lines(self):
|
||||
"""无标题结构的 L3 级文档 → L2 节点来自 LLM 大纲行,section_path 为空"""
|
||||
summary = DocumentSummary(
|
||||
l1_summary="本文介绍两个主题。",
|
||||
l2_outline="1. 主题一\n2. 主题二",
|
||||
l3_content_outline="详细摘要内容,无分块标题。",
|
||||
level=SummaryLevel.L3,
|
||||
)
|
||||
text = "这是一段没有标题结构的正文内容," * 40
|
||||
doc = DocumentInput(text=text, title="")
|
||||
qdrant = FakeQdrant()
|
||||
ingester = _make_ingester(summary, _category(), qdrant)
|
||||
|
||||
await ingester.ingest(doc)
|
||||
|
||||
l2_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L2]
|
||||
assert len(l2_calls) == 1
|
||||
l2_nodes = l2_calls[0]
|
||||
assert [n["text"] for n in l2_nodes] == ["1. 主题一", "2. 主题二"]
|
||||
assert all(n["section_path"] == "" for n in l2_nodes)
|
||||
|
||||
|
||||
class TestQdrantFailure:
|
||||
"""Qdrant 写入失败:抛 IngestionError,stage=qdrant,总结不丢可重试"""
|
||||
|
||||
async def test_chunks_upsert_failure(self):
|
||||
doc = _structured_doc()
|
||||
summary = _structured_summary()
|
||||
qdrant = FakeQdrant(fail_on="chunks")
|
||||
ingester = _make_ingester(summary, _category(), qdrant)
|
||||
|
||||
with pytest.raises(IngestionError) as exc_info:
|
||||
await ingester.ingest(doc)
|
||||
|
||||
err = exc_info.value
|
||||
assert err.stage == "qdrant"
|
||||
assert err.summary is summary
|
||||
# L1/L2/L3 已写入,失败发生在 chunks 层
|
||||
assert len(qdrant.l1_calls) == 1
|
||||
assert {c for c, _ in qdrant.nodes_calls} == {COLLECTION_L2, COLLECTION_L3}
|
||||
assert qdrant.chunks_calls == []
|
||||
@@ -0,0 +1,99 @@
|
||||
"""LLM-as-judge 评测指标单元测试(FakeOllama 替身,不依赖真实模型)"""
|
||||
|
||||
import pytest
|
||||
|
||||
from scripts.eval.judge import hallucination_rate, taxonomy_consistency
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""记录 prompt 并按序返回预设响应(或抛出预设异常)的 OllamaClient 替身"""
|
||||
|
||||
def __init__(self, responses: list[str] | None = None, error: Exception | None = None) -> None:
|
||||
self.prompts: list[str] = []
|
||||
self.json_modes: list[bool] = []
|
||||
self._responses = list(responses or [])
|
||||
self._error = error
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
self.prompts.append(prompt)
|
||||
self.json_modes.append(json_mode)
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._responses.pop(0)
|
||||
|
||||
|
||||
def _assertions_json(*supported: bool) -> str:
|
||||
"""构造 hallucination 判定的 JSON 输出,每个 bool 对应一条断言的 supported"""
|
||||
claims = ", ".join(f'{{"claim": "断言{i}", "supported": {str(flag).lower()}}}' for i, flag in enumerate(supported))
|
||||
return f'{{"assertions": [{claims}]}}'
|
||||
|
||||
|
||||
class TestHallucinationRate:
|
||||
async def test_all_supported_returns_zero(self) -> None:
|
||||
ollama = FakeOllama([_assertions_json(True, True, True)])
|
||||
assert await hallucination_rate("摘要", "原文", ollama) == 0.0 # type: ignore[arg-type]
|
||||
|
||||
async def test_prompt_contains_source_and_summary(self) -> None:
|
||||
ollama = FakeOllama([_assertions_json(True)])
|
||||
await hallucination_rate("这是摘要内容", "这是原文内容", ollama) # type: ignore[arg-type]
|
||||
assert len(ollama.prompts) == 1
|
||||
assert "这是摘要内容" in ollama.prompts[0]
|
||||
assert "这是原文内容" in ollama.prompts[0]
|
||||
assert ollama.json_modes == [True]
|
||||
|
||||
async def test_two_of_five_unsupported(self) -> None:
|
||||
ollama = FakeOllama([_assertions_json(True, False, True, False, True)])
|
||||
rate = await hallucination_rate("摘要", "原文", ollama) # type: ignore[arg-type]
|
||||
assert rate == pytest.approx(0.4)
|
||||
|
||||
async def test_non_json_output_returns_zero(self) -> None:
|
||||
ollama = FakeOllama(["我无法完成这个判定任务。"])
|
||||
assert await hallucination_rate("摘要", "原文", ollama) == 0.0 # type: ignore[arg-type]
|
||||
|
||||
async def test_missing_assertions_field_returns_zero(self) -> None:
|
||||
ollama = FakeOllama(['{"result": "ok"}'])
|
||||
assert await hallucination_rate("摘要", "原文", ollama) == 0.0 # type: ignore[arg-type]
|
||||
|
||||
async def test_generate_error_returns_zero(self) -> None:
|
||||
ollama = FakeOllama(error=RuntimeError("Ollama 不可用"))
|
||||
assert await hallucination_rate("摘要", "原文", ollama) == 0.0 # type: ignore[arg-type]
|
||||
|
||||
async def test_noisy_output_json_extracted(self) -> None:
|
||||
# 模型输出前后带噪声时仍能提取首个 JSON 对象
|
||||
ollama = FakeOllama([f"先分析一下。\n{_assertions_json(True, False)}\n以上。"])
|
||||
rate = await hallucination_rate("摘要", "原文", ollama) # type: ignore[arg-type]
|
||||
assert rate == pytest.approx(0.5)
|
||||
|
||||
async def test_claims_truncated_to_max(self) -> None:
|
||||
# 超过 _MAX_CLAIMS(5) 的断言不计入:前 5 条全支持,第 6/7 条不支持 → 0.0
|
||||
ollama = FakeOllama([_assertions_json(True, True, True, True, True, False, False)])
|
||||
assert await hallucination_rate("摘要", "原文", ollama) == 0.0 # type: ignore[arg-type]
|
||||
|
||||
|
||||
class TestTaxonomyConsistency:
|
||||
async def test_consistent_true(self) -> None:
|
||||
ollama = FakeOllama(['{"consistent": true}'])
|
||||
assert await taxonomy_consistency("摘要", "制度", ollama) is True # type: ignore[arg-type]
|
||||
|
||||
async def test_consistent_false(self) -> None:
|
||||
ollama = FakeOllama(['{"consistent": false}'])
|
||||
assert await taxonomy_consistency("摘要", "制度", ollama) is False # type: ignore[arg-type]
|
||||
|
||||
async def test_prompt_contains_category_and_summary(self) -> None:
|
||||
ollama = FakeOllama(['{"consistent": true}'])
|
||||
await taxonomy_consistency("这是摘要内容", "人事制度", ollama) # type: ignore[arg-type]
|
||||
assert "人事制度" in ollama.prompts[0]
|
||||
assert "这是摘要内容" in ollama.prompts[0]
|
||||
assert ollama.json_modes == [True]
|
||||
|
||||
async def test_non_json_output_returns_default_true(self) -> None:
|
||||
ollama = FakeOllama(["无法判定"])
|
||||
assert await taxonomy_consistency("摘要", "制度", ollama) is True # type: ignore[arg-type]
|
||||
|
||||
async def test_missing_consistent_field_returns_default_true(self) -> None:
|
||||
ollama = FakeOllama(['{"ok": 1}'])
|
||||
assert await taxonomy_consistency("摘要", "制度", ollama) is True # type: ignore[arg-type]
|
||||
|
||||
async def test_generate_error_returns_default_true(self) -> None:
|
||||
ollama = FakeOllama(error=RuntimeError("Ollama 不可用"))
|
||||
assert await taxonomy_consistency("摘要", "制度", ollama) is True # type: ignore[arg-type]
|
||||
@@ -0,0 +1,38 @@
|
||||
"""知识分类 API 测试(Qdrant 为 mock,不真实联网)"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.main import app
|
||||
from app.models.knowledge import UNCATEGORIZED, load_taxonomy
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
def test_list_categories(client: TestClient) -> None:
|
||||
"""categories 返回完整 taxonomy,且含 uncategorized"""
|
||||
resp = client.get("/api/v1/knowledge/categories")
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
|
||||
expected = load_taxonomy()
|
||||
data = body["data"]
|
||||
assert data["count"] == len(expected)
|
||||
names = [c["name"] for c in data["categories"]]
|
||||
assert names == [c.name for c in expected]
|
||||
assert UNCATEGORIZED in names
|
||||
@@ -0,0 +1,175 @@
|
||||
"""知识统计 API 测试(GET /api/v1/knowledge/stats)
|
||||
|
||||
两层覆盖:
|
||||
1. TestClient + monkeypatch 注入假 Qdrant 服务:响应结构与异常路径
|
||||
2. 真实内存 Qdrant(location=":memory:")集成:stats 数值与实际写入一致
|
||||
"""
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.api.v1 import knowledge as knowledge_module
|
||||
from app.config import settings
|
||||
from app.main import app
|
||||
from app.services.qdrant import (
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
DIM = settings.embedding_dimension
|
||||
|
||||
|
||||
def _dense(seed: float) -> list[float]:
|
||||
"""构造确定性 dense 向量"""
|
||||
vec = [0.0] * DIM
|
||||
vec[0] = seed
|
||||
vec[1] = 1.0 - seed
|
||||
vec[2] = seed * 0.5
|
||||
vec[3] = 0.1
|
||||
return vec
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]:
|
||||
"""TestClient,lifespan 中的 Qdrant 集合初始化替换为空操作"""
|
||||
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
class _FakeQdrantService:
|
||||
"""假 Qdrant 服务:固定集合计数,scroll_l1 分 2 页返回类目"""
|
||||
|
||||
COUNTS = {COLLECTION_L1: 3, COLLECTION_L2: 6, COLLECTION_L3: 4, COLLECTION_CHUNKS: 9}
|
||||
PAGES = [
|
||||
([{"category": "技术文档"}, {"category": "技术文档"}], "cursor-1"),
|
||||
([{"category": "uncategorized"}], None),
|
||||
]
|
||||
|
||||
async def count(self, collection: str) -> int:
|
||||
return self.COUNTS[collection]
|
||||
|
||||
async def scroll_l1(
|
||||
self, limit: int = 100, offset: str | None = None
|
||||
) -> tuple[list[dict[str, Any]], str | None]:
|
||||
assert limit == 100
|
||||
return self.PAGES[0] if offset is None else self.PAGES[1]
|
||||
|
||||
|
||||
class _FailingQdrantService:
|
||||
"""count 即抛异常的假服务,用于验证错误包装"""
|
||||
|
||||
async def count(self, collection: str) -> int:
|
||||
raise RuntimeError("qdrant 连接失败")
|
||||
|
||||
async def scroll_l1(
|
||||
self, limit: int = 100, offset: str | None = None
|
||||
) -> tuple[list[dict[str, Any]], str | None]:
|
||||
raise RuntimeError("qdrant 连接失败")
|
||||
|
||||
|
||||
def test_stats_ok(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""正常返回:四字段结构完整,数值来自假服务,类目跨页聚合"""
|
||||
monkeypatch.setattr(knowledge_module, "_get_qdrant", lambda: _FakeQdrantService())
|
||||
|
||||
resp = client.get("/api/v1/knowledge/stats")
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert set(data.keys()) == {"collections", "categories", "uncategorized_count", "documents_total"}
|
||||
assert data["collections"] == {"doc_l1": 3, "doc_l2": 6, "doc_l3": 4, "chunks": 9}
|
||||
assert data["categories"] == {"技术文档": 2, "uncategorized": 1}
|
||||
assert data["uncategorized_count"] == 1
|
||||
assert data["documents_total"] == 3
|
||||
|
||||
|
||||
def test_stats_qdrant_error(client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Qdrant 异常 → code 2000 统一错误响应"""
|
||||
monkeypatch.setattr(knowledge_module, "_get_qdrant", lambda: _FailingQdrantService())
|
||||
|
||||
resp = client.get("/api/v1/knowledge/stats")
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 2000
|
||||
assert body["data"] is None
|
||||
assert "获取知识库统计失败" in body["message"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def memory_service() -> QdrantService:
|
||||
"""真实内存 Qdrant 服务,集合并写入 3 篇文档与若干 chunk"""
|
||||
service = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await service.ensure_collections()
|
||||
|
||||
# 3 篇 L1:2 篇同类目 + 1 篇 uncategorized
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-a",
|
||||
title="标题A",
|
||||
summary="总结A",
|
||||
category="技术文档",
|
||||
tags=["api"],
|
||||
dense_vector=_dense(0.9),
|
||||
)
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-b",
|
||||
title="标题B",
|
||||
summary="总结B",
|
||||
category="技术文档",
|
||||
tags=["sdk"],
|
||||
dense_vector=_dense(0.8),
|
||||
)
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-c",
|
||||
title="标题C",
|
||||
summary="总结C",
|
||||
category="uncategorized",
|
||||
tags=[],
|
||||
dense_vector=_dense(0.7),
|
||||
)
|
||||
# 5 个 chunk(doc-a 3 个、doc-c 2 个)
|
||||
chunks = [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"chunk_index": i,
|
||||
"text": f"chunk{i}-{doc_id}",
|
||||
"section_path": "1",
|
||||
"title": f"标题-{doc_id}",
|
||||
"category": category,
|
||||
"tags": [],
|
||||
"dense_vector": _dense(0.5 + i * 0.1),
|
||||
}
|
||||
for doc_id, category, n in (("doc-a", "技术文档", 3), ("doc-c", "uncategorized", 2))
|
||||
for i in range(n)
|
||||
]
|
||||
await service.upsert_chunks(chunks)
|
||||
return service
|
||||
|
||||
|
||||
async def test_stats_with_real_memory_qdrant(
|
||||
memory_service: QdrantService, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""stats 数值与真实内存 Qdrant 写入一致(直接调路由处理函数)"""
|
||||
monkeypatch.setattr(knowledge_module, "_get_qdrant", lambda: memory_service)
|
||||
|
||||
body = await knowledge_module.knowledge_stats()
|
||||
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["collections"] == {"doc_l1": 3, "doc_l2": 0, "doc_l3": 0, "chunks": 5}
|
||||
assert data["documents_total"] == 3
|
||||
assert data["categories"] == {"技术文档": 2, "uncategorized": 1}
|
||||
assert data["uncategorized_count"] == 1
|
||||
@@ -0,0 +1,134 @@
|
||||
"""数据模型与 taxonomy 加载的单元测试"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.models.document import ChunkModel
|
||||
from app.models.knowledge import (
|
||||
UNCATEGORIZED,
|
||||
CategoryResult,
|
||||
TaxonomyCategory,
|
||||
load_taxonomy,
|
||||
)
|
||||
from app.models.search import SearchHit, SearchRequest, SearchResponse
|
||||
|
||||
|
||||
class TestLoadTaxonomy:
|
||||
"""taxonomy 加载与校验"""
|
||||
|
||||
def test_default_taxonomy(self):
|
||||
"""空路径时使用内置默认类目集"""
|
||||
taxonomy = load_taxonomy("")
|
||||
assert len(taxonomy) >= 6
|
||||
names = [c.name for c in taxonomy]
|
||||
assert UNCATEGORIZED in names
|
||||
assert len(names) == len(set(names))
|
||||
assert all(isinstance(c, TaxonomyCategory) for c in taxonomy)
|
||||
|
||||
def test_load_from_json_file(self, tmp_path):
|
||||
"""从 JSON 文件加载自定义类目集"""
|
||||
data = [
|
||||
{"name": "技术", "description": "技术类文档"},
|
||||
{"name": UNCATEGORIZED, "description": "未分类"},
|
||||
]
|
||||
file = tmp_path / "taxonomy.json"
|
||||
file.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
taxonomy = load_taxonomy(str(file))
|
||||
assert [c.name for c in taxonomy] == ["技术", UNCATEGORIZED]
|
||||
assert taxonomy[0].description == "技术类文档"
|
||||
|
||||
def test_append_uncategorized_when_missing(self, tmp_path):
|
||||
"""JSON 文件缺少 uncategorized 时自动追加"""
|
||||
data = [{"name": "技术"}, {"name": "产品"}]
|
||||
file = tmp_path / "taxonomy.json"
|
||||
file.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
taxonomy = load_taxonomy(str(file))
|
||||
names = [c.name for c in taxonomy]
|
||||
assert names == ["技术", "产品", UNCATEGORIZED]
|
||||
|
||||
def test_duplicate_name_raises(self, tmp_path):
|
||||
"""类目 name 重复时报错"""
|
||||
data = [{"name": "技术"}, {"name": "技术"}]
|
||||
file = tmp_path / "taxonomy.json"
|
||||
file.write_text(json.dumps(data, ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="重复"):
|
||||
load_taxonomy(str(file))
|
||||
|
||||
|
||||
class TestCategoryResult:
|
||||
"""CategoryResult 字段校验"""
|
||||
|
||||
def test_defaults(self):
|
||||
result = CategoryResult(main_category="技术文档", confidence=0.9)
|
||||
assert result.tags == []
|
||||
assert result.confidence == 0.9
|
||||
|
||||
def test_confidence_out_of_range(self):
|
||||
"""confidence 越界(>1 或 <0)时校验失败"""
|
||||
with pytest.raises(ValidationError):
|
||||
CategoryResult(main_category="技术文档", confidence=1.5)
|
||||
with pytest.raises(ValidationError):
|
||||
CategoryResult(main_category="技术文档", confidence=-0.1)
|
||||
|
||||
def test_confidence_boundary_values(self):
|
||||
"""confidence 边界值 0 和 1 合法"""
|
||||
assert CategoryResult(main_category="a", confidence=0).confidence == 0
|
||||
assert CategoryResult(main_category="a", confidence=1).confidence == 1
|
||||
|
||||
|
||||
class TestSearchModels:
|
||||
"""检索模型字段校验"""
|
||||
|
||||
def test_search_hit(self):
|
||||
hit = SearchHit(text="内容", doc_id="doc-1", score=0.85)
|
||||
assert hit.title == ""
|
||||
assert hit.section_path == ""
|
||||
assert hit.doc_summary == ""
|
||||
|
||||
def test_search_hit_full_fields(self):
|
||||
hit = SearchHit(
|
||||
text="内容",
|
||||
doc_id="doc-1",
|
||||
title="标题",
|
||||
section_path="第一章/第一节",
|
||||
score=0.85,
|
||||
doc_summary="一句话总结",
|
||||
)
|
||||
assert hit.section_path == "第一章/第一节"
|
||||
|
||||
def test_search_hit_missing_required(self):
|
||||
"""缺少必填字段时校验失败"""
|
||||
with pytest.raises(ValidationError):
|
||||
SearchHit(text="内容", doc_id="doc-1") # type: ignore[call-arg]
|
||||
|
||||
def test_search_request_default_top_k(self):
|
||||
req = SearchRequest(query="如何报销")
|
||||
assert req.top_k is None
|
||||
|
||||
def test_search_response_defaults(self):
|
||||
resp = SearchResponse(query="q", hits=[])
|
||||
assert resp.routed_categories == []
|
||||
assert resp.fallback is False
|
||||
|
||||
|
||||
class TestChunkModel:
|
||||
"""ChunkModel 字段校验"""
|
||||
|
||||
def test_defaults(self):
|
||||
chunk = ChunkModel(doc_id="doc-1", chunk_index=0, text="第一段")
|
||||
assert chunk.section_path == ""
|
||||
|
||||
def test_full_fields(self):
|
||||
chunk = ChunkModel(doc_id="doc-1", chunk_index=2, text="第三段", section_path="第二章")
|
||||
assert chunk.chunk_index == 2
|
||||
assert chunk.section_path == "第二章"
|
||||
|
||||
def test_missing_required(self):
|
||||
"""缺少必填字段时校验失败"""
|
||||
with pytest.raises(ValidationError):
|
||||
ChunkModel(doc_id="doc-1", chunk_index=0) # type: ignore[call-arg]
|
||||
@@ -0,0 +1,126 @@
|
||||
"""OllamaClient 单元测试(mock httpx,不发起真实网络请求)"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from app.services.ollama import OllamaClient
|
||||
|
||||
|
||||
def _make_response(status_code: int, data: dict | None = None) -> httpx.Response:
|
||||
"""构造带 request 上下文的 httpx.Response(raise_for_status 依赖 request)"""
|
||||
request = httpx.Request("POST", "http://localhost:11434/api/generate")
|
||||
if data is None:
|
||||
return httpx.Response(status_code, request=request)
|
||||
return httpx.Response(status_code, json=data, request=request)
|
||||
|
||||
|
||||
class _FakeAsyncClient:
|
||||
"""httpx.AsyncClient 替代品:记录请求参数,返回预设响应或抛出预设异常"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
captured: dict,
|
||||
response: httpx.Response | None = None,
|
||||
error: Exception | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self._captured = captured
|
||||
self._response = response
|
||||
self._error = error
|
||||
captured["timeout"] = kwargs.get("timeout")
|
||||
|
||||
async def __aenter__(self) -> "_FakeAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: object) -> bool:
|
||||
return False
|
||||
|
||||
async def post(self, url: str, json: dict | None = None) -> httpx.Response:
|
||||
self._captured["url"] = url
|
||||
self._captured["json"] = json
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
assert self._response is not None
|
||||
return self._response
|
||||
|
||||
async def get(self, url: str) -> httpx.Response:
|
||||
self._captured["url"] = url
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
assert self._response is not None
|
||||
return self._response
|
||||
|
||||
|
||||
class TestGenerate:
|
||||
async def test_returns_response_field(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(
|
||||
httpx,
|
||||
"AsyncClient",
|
||||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"response": "生成结果"}), **kw),
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434/", model="qwen2.5:1.5b")
|
||||
|
||||
result = await client.generate("你好")
|
||||
|
||||
assert result == "生成结果"
|
||||
# base_url 尾部斜杠被去除
|
||||
assert captured["url"] == "http://localhost:11434/api/generate"
|
||||
assert captured["json"] == {"model": "qwen2.5:1.5b", "prompt": "你好", "stream": False}
|
||||
assert "format" not in captured["json"]
|
||||
assert captured["timeout"] == 120.0
|
||||
|
||||
async def test_json_mode_adds_format(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(
|
||||
httpx,
|
||||
"AsyncClient",
|
||||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"response": "{}"}), **kw),
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||||
|
||||
await client.generate("你好", json_mode=True)
|
||||
|
||||
assert captured["json"]["format"] == "json"
|
||||
|
||||
async def test_http_error_status_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, response=_make_response(500), **kw)
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
await client.generate("你好")
|
||||
|
||||
|
||||
class TestIsAvailable:
|
||||
async def test_200_returns_true(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(
|
||||
httpx,
|
||||
"AsyncClient",
|
||||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"models": []}), **kw),
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||||
|
||||
assert await client.is_available() is True
|
||||
assert captured["url"] == "http://localhost:11434/api/tags"
|
||||
assert captured["timeout"] == 5.0
|
||||
|
||||
async def test_non_200_returns_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, response=_make_response(503), **kw)
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||||
|
||||
assert await client.is_available() is False
|
||||
|
||||
async def test_connect_error_returns_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, error=httpx.ConnectError("连接被拒绝"), **kw)
|
||||
)
|
||||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||||
|
||||
assert await client.is_available() is False
|
||||
@@ -0,0 +1,236 @@
|
||||
"""QdrantService 测试
|
||||
|
||||
使用 AsyncQdrantClient(location=":memory:") 本地模式,无需 Docker。
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.config import settings
|
||||
from app.services.qdrant import (
|
||||
ALL_COLLECTIONS,
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
DIM = settings.embedding_dimension
|
||||
|
||||
|
||||
def _dense(seed: float) -> list[float]:
|
||||
"""构造确定性 dense 向量:前 4 维取特征值,便于区分不同文档"""
|
||||
vec = [0.0] * DIM
|
||||
vec[0] = seed
|
||||
vec[1] = 1.0 - seed
|
||||
vec[2] = seed * 0.5
|
||||
vec[3] = 0.1
|
||||
return vec
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def service() -> QdrantService:
|
||||
svc = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await svc.ensure_collections()
|
||||
return svc
|
||||
|
||||
|
||||
async def test_ensure_collections_idempotent(service: QdrantService) -> None:
|
||||
"""重复调用 ensure_collections 不报错,且 4 个集合均存在"""
|
||||
await service.ensure_collections() # fixture 中已调一次,这里第二次
|
||||
collections = await service.client.get_collections()
|
||||
names = {c.name for c in collections.collections}
|
||||
assert set(ALL_COLLECTIONS) <= names
|
||||
|
||||
# L1 与 chunks 应配置 sparse 命名向量
|
||||
for name in (COLLECTION_L1, COLLECTION_CHUNKS):
|
||||
info = await service.client.get_collection(name)
|
||||
assert info.config.params.sparse_vectors is not None
|
||||
assert "sparse" in info.config.params.sparse_vectors
|
||||
|
||||
|
||||
async def test_upsert_l1_and_filter_by_doc_id(service: QdrantService) -> None:
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-1",
|
||||
title="文档一",
|
||||
summary="这是文档一的总结",
|
||||
category="tech",
|
||||
tags=["ai", "rag"],
|
||||
dense_vector=_dense(0.9),
|
||||
sparse_vector=([1, 2, 3], [0.5, 0.3, 0.2]),
|
||||
)
|
||||
# 重复写入(幂等覆盖)不报错
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-1",
|
||||
title="文档一",
|
||||
summary="这是文档一的总结(更新)",
|
||||
category="tech",
|
||||
tags=["ai", "rag"],
|
||||
dense_vector=_dense(0.9),
|
||||
)
|
||||
|
||||
results = await service.search_dense(
|
||||
COLLECTION_L1,
|
||||
_dense(0.9),
|
||||
limit=5,
|
||||
query_filter=QdrantService.build_filter(doc_ids=["doc-1"]),
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].payload is not None
|
||||
assert results[0].payload["doc_id"] == "doc-1"
|
||||
assert results[0].payload["text"] == "这是文档一的总结(更新)"
|
||||
|
||||
|
||||
async def test_upsert_nodes(service: QdrantService) -> None:
|
||||
nodes = [
|
||||
{
|
||||
"doc_id": "doc-2",
|
||||
"section_path": "1",
|
||||
"text": "第一章大纲",
|
||||
"category": "tech",
|
||||
"tags": ["db"],
|
||||
"dense_vector": _dense(0.8),
|
||||
},
|
||||
{
|
||||
"doc_id": "doc-2",
|
||||
"section_path": "2",
|
||||
"text": "第二章大纲",
|
||||
"category": "tech",
|
||||
"tags": ["db"],
|
||||
"dense_vector": _dense(0.7),
|
||||
},
|
||||
]
|
||||
await service.upsert_nodes(COLLECTION_L2, nodes)
|
||||
await service.upsert_nodes(COLLECTION_L3, nodes)
|
||||
|
||||
for collection in (COLLECTION_L2, COLLECTION_L3):
|
||||
results = await service.search_dense(
|
||||
collection,
|
||||
_dense(0.8),
|
||||
limit=10,
|
||||
query_filter=QdrantService.build_filter(doc_ids=["doc-2"]),
|
||||
)
|
||||
assert len(results) == 2
|
||||
paths = {r.payload["section_path"] for r in results if r.payload}
|
||||
assert paths == {"1", "2"}
|
||||
|
||||
# 非法集合应抛 ValueError
|
||||
with pytest.raises(ValueError):
|
||||
await service.upsert_nodes(COLLECTION_L1, nodes)
|
||||
|
||||
|
||||
async def test_upsert_chunks(service: QdrantService) -> None:
|
||||
chunks = [
|
||||
{
|
||||
"doc_id": "doc-3",
|
||||
"chunk_index": 0,
|
||||
"text": "第一段原文",
|
||||
"section_path": "1",
|
||||
"title": "文档三",
|
||||
"category": "finance",
|
||||
"tags": ["stock"],
|
||||
"dense_vector": _dense(0.6),
|
||||
"sparse_vector": ([10, 20], [1.0, 0.8]),
|
||||
},
|
||||
{
|
||||
"doc_id": "doc-3",
|
||||
"chunk_index": 1,
|
||||
"text": "第二段原文",
|
||||
"section_path": "2",
|
||||
"title": "文档三",
|
||||
"category": "finance",
|
||||
"tags": ["stock"],
|
||||
"dense_vector": _dense(0.4),
|
||||
},
|
||||
]
|
||||
await service.upsert_chunks(chunks)
|
||||
|
||||
results = await service.search_dense(
|
||||
COLLECTION_CHUNKS,
|
||||
_dense(0.6),
|
||||
limit=10,
|
||||
query_filter=QdrantService.build_filter(doc_ids=["doc-3"]),
|
||||
)
|
||||
assert len(results) == 2
|
||||
indices = {r.payload["chunk_index"] for r in results if r.payload}
|
||||
assert indices == {0, 1}
|
||||
|
||||
|
||||
def test_build_filter_empty() -> None:
|
||||
assert QdrantService.build_filter() is None
|
||||
assert QdrantService.build_filter(categories=[], doc_ids=[], section_paths=[]) is None
|
||||
|
||||
|
||||
async def test_build_filter_categories(service: QdrantService) -> None:
|
||||
"""categories 过滤:主类命中与标签命中的文档都能召回,无关类目被排除"""
|
||||
await service.upsert_l1("doc-cat", "主类命中", "总结", category="tech", tags=["x"], dense_vector=_dense(0.9))
|
||||
await service.upsert_l1(
|
||||
"doc-tag", "标签命中", "总结", category="life", tags=["tech", "y"], dense_vector=_dense(0.8)
|
||||
)
|
||||
await service.upsert_l1("doc-none", "无关文档", "总结", category="finance", tags=["z"], dense_vector=_dense(0.7))
|
||||
|
||||
query_filter = QdrantService.build_filter(categories=["tech"])
|
||||
assert query_filter is not None
|
||||
results = await service.search_dense(COLLECTION_L1, _dense(0.9), limit=10, query_filter=query_filter)
|
||||
doc_ids = {r.payload["doc_id"] for r in results if r.payload}
|
||||
assert doc_ids == {"doc-cat", "doc-tag"}
|
||||
|
||||
|
||||
async def test_build_filter_doc_ids(service: QdrantService) -> None:
|
||||
await service.upsert_l1("doc-a", "A", "总结A", category="t", tags=[], dense_vector=_dense(0.9))
|
||||
await service.upsert_l1("doc-b", "B", "总结B", category="t", tags=[], dense_vector=_dense(0.8))
|
||||
|
||||
results = await service.search_dense(
|
||||
COLLECTION_L1,
|
||||
_dense(0.9),
|
||||
limit=10,
|
||||
query_filter=QdrantService.build_filter(doc_ids=["doc-b"]),
|
||||
)
|
||||
assert len(results) == 1
|
||||
assert results[0].payload is not None
|
||||
assert results[0].payload["doc_id"] == "doc-b"
|
||||
|
||||
|
||||
async def test_search_hybrid_rrf(service: QdrantService) -> None:
|
||||
"""hybrid 查询:dense + sparse 两路 prefetch 走服务端 RRF 融合,应正常返回结果"""
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-h1",
|
||||
title="混合一",
|
||||
summary="混合检索文档一",
|
||||
category="tech",
|
||||
tags=["ai"],
|
||||
dense_vector=_dense(0.9),
|
||||
sparse_vector=([100, 200], [1.0, 0.5]),
|
||||
)
|
||||
await service.upsert_l1(
|
||||
doc_id="doc-h2",
|
||||
title="混合二",
|
||||
summary="混合检索文档二",
|
||||
category="tech",
|
||||
tags=["ai"],
|
||||
dense_vector=_dense(0.3),
|
||||
sparse_vector=([100, 300], [0.9, 0.7]),
|
||||
)
|
||||
|
||||
results = await service.search_hybrid(
|
||||
COLLECTION_L1,
|
||||
dense_vector=_dense(0.9),
|
||||
sparse=([100, 200], [1.0, 0.5]),
|
||||
limit=5,
|
||||
)
|
||||
assert len(results) >= 1
|
||||
doc_ids = {r.payload["doc_id"] for r in results if r.payload}
|
||||
assert "doc-h1" in doc_ids
|
||||
|
||||
# hybrid 带过滤也应正常工作
|
||||
filtered = await service.search_hybrid(
|
||||
COLLECTION_L1,
|
||||
dense_vector=_dense(0.9),
|
||||
sparse=([100], [1.0]),
|
||||
limit=5,
|
||||
query_filter=QdrantService.build_filter(doc_ids=["doc-h2"]),
|
||||
)
|
||||
assert len(filtered) == 1
|
||||
assert filtered[0].payload is not None
|
||||
assert filtered[0].payload["doc_id"] == "doc-h2"
|
||||
@@ -0,0 +1,197 @@
|
||||
"""QdrantService 管理操作测试
|
||||
|
||||
使用 AsyncQdrantClient(location=":memory:") 本地模式,无需 Docker。
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
|
||||
from app.config import settings
|
||||
from app.services.qdrant import (
|
||||
ALL_COLLECTIONS,
|
||||
COLLECTION_CHUNKS,
|
||||
COLLECTION_L1,
|
||||
COLLECTION_L2,
|
||||
COLLECTION_L3,
|
||||
QdrantService,
|
||||
)
|
||||
|
||||
DIM = settings.embedding_dimension
|
||||
|
||||
|
||||
def _dense(seed: float) -> list[float]:
|
||||
"""构造确定性 dense 向量:前 4 维取特征值,便于区分不同文档"""
|
||||
vec = [0.0] * DIM
|
||||
vec[0] = seed
|
||||
vec[1] = 1.0 - seed
|
||||
vec[2] = seed * 0.5
|
||||
vec[3] = 0.1
|
||||
return vec
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def service() -> QdrantService:
|
||||
svc = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||||
await svc.ensure_collections()
|
||||
return svc
|
||||
|
||||
|
||||
async def _seed_doc(service: QdrantService, doc_id: str, chunk_count: int = 2) -> None:
|
||||
"""写入一篇完整文档:L1 一条 + L2/L3 各 2 节点 + chunk_count 个 chunk"""
|
||||
await service.upsert_l1(
|
||||
doc_id=doc_id,
|
||||
title=f"标题-{doc_id}",
|
||||
summary=f"总结-{doc_id}",
|
||||
category="tech",
|
||||
tags=["t"],
|
||||
dense_vector=_dense(0.9),
|
||||
)
|
||||
nodes = [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"section_path": str(i),
|
||||
"text": f"节点{i}-{doc_id}",
|
||||
"category": "tech",
|
||||
"tags": ["t"],
|
||||
"dense_vector": _dense(0.8 - i * 0.1),
|
||||
}
|
||||
for i in range(1, 3)
|
||||
]
|
||||
await service.upsert_nodes(COLLECTION_L2, nodes)
|
||||
await service.upsert_nodes(COLLECTION_L3, nodes)
|
||||
chunks = [
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"chunk_index": i,
|
||||
"text": f"chunk{i}-{doc_id}",
|
||||
"section_path": "1",
|
||||
"title": f"标题-{doc_id}",
|
||||
"category": "tech",
|
||||
"tags": ["t"],
|
||||
"dense_vector": _dense(0.5 + i * 0.1),
|
||||
}
|
||||
for i in range(chunk_count)
|
||||
]
|
||||
await service.upsert_chunks(chunks)
|
||||
|
||||
|
||||
async def test_count_empty_and_after_upsert(service: QdrantService) -> None:
|
||||
for collection in ALL_COLLECTIONS:
|
||||
assert await service.count(collection) == 0
|
||||
|
||||
await _seed_doc(service, "doc-count", chunk_count=3)
|
||||
assert await service.count(COLLECTION_L1) == 1
|
||||
assert await service.count(COLLECTION_L2) == 2
|
||||
assert await service.count(COLLECTION_L3) == 2
|
||||
assert await service.count(COLLECTION_CHUNKS) == 3
|
||||
|
||||
|
||||
async def test_scroll_l1_empty(service: QdrantService) -> None:
|
||||
items, next_offset = await service.scroll_l1()
|
||||
assert items == []
|
||||
assert next_offset is None
|
||||
|
||||
|
||||
async def test_scroll_l1_pagination(service: QdrantService) -> None:
|
||||
"""写入 3 篇文档,limit=2 翻页应取全且无重复"""
|
||||
for i in range(3):
|
||||
await service.upsert_l1(
|
||||
doc_id=f"doc-{i}",
|
||||
title=f"标题{i}",
|
||||
summary=f"总结{i}",
|
||||
category="tech",
|
||||
tags=["x"],
|
||||
dense_vector=_dense(0.5 + i * 0.1),
|
||||
)
|
||||
|
||||
seen: list[dict] = []
|
||||
offset: str | None = None
|
||||
pages = 0
|
||||
while True:
|
||||
items, offset = await service.scroll_l1(limit=2, offset=offset)
|
||||
seen.extend(items)
|
||||
pages += 1
|
||||
if offset is None:
|
||||
break
|
||||
assert pages <= 3 # 防止游标异常导致死循环
|
||||
assert pages == 2
|
||||
|
||||
assert len(seen) == 3
|
||||
doc_ids = [item["doc_id"] for item in seen]
|
||||
assert len(set(doc_ids)) == 3 # 无重复
|
||||
assert set(doc_ids) == {"doc-0", "doc-1", "doc-2"}
|
||||
|
||||
# item 字段完整,summary 映射自 payload text
|
||||
for item in seen:
|
||||
assert set(item.keys()) == {"doc_id", "title", "category", "tags", "summary"}
|
||||
i = int(item["doc_id"].rsplit("-", 1)[1])
|
||||
assert item["title"] == f"标题{i}"
|
||||
assert item["summary"] == f"总结{i}"
|
||||
assert item["category"] == "tech"
|
||||
assert item["tags"] == ["x"]
|
||||
|
||||
|
||||
async def test_get_doc_detail_not_found(service: QdrantService) -> None:
|
||||
assert await service.get_doc_detail("doc-missing") is None
|
||||
|
||||
|
||||
async def test_get_doc_detail(service: QdrantService) -> None:
|
||||
await _seed_doc(service, "doc-detail", chunk_count=3)
|
||||
# 干扰数据:不应混入结果
|
||||
await _seed_doc(service, "doc-other", chunk_count=1)
|
||||
|
||||
detail = await service.get_doc_detail("doc-detail")
|
||||
assert detail is not None
|
||||
|
||||
l1 = detail["l1"]
|
||||
assert l1["doc_id"] == "doc-detail"
|
||||
assert l1["title"] == "标题-doc-detail"
|
||||
assert l1["text"] == "总结-doc-detail"
|
||||
assert l1["category"] == "tech"
|
||||
|
||||
assert len(detail["l2_nodes"]) == 2
|
||||
assert len(detail["l3_nodes"]) == 2
|
||||
for nodes in (detail["l2_nodes"], detail["l3_nodes"]):
|
||||
assert {n["section_path"] for n in nodes} == {"1", "2"}
|
||||
for node in nodes:
|
||||
assert node["doc_id"] == "doc-detail"
|
||||
assert "text" in node
|
||||
|
||||
assert detail["chunks_count"] == 3
|
||||
|
||||
|
||||
async def test_delete_by_doc_id(service: QdrantService) -> None:
|
||||
await _seed_doc(service, "doc-del", chunk_count=2)
|
||||
await _seed_doc(service, "doc-keep", chunk_count=1)
|
||||
|
||||
deleted = await service.delete_by_doc_id("doc-del")
|
||||
assert deleted == {
|
||||
COLLECTION_L1: 1,
|
||||
COLLECTION_L2: 2,
|
||||
COLLECTION_L3: 2,
|
||||
COLLECTION_CHUNKS: 2,
|
||||
}
|
||||
|
||||
# 四层该 doc 的点全部清空
|
||||
assert await service.get_doc_detail("doc-del") is None
|
||||
doc_filter = QdrantService.build_filter(doc_ids=["doc-del"])
|
||||
for collection in ALL_COLLECTIONS:
|
||||
result = await service.client.count(collection, count_filter=doc_filter, exact=True)
|
||||
assert result.count == 0
|
||||
|
||||
# 不影响其他 doc 的数据
|
||||
keep_detail = await service.get_doc_detail("doc-keep")
|
||||
assert keep_detail is not None
|
||||
assert keep_detail["chunks_count"] == 1
|
||||
assert len(keep_detail["l2_nodes"]) == 2
|
||||
|
||||
|
||||
async def test_delete_by_doc_id_nonexistent(service: QdrantService) -> None:
|
||||
"""删除不存在的 doc_id:各集合返回 0,已有数据不受影响(幂等)"""
|
||||
await _seed_doc(service, "doc-alive", chunk_count=1)
|
||||
|
||||
deleted = await service.delete_by_doc_id("doc-missing")
|
||||
assert deleted == {name: 0 for name in ALL_COLLECTIONS}
|
||||
|
||||
assert await service.count(COLLECTION_L1) == 1
|
||||
assert await service.count(COLLECTION_CHUNKS) == 1
|
||||
@@ -0,0 +1,238 @@
|
||||
"""query 解析与分类路由的单元测试(mock OllamaClient,不真实联网)"""
|
||||
|
||||
import json
|
||||
|
||||
from app.core.query_parser import CategoryHit, ParsedQuery, QueryParser, decide_route
|
||||
from app.models.knowledge import UNCATEGORIZED, TaxonomyCategory
|
||||
|
||||
|
||||
def _taxonomy() -> list[TaxonomyCategory]:
|
||||
"""测试用 taxonomy 类目集"""
|
||||
return [
|
||||
TaxonomyCategory(name="技术文档", description="架构设计、API 文档、开发规范等技术资料"),
|
||||
TaxonomyCategory(name="产品手册", description="产品功能介绍、使用说明"),
|
||||
TaxonomyCategory(name="财务行政", description="财务制度、报销流程、行政通知"),
|
||||
TaxonomyCategory(name=UNCATEGORIZED, description="无法归入其他类目的文档"),
|
||||
]
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""返回固定响应的假 OllamaClient,记录调用参数"""
|
||||
|
||||
def __init__(self, response: str) -> None:
|
||||
self.response = response
|
||||
self.calls: list[dict] = []
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
self.calls.append({"prompt": prompt, "json_mode": json_mode})
|
||||
return self.response
|
||||
|
||||
|
||||
def _make_parser(response: str) -> QueryParser:
|
||||
return QueryParser(ollama=FakeOllama(response), taxonomy=_taxonomy()) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class TestQueryParserParse:
|
||||
"""QueryParser.parse 的 JSON 解析与容错"""
|
||||
|
||||
async def test_parse_valid_json(self):
|
||||
"""正常 JSON 响应 → 正确解析出 categories/rewrite/keywords"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [{"name": "技术文档", "confidence": 0.9}],
|
||||
"rewrite": "如何设计系统架构",
|
||||
"keywords": ["架构", "设计"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("怎么做架构设计")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert parsed.raw_query == "怎么做架构设计"
|
||||
assert parsed.rewrite == "如何设计系统架构"
|
||||
assert parsed.keywords == ["架构", "设计"]
|
||||
assert len(parsed.categories) == 1
|
||||
assert parsed.categories[0].name == "技术文档"
|
||||
assert parsed.categories[0].confidence == 0.9
|
||||
|
||||
async def test_parse_uses_json_mode_and_prompt_contains_taxonomy(self):
|
||||
"""以 json_mode 调用 LLM,且 prompt 中包含 taxonomy 类目名与描述"""
|
||||
parser = _make_parser('{"categories": [], "rewrite": "q", "keywords": []}')
|
||||
|
||||
await parser.parse("报销流程是什么")
|
||||
|
||||
ollama = parser.ollama
|
||||
assert ollama.calls[0]["json_mode"] is True
|
||||
prompt = ollama.calls[0]["prompt"]
|
||||
assert "技术文档" in prompt and "架构设计" in prompt
|
||||
assert "财务行政" in prompt and "报销流程" in prompt
|
||||
assert "报销流程是什么" in prompt
|
||||
|
||||
async def test_parse_json_with_surrounding_text(self):
|
||||
"""响应带前后多余文本 → 正则提取第一个 {...} 块成功"""
|
||||
response = (
|
||||
"好的,分析结果如下:\n"
|
||||
'{"categories": [{"name": "财务行政", "confidence": 0.8}], "rewrite": "报销流程", "keywords": ["报销"]}\n'
|
||||
"以上就是分析结果。"
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("怎么报销")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert parsed.rewrite == "报销流程"
|
||||
assert parsed.keywords == ["报销"]
|
||||
assert [c.name for c in parsed.categories] == ["财务行政"]
|
||||
|
||||
async def test_parse_non_json_response(self):
|
||||
"""完全非 JSON 响应 → parse_failed=True,rewrite 回退为原 query"""
|
||||
parser = _make_parser("抱歉,我无法理解这个问题。")
|
||||
|
||||
parsed = await parser.parse("blah blah")
|
||||
|
||||
assert parsed.parse_failed is True
|
||||
assert parsed.rewrite == "blah blah"
|
||||
assert parsed.keywords == []
|
||||
assert parsed.categories == []
|
||||
|
||||
async def test_parse_missing_fields(self):
|
||||
"""JSON 合法但必填字段缺失 → parse_failed=True"""
|
||||
parser = _make_parser('{"rewrite": "只有 rewrite"}')
|
||||
|
||||
parsed = await parser.parse("q")
|
||||
|
||||
assert parsed.parse_failed is True
|
||||
assert parsed.rewrite == "q"
|
||||
assert parsed.categories == []
|
||||
|
||||
async def test_unknown_category_dropped(self):
|
||||
"""未知类目名(含 uncategorized)被丢弃,合法类目保留"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [
|
||||
{"name": "不存在的类目", "confidence": 0.9},
|
||||
{"name": UNCATEGORIZED, "confidence": 0.8},
|
||||
{"name": "产品手册", "confidence": 0.7},
|
||||
],
|
||||
"rewrite": "r",
|
||||
"keywords": [],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("q")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert [c.name for c in parsed.categories] == ["产品手册"]
|
||||
|
||||
async def test_out_of_range_confidence_dropped(self):
|
||||
"""confidence 越界(>1 或 <0)的条目被丢弃"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [
|
||||
{"name": "技术文档", "confidence": 1.5},
|
||||
{"name": "产品手册", "confidence": -0.2},
|
||||
{"name": "财务行政", "confidence": 0.6},
|
||||
],
|
||||
"rewrite": "r",
|
||||
"keywords": [],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("q")
|
||||
|
||||
assert [c.name for c in parsed.categories] == ["财务行政"]
|
||||
|
||||
|
||||
class TestDecideRoute:
|
||||
"""decide_route 纯函数的四个分支"""
|
||||
|
||||
def _parsed(self, categories: list[CategoryHit], parse_failed: bool = False) -> ParsedQuery:
|
||||
return ParsedQuery(raw_query="q", rewrite="q", categories=categories, parse_failed=parse_failed)
|
||||
|
||||
def test_parse_failed_fallback(self):
|
||||
"""解析失败 → 全库兜底,reason=parse_failed"""
|
||||
decision = decide_route(self._parsed([], parse_failed=True), threshold=0.6, max_categories=3)
|
||||
|
||||
assert decision.fallback is True
|
||||
assert decision.filter_categories is None
|
||||
assert decision.reason == "parse_failed"
|
||||
|
||||
def test_low_confidence_fallback(self):
|
||||
"""所有类目 confidence 低于阈值 → 全库兜底,reason=low_confidence"""
|
||||
categories = [CategoryHit(name="技术文档", confidence=0.5), CategoryHit(name="产品手册", confidence=0.3)]
|
||||
decision = decide_route(self._parsed(categories), threshold=0.6, max_categories=3)
|
||||
|
||||
assert decision.fallback is True
|
||||
assert decision.filter_categories is None
|
||||
assert decision.reason == "low_confidence"
|
||||
|
||||
def test_empty_categories_fallback(self):
|
||||
"""无命中类目 → 全库兜底,reason=low_confidence"""
|
||||
decision = decide_route(self._parsed([]), threshold=0.6, max_categories=3)
|
||||
|
||||
assert decision.fallback is True
|
||||
assert decision.filter_categories is None
|
||||
assert decision.reason == "low_confidence"
|
||||
|
||||
def test_too_many_categories_fallback(self):
|
||||
"""过阈值类目数超过上限 → 全库兜底,reason=too_many_categories"""
|
||||
categories = [
|
||||
CategoryHit(name="技术文档", confidence=0.9),
|
||||
CategoryHit(name="产品手册", confidence=0.8),
|
||||
CategoryHit(name="财务行政", confidence=0.7),
|
||||
]
|
||||
decision = decide_route(self._parsed(categories), threshold=0.6, max_categories=2)
|
||||
|
||||
assert decision.fallback is True
|
||||
assert decision.filter_categories is None
|
||||
assert decision.reason == "too_many_categories"
|
||||
|
||||
def test_routed_sorted_by_confidence_desc(self):
|
||||
"""正常路由 → 按 confidence 降序输出过滤类目,低于阈值的类目被剔除"""
|
||||
categories = [
|
||||
CategoryHit(name="产品手册", confidence=0.7),
|
||||
CategoryHit(name="财务行政", confidence=0.5), # 低于阈值,应被剔除
|
||||
CategoryHit(name="技术文档", confidence=0.9),
|
||||
]
|
||||
decision = decide_route(self._parsed(categories), threshold=0.6, max_categories=3)
|
||||
|
||||
assert decision.fallback is False
|
||||
assert decision.filter_categories == ["技术文档", "产品手册"]
|
||||
assert decision.reason == "routed"
|
||||
|
||||
|
||||
class TestParseAndRoute:
|
||||
"""parse_and_route 便捷方法(threshold/max_categories 取自 settings,默认 0.6/3)"""
|
||||
|
||||
async def test_parse_and_route_routed(self):
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [{"name": "技术文档", "confidence": 0.9}],
|
||||
"rewrite": "r",
|
||||
"keywords": ["k"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
decision = await parser.parse_and_route("架构设计文档在哪")
|
||||
|
||||
assert decision.fallback is False
|
||||
assert decision.filter_categories == ["技术文档"]
|
||||
assert decision.reason == "routed"
|
||||
assert decision.parsed.raw_query == "架构设计文档在哪"
|
||||
|
||||
async def test_parse_and_route_fallback_on_parse_failure(self):
|
||||
parser = _make_parser("不是 JSON")
|
||||
|
||||
decision = await parser.parse_and_route("q")
|
||||
|
||||
assert decision.fallback is True
|
||||
assert decision.filter_categories is None
|
||||
assert decision.reason == "parse_failed"
|
||||
@@ -0,0 +1,79 @@
|
||||
"""RRF 融合与截断的单元测试"""
|
||||
|
||||
import pytest
|
||||
from qdrant_client import models
|
||||
|
||||
from app.core.ranker import finalize, rrf_fuse
|
||||
|
||||
|
||||
def _point(pid: str, score: float = 1.0) -> models.ScoredPoint:
|
||||
return models.ScoredPoint(id=pid, version=0, score=score, payload={}, vector=None)
|
||||
|
||||
|
||||
class TestRrfFuse:
|
||||
"""rrf_fuse 的排序 / 去重 / 空输入 / k 参数行为"""
|
||||
|
||||
def test_empty_input(self):
|
||||
"""空列表输入返回 []"""
|
||||
assert rrf_fuse([]) == []
|
||||
assert rrf_fuse([[], []]) == []
|
||||
|
||||
def test_single_list_fused_scores(self):
|
||||
"""单路结果按 rank 计算 1/(k+rank) 融合分并降序返回"""
|
||||
fused = rrf_fuse([[_point("a"), _point("b"), _point("c")]], k=60)
|
||||
assert [p.id for p in fused] == ["a", "b", "c"]
|
||||
assert fused[0].score == pytest.approx(1 / 61)
|
||||
assert fused[1].score == pytest.approx(1 / 62)
|
||||
assert fused[2].score == pytest.approx(1 / 63)
|
||||
|
||||
def test_dedup_merges_scores(self):
|
||||
"""同一 point id 出现在多路结果中,只保留一条且分数累加"""
|
||||
list1 = [_point("a"), _point("b")]
|
||||
list2 = [_point("b"), _point("c")]
|
||||
fused = rrf_fuse([list1, list2], k=60)
|
||||
# b 在两路中分别 rank2/rank1,融合分最高;a 与 c 各 1/61
|
||||
assert [p.id for p in fused] == ["b", "a", "c"]
|
||||
assert fused[0].score == pytest.approx(1 / 62 + 1 / 61)
|
||||
assert fused[1].score == pytest.approx(1 / 61)
|
||||
assert fused[2].score == pytest.approx(1 / 62)
|
||||
|
||||
def test_k_affects_ranking(self):
|
||||
"""k 参数改变融合权重,可改变最终排序
|
||||
|
||||
A 排名 [1, 10],B 排名 [5, 5]:
|
||||
- k=60 时 B 总分更高(两路均衡占优)
|
||||
- k=1 时 A 总分更高(头部 rank 权重被放大)
|
||||
"""
|
||||
l1 = [_point("a"), *[_point(f"n{i}") for i in range(3)], _point("b"), *[_point(f"m{i}") for i in range(5)]]
|
||||
l2 = [*[_point(f"x{i}") for i in range(4)], _point("b"), *[_point(f"y{i}") for i in range(4)], _point("a")]
|
||||
# l1:a rank1,b rank5;l2:b rank5,a rank10
|
||||
|
||||
fused_k60 = rrf_fuse([l1, l2], k=60)
|
||||
assert fused_k60[0].id == "b"
|
||||
|
||||
fused_k1 = rrf_fuse([l1, l2], k=1)
|
||||
assert fused_k1[0].id == "a"
|
||||
|
||||
def test_result_is_new_objects(self):
|
||||
"""返回的是写回融合分的新对象,不修改原 point"""
|
||||
original = _point("a", score=0.99)
|
||||
fused = rrf_fuse([[original]], k=60)
|
||||
assert fused[0].score != 0.99
|
||||
assert original.score == 0.99
|
||||
|
||||
|
||||
class TestFinalize:
|
||||
"""finalize 截断行为"""
|
||||
|
||||
def test_truncate(self):
|
||||
points = [_point(str(i)) for i in range(10)]
|
||||
result = finalize(points, 5)
|
||||
assert [p.id for p in result] == ["0", "1", "2", "3", "4"]
|
||||
|
||||
def test_truncate_larger_than_input(self):
|
||||
"""final_k 大于列表长度时返回全部"""
|
||||
points = [_point("a"), _point("b")]
|
||||
assert finalize(points, 5) == points
|
||||
|
||||
def test_empty(self):
|
||||
assert finalize([], 5) == []
|
||||
@@ -0,0 +1,247 @@
|
||||
"""Redis 缓存与缓存集成测试(AsyncMock redis 客户端,不连真实 Redis)"""
|
||||
|
||||
import json
|
||||
from hashlib import sha256
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.config import settings
|
||||
from app.core.query_parser import CategoryHit, ParsedQuery, QueryParser
|
||||
from app.main import app
|
||||
from app.models.knowledge import UNCATEGORIZED, TaxonomyCategory
|
||||
from app.models.search import SearchRequest, SearchResponse
|
||||
from app.services.redis import RedisCache
|
||||
|
||||
|
||||
def _cache_with_client(client: AsyncMock) -> RedisCache:
|
||||
"""构造注入 mock 客户端的 RedisCache(绕过真实 Redis 连接)"""
|
||||
cache = RedisCache()
|
||||
cache._client = client
|
||||
return cache
|
||||
|
||||
|
||||
class _MemoryCache:
|
||||
"""内存版假缓存,接口与 RedisCache 一致"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.store: dict[str, dict[str, Any]] = {}
|
||||
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
return self.store.get(key)
|
||||
|
||||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
self.store[key] = value
|
||||
return True
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class TestRedisCacheGetJson:
|
||||
"""RedisCache.get_json:命中 / 未命中 / 数据异常 / 连接异常"""
|
||||
|
||||
async def test_hit(self):
|
||||
client = AsyncMock()
|
||||
client.get.return_value = json.dumps({"code": 0, "data": {"x": 1}}, ensure_ascii=False)
|
||||
cache = _cache_with_client(client)
|
||||
|
||||
assert await cache.get_json("k") == {"code": 0, "data": {"x": 1}}
|
||||
client.get.assert_awaited_once_with("k")
|
||||
|
||||
async def test_miss(self):
|
||||
client = AsyncMock()
|
||||
client.get.return_value = None
|
||||
|
||||
assert await _cache_with_client(client).get_json("k") is None
|
||||
|
||||
async def test_invalid_json_returns_none(self):
|
||||
client = AsyncMock()
|
||||
client.get.return_value = "not-json{"
|
||||
|
||||
assert await _cache_with_client(client).get_json("k") is None
|
||||
|
||||
async def test_non_dict_json_returns_none(self):
|
||||
client = AsyncMock()
|
||||
client.get.return_value = json.dumps([1, 2, 3])
|
||||
|
||||
assert await _cache_with_client(client).get_json("k") is None
|
||||
|
||||
async def test_error_returns_none(self):
|
||||
client = AsyncMock()
|
||||
client.get.side_effect = ConnectionError("redis down")
|
||||
|
||||
assert await _cache_with_client(client).get_json("k") is None
|
||||
|
||||
|
||||
class TestRedisCacheSetJson:
|
||||
"""RedisCache.set_json:默认 TTL / 自定义 TTL / 异常容错"""
|
||||
|
||||
async def test_ok_uses_default_ttl(self):
|
||||
client = AsyncMock()
|
||||
cache = _cache_with_client(client)
|
||||
|
||||
assert await cache.set_json("k", {"a": 1}) is True
|
||||
client.setex.assert_awaited_once_with("k", settings.cache_ttl, json.dumps({"a": 1}, ensure_ascii=False))
|
||||
|
||||
async def test_ok_custom_ttl(self):
|
||||
client = AsyncMock()
|
||||
cache = _cache_with_client(client)
|
||||
|
||||
assert await cache.set_json("k", {"a": 1}, ttl=10) is True
|
||||
client.setex.assert_awaited_once_with("k", 10, json.dumps({"a": 1}, ensure_ascii=False))
|
||||
|
||||
async def test_error_returns_false(self):
|
||||
client = AsyncMock()
|
||||
client.setex.side_effect = ConnectionError("redis down")
|
||||
|
||||
assert await _cache_with_client(client).set_json("k", {"a": 1}) is False
|
||||
|
||||
|
||||
class _CountingRetriever:
|
||||
"""记录 search 调用次数的假 Retriever"""
|
||||
|
||||
def __init__(self, response: SearchResponse) -> None:
|
||||
self.response = response
|
||||
self.calls = 0
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
self.calls += 1
|
||||
return self.response
|
||||
|
||||
|
||||
def _search_cache_key(query: str, top_k: int | None = None) -> str:
|
||||
"""与路由侧一致的检索缓存键"""
|
||||
return f"search:{sha256((query + '|' + str(top_k)).encode()).hexdigest()[:16]}"
|
||||
|
||||
|
||||
class TestSearchApiCache:
|
||||
"""检索 API 缓存层:命中短路 Retriever,Redis 异常不影响检索"""
|
||||
|
||||
def test_second_request_hits_cache(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""第一次调 Retriever 并回写缓存,第二次缓存命中不再调用"""
|
||||
retriever = _CountingRetriever(SearchResponse(query="q"))
|
||||
cache = _MemoryCache()
|
||||
monkeypatch.setattr("app.api.v1.search._retriever", retriever)
|
||||
monkeypatch.setattr("app.api.v1.search.get_cache", lambda: cache)
|
||||
|
||||
client = TestClient(app)
|
||||
body1 = client.post("/api/v1/search", json={"query": "q"}).json()
|
||||
body2 = client.post("/api/v1/search", json={"query": "q"}).json()
|
||||
|
||||
assert retriever.calls == 1
|
||||
assert body1 == body2
|
||||
assert body2["code"] == 0
|
||||
assert _search_cache_key("q") in cache.store
|
||||
|
||||
def test_different_top_k_uses_different_key(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""top_k 参与缓存键计算:同 query 不同 top_k 不共享缓存"""
|
||||
retriever = _CountingRetriever(SearchResponse(query="q"))
|
||||
cache = _MemoryCache()
|
||||
monkeypatch.setattr("app.api.v1.search._retriever", retriever)
|
||||
monkeypatch.setattr("app.api.v1.search.get_cache", lambda: cache)
|
||||
|
||||
client = TestClient(app)
|
||||
client.post("/api/v1/search", json={"query": "q"})
|
||||
client.post("/api/v1/search", json={"query": "q", "top_k": 3})
|
||||
|
||||
assert retriever.calls == 2
|
||||
assert _search_cache_key("q") in cache.store
|
||||
assert _search_cache_key("q", 3) in cache.store
|
||||
|
||||
def test_redis_error_still_returns_result(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Redis 读写全部抛异常 → 降级为无缓存,检索正常返回"""
|
||||
broken_client = AsyncMock()
|
||||
broken_client.get.side_effect = ConnectionError("redis down")
|
||||
broken_client.setex.side_effect = ConnectionError("redis down")
|
||||
retriever = _CountingRetriever(SearchResponse(query="q"))
|
||||
monkeypatch.setattr("app.api.v1.search._retriever", retriever)
|
||||
monkeypatch.setattr("app.api.v1.search.get_cache", lambda: _cache_with_client(broken_client))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/search", json={"query": "q"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"]["query"] == "q"
|
||||
assert retriever.calls == 1
|
||||
|
||||
|
||||
def _taxonomy() -> list[TaxonomyCategory]:
|
||||
return [
|
||||
TaxonomyCategory(name="技术文档", description="技术资料"),
|
||||
TaxonomyCategory(name=UNCATEGORIZED, description="未分类"),
|
||||
]
|
||||
|
||||
|
||||
class _FakeOllama:
|
||||
"""返回固定响应的假 OllamaClient,记录调用次数"""
|
||||
|
||||
def __init__(self, response: str) -> None:
|
||||
self.response = response
|
||||
self.calls: list[str] = []
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
self.calls.append(prompt)
|
||||
return self.response
|
||||
|
||||
|
||||
def _qparse_cache_key(query: str) -> str:
|
||||
"""与 QueryParser 一致的解析缓存键"""
|
||||
return f"qparse:{sha256(query.encode()).hexdigest()[:16]}"
|
||||
|
||||
|
||||
class TestQueryParserCache:
|
||||
"""QueryParser 解析缓存:命中跳过 LLM,异常回退 LLM"""
|
||||
|
||||
async def test_cache_hit_skips_ollama(self):
|
||||
"""缓存命中时直接用缓存重建 ParsedQuery,不调用 Ollama;decide_route 仍现算"""
|
||||
parsed = ParsedQuery(
|
||||
raw_query="q",
|
||||
rewrite="缓存里的 rewrite",
|
||||
keywords=["k"],
|
||||
categories=[CategoryHit(name="技术文档", confidence=0.9)],
|
||||
)
|
||||
cache = _MemoryCache()
|
||||
cache.store[_qparse_cache_key("q")] = parsed.model_dump()
|
||||
ollama = _FakeOllama("不应被调用")
|
||||
parser = QueryParser(ollama=ollama, taxonomy=_taxonomy(), cache=cache) # type: ignore[arg-type]
|
||||
|
||||
decision = await parser.parse_and_route("q")
|
||||
|
||||
assert ollama.calls == []
|
||||
assert decision.parsed.rewrite == "缓存里的 rewrite"
|
||||
assert decision.fallback is False
|
||||
assert decision.reason == "routed"
|
||||
assert decision.filter_categories == ["技术文档"]
|
||||
|
||||
async def test_llm_result_written_to_cache(self):
|
||||
"""首次走 LLM 并回写缓存,第二次同 query 命中缓存不再调 LLM"""
|
||||
ollama = _FakeOllama('{"categories": [], "rewrite": "r", "keywords": []}')
|
||||
cache = _MemoryCache()
|
||||
parser = QueryParser(ollama=ollama, taxonomy=_taxonomy(), cache=cache) # type: ignore[arg-type]
|
||||
|
||||
await parser.parse_and_route("q")
|
||||
await parser.parse_and_route("q")
|
||||
|
||||
assert len(ollama.calls) == 1
|
||||
assert _qparse_cache_key("q") in cache.store
|
||||
|
||||
async def test_cache_error_falls_back_to_llm(self):
|
||||
"""Redis 读写全部抛异常 → 降级为无缓存,正常走 LLM 解析"""
|
||||
broken_client = AsyncMock()
|
||||
broken_client.get.side_effect = ConnectionError("redis down")
|
||||
broken_client.setex.side_effect = ConnectionError("redis down")
|
||||
ollama = _FakeOllama('{"categories": [], "rewrite": "r", "keywords": []}')
|
||||
parser = QueryParser(
|
||||
ollama=ollama, taxonomy=_taxonomy(), cache=_cache_with_client(broken_client) # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
decision = await parser.parse_and_route("q")
|
||||
|
||||
assert len(ollama.calls) == 1
|
||||
assert decision.parsed.rewrite == "r"
|
||||
assert decision.reason == "low_confidence"
|
||||
@@ -0,0 +1,42 @@
|
||||
"""统一 API 响应包装与业务异常单元测试"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.api.response import ApiError, error, ok
|
||||
|
||||
|
||||
class TestOk:
|
||||
def test_ok_structure(self) -> None:
|
||||
data = {"items": [1, 2, 3]}
|
||||
assert ok(data) == {"code": 0, "data": {"items": [1, 2, 3]}, "message": "ok"}
|
||||
|
||||
def test_ok_none_data(self) -> None:
|
||||
assert ok(None) == {"code": 0, "data": None, "message": "ok"}
|
||||
|
||||
|
||||
class TestError:
|
||||
def test_error_structure(self) -> None:
|
||||
assert error(1004, "文档不存在") == {"code": 1004, "data": None, "message": "文档不存在"}
|
||||
|
||||
def test_error_code_passthrough(self) -> None:
|
||||
result = error(2000, "服务器内部错误")
|
||||
assert result["code"] == 2000
|
||||
assert result["message"] == "服务器内部错误"
|
||||
|
||||
|
||||
class TestApiError:
|
||||
def test_attributes(self) -> None:
|
||||
exc = ApiError(1004, "文档不存在")
|
||||
assert exc.code == 1004
|
||||
assert exc.message == "文档不存在"
|
||||
|
||||
def test_is_exception_with_message(self) -> None:
|
||||
exc = ApiError(1004, "文档不存在")
|
||||
assert isinstance(exc, Exception)
|
||||
assert str(exc) == "文档不存在"
|
||||
|
||||
def test_raise_and_catch(self) -> None:
|
||||
with pytest.raises(ApiError) as exc_info:
|
||||
raise ApiError(1001, "请求参数校验失败")
|
||||
assert exc_info.value.code == 1001
|
||||
assert exc_info.value.message == "请求参数校验失败"
|
||||
@@ -0,0 +1,350 @@
|
||||
"""分层检索引擎与检索 API 的单元测试(全部 mock,不联网、不起 Docker)"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from qdrant_client import models
|
||||
|
||||
from app.config import settings
|
||||
from app.core.query_parser import ParsedQuery, RouteDecision
|
||||
from app.core.retriever import Retriever
|
||||
from app.core.sparse import SparseEncoder
|
||||
from app.main import app
|
||||
from app.models.search import SearchHit, SearchRequest, SearchResponse
|
||||
from app.services.qdrant import COLLECTION_CHUNKS, COLLECTION_L1, COLLECTION_L2, COLLECTION_L3
|
||||
|
||||
|
||||
def _point(
|
||||
pid: str, doc_id: str = "", section_path: str | None = None, score: float = 1.0, **extra: Any
|
||||
) -> models.ScoredPoint:
|
||||
"""构造测试用 ScoredPoint,payload 模拟 chunk/节点结构"""
|
||||
payload: dict[str, Any] = {"doc_id": doc_id, "text": f"text-{pid}", "title": f"title-{doc_id}", **extra}
|
||||
if section_path is not None:
|
||||
payload["section_path"] = section_path
|
||||
return models.ScoredPoint(id=pid, version=0, score=score, payload=payload, vector=None)
|
||||
|
||||
|
||||
class FakeQdrant:
|
||||
"""按集合 + 调用顺序返回预置结果的假 QdrantService,记录每次调用的参数"""
|
||||
|
||||
def __init__(self, results: dict[str, list[list[models.ScoredPoint]]]) -> None:
|
||||
self._results = results
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
def _next(self, collection: str) -> list[models.ScoredPoint]:
|
||||
queue = self._results.get(collection, [])
|
||||
return queue.pop(0) if queue else []
|
||||
|
||||
async def search_dense(
|
||||
self, collection: str, vector: list[float], limit: int, query_filter: models.Filter | None = None
|
||||
) -> list[models.ScoredPoint]:
|
||||
self.calls.append({"collection": collection, "method": "dense", "limit": limit, "filter": query_filter})
|
||||
return self._next(collection)
|
||||
|
||||
async def search_hybrid(
|
||||
self,
|
||||
collection: str,
|
||||
dense_vector: list[float],
|
||||
sparse: tuple[list[int], list[float]],
|
||||
limit: int,
|
||||
query_filter: models.Filter | None = None,
|
||||
) -> list[models.ScoredPoint]:
|
||||
self.calls.append({"collection": collection, "method": "hybrid", "limit": limit, "filter": query_filter})
|
||||
return self._next(collection)
|
||||
|
||||
|
||||
class FakeParser:
|
||||
"""返回固定路由决策的假 QueryParser"""
|
||||
|
||||
def __init__(self, route: RouteDecision) -> None:
|
||||
self.route = route
|
||||
|
||||
async def parse_and_route(self, query: str) -> RouteDecision:
|
||||
return self.route
|
||||
|
||||
|
||||
class FakeEmbedding:
|
||||
"""返回固定向量的假 EmbeddingService"""
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
return [[0.1, 0.2, 0.3] for _ in texts]
|
||||
|
||||
|
||||
def _route(
|
||||
fallback: bool = False, categories: tuple[str, ...] = ("技术文档",), rewrite: str = "rewrite query"
|
||||
) -> RouteDecision:
|
||||
return RouteDecision(
|
||||
fallback=fallback,
|
||||
filter_categories=None if fallback else list(categories),
|
||||
reason="low_confidence" if fallback else "routed",
|
||||
parsed=ParsedQuery(raw_query="q", rewrite=rewrite),
|
||||
)
|
||||
|
||||
|
||||
def _make_retriever(qdrant: FakeQdrant, route: RouteDecision) -> Retriever:
|
||||
return Retriever(
|
||||
qdrant=qdrant, # type: ignore[arg-type]
|
||||
query_parser=FakeParser(route), # type: ignore[arg-type]
|
||||
embedding=FakeEmbedding(),
|
||||
sparse_encoder=SparseEncoder(),
|
||||
)
|
||||
|
||||
|
||||
def _calls(qdrant: FakeQdrant, collection: str) -> list[dict[str, Any]]:
|
||||
return [c for c in qdrant.calls if c["collection"] == collection]
|
||||
|
||||
|
||||
def _must_match_any(flt: models.Filter | None, key: str) -> list[str] | None:
|
||||
"""提取 must 中指定 key 的 MatchAny 值,不存在返回 None"""
|
||||
if flt is None:
|
||||
return None
|
||||
for cond in flt.must or []:
|
||||
if cond.key == key:
|
||||
return list(cond.match.any)
|
||||
return None
|
||||
|
||||
|
||||
def _filter_categories(flt: models.Filter | None) -> list[str] | None:
|
||||
"""提取 min_should 中 category 条件的 MatchAny 值,不存在返回 None"""
|
||||
if flt is None or flt.min_should is None:
|
||||
return None
|
||||
for cond in flt.min_should.conditions:
|
||||
if cond.key == "category":
|
||||
return list(cond.match.any)
|
||||
return None
|
||||
|
||||
|
||||
class TestRetriever:
|
||||
"""分层检索流程:各层 filter 参数与回退路径"""
|
||||
|
||||
async def test_full_pipeline(self):
|
||||
"""正常三级逐层:L1 候选 → L2 收窄 → L3 两路(含 2.5 级文档 b 路)→ chunk"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||||
COLLECTION_L2: [[_point("l2a", "d1", "章节A")]],
|
||||
COLLECTION_L3: [
|
||||
[_point("l3a", "d1", "章节A")], # a 路:L2 命中文档
|
||||
[_point("l3b", "d2", "")], # b 路:2.5 级文档(无 L2 节点)
|
||||
],
|
||||
COLLECTION_CHUNKS: [
|
||||
[
|
||||
_point("c1", "d1", "章节A", doc_summary="总结1"),
|
||||
_point("c2", "d2", "", doc_summary="总结2"),
|
||||
]
|
||||
],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route())
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||||
|
||||
# L1:categories 过滤、无 doc 限制,limit=l1_doc_top_n
|
||||
l1_call = _calls(qdrant, COLLECTION_L1)[0]
|
||||
assert l1_call["limit"] == settings.l1_doc_top_n
|
||||
assert _filter_categories(l1_call["filter"]) == ["技术文档"]
|
||||
assert _must_match_any(l1_call["filter"], "doc_id") is None
|
||||
|
||||
# L2:doc_ids 为 L1 候选(保序),limit=l2_section_top_n*候选数
|
||||
l2_call = _calls(qdrant, COLLECTION_L2)[0]
|
||||
assert l2_call["method"] == "dense"
|
||||
assert l2_call["limit"] == settings.l2_section_top_n * 2
|
||||
assert _must_match_any(l2_call["filter"], "doc_id") == ["d1", "d2"]
|
||||
assert _filter_categories(l2_call["filter"]) == ["技术文档"]
|
||||
|
||||
# L3 两路:a 路 d1 + section「章节A」;b 路 d2(2.5 级文档)仅 doc 过滤
|
||||
l3_calls = _calls(qdrant, COLLECTION_L3)
|
||||
assert len(l3_calls) == 2
|
||||
assert all(c["limit"] == settings.l3_top_n for c in l3_calls)
|
||||
assert _must_match_any(l3_calls[0]["filter"], "doc_id") == ["d1"]
|
||||
assert _must_match_any(l3_calls[0]["filter"], "section_path") == ["章节A"]
|
||||
assert _must_match_any(l3_calls[1]["filter"], "doc_id") == ["d2"]
|
||||
assert _must_match_any(l3_calls[1]["filter"], "section_path") is None
|
||||
|
||||
# chunk:doc_ids 为 L3 命中文档,section_paths 仅非空值
|
||||
chunk_call = _calls(qdrant, COLLECTION_CHUNKS)[0]
|
||||
assert chunk_call["limit"] == settings.retrieval_top_k
|
||||
assert sorted(_must_match_any(chunk_call["filter"], "doc_id") or []) == ["d1", "d2"]
|
||||
assert _must_match_any(chunk_call["filter"], "section_path") == ["章节A"]
|
||||
|
||||
# 响应组装
|
||||
assert resp.query == "测试查询"
|
||||
assert resp.fallback is False
|
||||
assert resp.routed_categories == ["技术文档"]
|
||||
assert [h.doc_id for h in resp.hits] == ["d1", "d2"]
|
||||
assert resp.hits[0].text == "text-c1"
|
||||
assert resp.hits[0].title == "title-d1"
|
||||
assert resp.hits[0].section_path == "章节A"
|
||||
assert resp.hits[0].doc_summary == "总结1"
|
||||
assert resp.hits[0].score > 0
|
||||
|
||||
async def test_l2_empty_all_docs_go_l3_b_path(self):
|
||||
"""L2 整体无命中:全部候选文档回退为 L3 单路 doc 级查询"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||||
COLLECTION_L2: [[]],
|
||||
COLLECTION_L3: [[_point("l3a", "d1", "章节A")]],
|
||||
COLLECTION_CHUNKS: [[_point("c1", "d1", "章节A")]],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route())
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||||
|
||||
l3_calls = _calls(qdrant, COLLECTION_L3)
|
||||
assert len(l3_calls) == 1
|
||||
assert _must_match_any(l3_calls[0]["filter"], "doc_id") == ["d1", "d2"]
|
||||
assert _must_match_any(l3_calls[0]["filter"], "section_path") is None
|
||||
assert [h.doc_id for h in resp.hits] == ["d1"]
|
||||
|
||||
async def test_l3_empty_fallback_doc_level_chunks(self):
|
||||
"""L3 两路均无命中:chunk 层回退为 L1 候选文档级检索(无 section 过滤)"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||||
COLLECTION_L2: [[_point("l2a", "d1", "章节A")]],
|
||||
COLLECTION_L3: [[], []],
|
||||
COLLECTION_CHUNKS: [[_point("c1", "d2", "")]],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route())
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||||
|
||||
chunk_call = _calls(qdrant, COLLECTION_CHUNKS)[0]
|
||||
assert _must_match_any(chunk_call["filter"], "doc_id") == ["d1", "d2"]
|
||||
assert _must_match_any(chunk_call["filter"], "section_path") is None
|
||||
assert [h.doc_id for h in resp.hits] == ["d2"]
|
||||
|
||||
async def test_l1_empty_global_chunk_fallback(self):
|
||||
"""L1 无候选文档:直接全库 chunk 兜底(无 filter),fallback=True"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[]],
|
||||
COLLECTION_CHUNKS: [[_point("c1", "d9", "章节X", doc_summary="总结9")]],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route())
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||||
|
||||
# 不再触发 L2/L3 查询
|
||||
assert _calls(qdrant, COLLECTION_L2) == []
|
||||
assert _calls(qdrant, COLLECTION_L3) == []
|
||||
# 全库 chunk 检索:无 filter
|
||||
chunk_calls = _calls(qdrant, COLLECTION_CHUNKS)
|
||||
assert len(chunk_calls) == 1
|
||||
assert chunk_calls[0]["filter"] is None
|
||||
assert chunk_calls[0]["limit"] == settings.retrieval_top_k
|
||||
assert resp.fallback is True
|
||||
assert [h.doc_id for h in resp.hits] == ["d9"]
|
||||
|
||||
async def test_route_fallback_no_category_filter(self):
|
||||
"""路由兜底时 categories=None,各层均不做类目过滤"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[_point("l1a", "d1")]],
|
||||
COLLECTION_L2: [[]],
|
||||
COLLECTION_L3: [[_point("l3a", "d1", "")]],
|
||||
COLLECTION_CHUNKS: [[_point("c1", "d1", "")]],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route(fallback=True))
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||||
|
||||
# L1:categories 与 doc_ids 均空 → filter 为 None
|
||||
assert _calls(qdrant, COLLECTION_L1)[0]["filter"] is None
|
||||
# L2/L3/chunk:仅有 doc 级 must 条件,无类目 min_should
|
||||
for collection in (COLLECTION_L2, COLLECTION_L3, COLLECTION_CHUNKS):
|
||||
for call in _calls(qdrant, collection):
|
||||
assert _filter_categories(call["filter"]) is None
|
||||
assert resp.fallback is True
|
||||
assert resp.routed_categories == []
|
||||
|
||||
async def test_top_k_override(self):
|
||||
"""request.top_k 优先于 settings.retrieval_final_k"""
|
||||
qdrant = FakeQdrant(
|
||||
{
|
||||
COLLECTION_L1: [[_point("l1a", "d1")]],
|
||||
COLLECTION_L2: [[]],
|
||||
COLLECTION_L3: [[_point("l3a", "d1", "")]],
|
||||
COLLECTION_CHUNKS: [[_point(f"c{i}", "d1", "") for i in range(5)]],
|
||||
}
|
||||
)
|
||||
retriever = _make_retriever(qdrant, _route())
|
||||
|
||||
resp = await retriever.search(SearchRequest(query="测试查询", top_k=2))
|
||||
assert len(resp.hits) == 2
|
||||
|
||||
|
||||
class _FakeRetriever:
|
||||
"""API 测试用假 Retriever:返回固定响应或抛异常"""
|
||||
|
||||
def __init__(self, response: SearchResponse | None = None, exc: Exception | None = None) -> None:
|
||||
self.response = response
|
||||
self.exc = exc
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
if self.exc is not None:
|
||||
raise self.exc
|
||||
assert self.response is not None
|
||||
return self.response
|
||||
|
||||
|
||||
class TestSearchApi:
|
||||
"""检索 API 统一响应包装"""
|
||||
|
||||
def test_search_ok(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""正常检索 → {"code": 0, "data": ..., "message": "ok"}"""
|
||||
response = SearchResponse(
|
||||
query="q",
|
||||
hits=[SearchHit(text="t", doc_id="d1", title="标题", section_path="", score=0.5, doc_summary="s")],
|
||||
routed_categories=["技术文档"],
|
||||
fallback=False,
|
||||
)
|
||||
monkeypatch.setattr("app.api.v1.search._retriever", _FakeRetriever(response=response))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/search", json={"query": "q"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["message"] == "ok"
|
||||
assert body["data"]["query"] == "q"
|
||||
assert body["data"]["hits"][0]["doc_id"] == "d1"
|
||||
assert body["data"]["routed_categories"] == ["技术文档"]
|
||||
assert body["data"]["fallback"] is False
|
||||
|
||||
def test_search_error(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""检索内部异常 → code 2000"""
|
||||
monkeypatch.setattr("app.api.v1.search._retriever", _FakeRetriever(exc=RuntimeError("boom")))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/search", json={"query": "q"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 2000
|
||||
assert body["data"] is None
|
||||
assert "boom" in body["message"]
|
||||
|
||||
def test_validation_error(self):
|
||||
"""缺少必填 query 字段 → code 1001"""
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/search", json={})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert body["data"] is None
|
||||
|
||||
def test_health_kept(self):
|
||||
"""现有健康检查接口不受影响"""
|
||||
client = TestClient(app)
|
||||
resp = client.get("/api/v1/health")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json() == {"status": "ok"}
|
||||
@@ -0,0 +1,131 @@
|
||||
"""run_eval 门槛常量与报告组装纯函数单元测试"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from scripts.eval.run_eval import (
|
||||
THRESHOLD_HALLUCINATION,
|
||||
THRESHOLD_L1_ER,
|
||||
THRESHOLD_L3_ER,
|
||||
THRESHOLD_PRUNING_LOSS,
|
||||
_build_report,
|
||||
_fmt,
|
||||
_mark,
|
||||
)
|
||||
|
||||
|
||||
class TestThresholds:
|
||||
"""Spec 规定的门槛值:L1 ER ≥ 0.85 / L3 ER ≥ 0.9 / 幻觉率 < 2% / Pruning Loss < 8%"""
|
||||
|
||||
def test_l1_entity_recall_threshold(self) -> None:
|
||||
assert THRESHOLD_L1_ER == 0.85
|
||||
|
||||
def test_l3_entity_recall_threshold(self) -> None:
|
||||
assert THRESHOLD_L3_ER == 0.9
|
||||
|
||||
def test_hallucination_threshold(self) -> None:
|
||||
assert THRESHOLD_HALLUCINATION == 0.02
|
||||
|
||||
def test_pruning_loss_threshold(self) -> None:
|
||||
assert THRESHOLD_PRUNING_LOSS == 0.08
|
||||
|
||||
|
||||
class TestMark:
|
||||
def test_pass(self) -> None:
|
||||
assert _mark(True) == "✅"
|
||||
|
||||
def test_fail(self) -> None:
|
||||
assert _mark(False) == "❌"
|
||||
|
||||
|
||||
class TestFmt:
|
||||
def test_none_returns_na(self) -> None:
|
||||
assert _fmt(None) == "N/A"
|
||||
assert _fmt(None, percent=True) == "N/A"
|
||||
|
||||
def test_decimal_format(self) -> None:
|
||||
assert _fmt(0.85) == "0.8500"
|
||||
|
||||
def test_percent_format(self) -> None:
|
||||
assert _fmt(0.1234, percent=True) == "12.3%"
|
||||
assert _fmt(0.02, percent=True) == "2.0%"
|
||||
|
||||
|
||||
def _doc_record(**overrides: Any) -> dict:
|
||||
record: dict[str, Any] = {
|
||||
"id": "d1",
|
||||
"title": "考勤制度",
|
||||
"golden_category": "制度",
|
||||
"category": "制度",
|
||||
"entity_recall_l1": 0.9,
|
||||
"entity_recall_l3": 0.95,
|
||||
"hallucination_rate": 0.01,
|
||||
"taxonomy_consistency": True,
|
||||
}
|
||||
record.update(overrides)
|
||||
return record
|
||||
|
||||
|
||||
def _report_kwargs(**overrides: Any) -> dict:
|
||||
kwargs: dict[str, Any] = {
|
||||
"regression_path": Path("regression_set.json"),
|
||||
"doc_records": [_doc_record()],
|
||||
"query_summary": {"total": 4, "positive": 3, "negative": 1, "negative_false_alarm": 0},
|
||||
"hier_metrics": {"precision@5": 0.6, "recall@10": 0.8},
|
||||
"baseline_metrics": {"precision@5": 0.4, "recall@10": 0.7},
|
||||
"routing": {"precision": 0.75, "recall": 0.8, "f1": 0.77},
|
||||
"prune_loss": 0.05,
|
||||
"judge_enabled": True,
|
||||
"kept_data": False,
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return kwargs
|
||||
|
||||
|
||||
class TestBuildReport:
|
||||
def test_baseline_comparison_columns(self) -> None:
|
||||
report = _build_report(**_report_kwargs())
|
||||
assert "| 指标 | 分层检索 | 平铺 baseline |" in report
|
||||
assert "| Precision@5 | 0.6000 | 0.4000 |" in report
|
||||
assert "| Recall@10 | 0.8000 | 0.7000 |" in report
|
||||
|
||||
def test_doc_table_row(self) -> None:
|
||||
report = _build_report(**_report_kwargs())
|
||||
assert "| d1 | 考勤制度 | 制度 | 制度 | 0.9000 | 0.9500 | 1.0% | 是 |" in report
|
||||
|
||||
def test_threshold_pass_marks(self) -> None:
|
||||
report = _build_report(**_report_kwargs())
|
||||
assert "| L1 Entity Recall | 0.9000 | ≥ 0.85 | ✅ |" in report
|
||||
assert "| L3 Entity Recall | 0.9500 | ≥ 0.9 | ✅ |" in report
|
||||
assert "| Hallucination Rate | 1.0% | < 2% | ✅ |" in report
|
||||
assert "| Pruning Loss | 5.0% | < 8% | ✅ |" in report
|
||||
|
||||
def test_threshold_fail_marks(self) -> None:
|
||||
report = _build_report(
|
||||
**_report_kwargs(
|
||||
doc_records=[
|
||||
_doc_record(
|
||||
entity_recall_l1=0.8,
|
||||
entity_recall_l3=0.7,
|
||||
hallucination_rate=0.05,
|
||||
taxonomy_consistency=False,
|
||||
)
|
||||
],
|
||||
prune_loss=0.2,
|
||||
)
|
||||
)
|
||||
assert "| L1 Entity Recall | 0.8000 | ≥ 0.85 | ❌ |" in report
|
||||
assert "| L3 Entity Recall | 0.7000 | ≥ 0.9 | ❌ |" in report
|
||||
assert "| Hallucination Rate | 5.0% | < 2% | ❌ |" in report
|
||||
assert "| Pruning Loss | 20.0% | < 8% | ❌ |" in report
|
||||
assert " | 否 |" in report
|
||||
|
||||
def test_judge_skipped_renders_na(self) -> None:
|
||||
report = _build_report(
|
||||
**_report_kwargs(
|
||||
doc_records=[_doc_record(hallucination_rate=None, taxonomy_consistency=None)],
|
||||
judge_enabled=False,
|
||||
)
|
||||
)
|
||||
assert "N/A" in report
|
||||
assert "跳过" in report
|
||||
@@ -0,0 +1,57 @@
|
||||
"""稀疏向量编码器单元测试"""
|
||||
|
||||
from app.core.sparse import SPARSE_DIM, SparseEncoder
|
||||
|
||||
|
||||
class TestSparseEncoder:
|
||||
"""SparseEncoder.encode / encode_batch 行为"""
|
||||
|
||||
def setup_method(self):
|
||||
self.encoder = SparseEncoder()
|
||||
|
||||
def test_deterministic(self):
|
||||
"""同一文本多次编码结果完全一致"""
|
||||
text = "人工智能 artificial intelligence 2024"
|
||||
assert self.encoder.encode(text) == self.encoder.encode(text)
|
||||
|
||||
def test_indices_sorted_unique(self):
|
||||
"""indices 升序且无重复,与 values 等长,且都在 [0, SPARSE_DIM) 内"""
|
||||
indices, values = self.encoder.encode("向量检索与关键词检索相结合 hybrid search 123")
|
||||
assert indices == sorted(indices)
|
||||
assert len(indices) == len(set(indices))
|
||||
assert len(indices) == len(values)
|
||||
assert all(0 <= idx < SPARSE_DIM for idx in indices)
|
||||
assert all(v > 0 for v in values)
|
||||
|
||||
def test_mixed_text_non_empty(self):
|
||||
"""中英文混合文本产生非空向量"""
|
||||
indices, values = self.encoder.encode("使用 Qdrant 存储向量")
|
||||
assert len(indices) > 0
|
||||
assert len(values) > 0
|
||||
|
||||
def test_empty_string(self):
|
||||
"""空字符串返回空向量"""
|
||||
assert self.encoder.encode("") == ([], [])
|
||||
|
||||
def test_different_texts_differ(self):
|
||||
"""不同文本产生不同向量"""
|
||||
assert self.encoder.encode("机器学习") != self.encoder.encode("深度学习")
|
||||
|
||||
def test_english_case_insensitive(self):
|
||||
"""英文词小写化:大小写不同编码结果一致"""
|
||||
assert self.encoder.encode("Hello") == self.encoder.encode("hello")
|
||||
|
||||
def test_tf_weighting(self):
|
||||
"""词频越高权重越大(1 + log(tf))"""
|
||||
once = dict(zip(*self.encoder.encode("apple"), strict=True))
|
||||
twice = dict(zip(*self.encoder.encode("apple apple"), strict=True))
|
||||
assert once.keys() == twice.keys()
|
||||
for idx in once:
|
||||
assert twice[idx] > once[idx]
|
||||
|
||||
def test_batch_matches_single(self):
|
||||
"""encode_batch 与逐条 encode 结果一致"""
|
||||
texts = ["人工智能", "vector search", ""]
|
||||
batch = self.encoder.encode_batch(texts)
|
||||
singles = [self.encoder.encode(t) for t in texts]
|
||||
assert batch == singles
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Summarizer 标题树接入的单元测试"""
|
||||
|
||||
from app.core.summarizer import Summarizer
|
||||
from app.models.document import SummaryLevel
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""记录调用次数的 OllamaClient 替身"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.prompts: list[str] = []
|
||||
|
||||
async def generate(self, prompt: str) -> str:
|
||||
self.prompts.append(prompt)
|
||||
return "LLM 生成结果"
|
||||
|
||||
|
||||
def _structured_text() -> str:
|
||||
"""构造带标题结构且长度 ≥ 500 的文档(不触发 2.5 回退)"""
|
||||
paragraph = "这是章节正文内容,包含足够多的信息量。" * 10
|
||||
return f"# 安装指南\n{paragraph}\n## 环境准备\n{paragraph}\n## 安装步骤\n{paragraph}"
|
||||
|
||||
|
||||
class TestStructuredDocument:
|
||||
"""结构化文档:L2 直接使用标题树,不调 LLM"""
|
||||
|
||||
async def test_l2_from_headings_skips_llm(self):
|
||||
ollama = FakeOllama()
|
||||
summarizer = Summarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
|
||||
summary = await summarizer.summarize(_structured_text(), title="安装文档")
|
||||
|
||||
assert summary.level == SummaryLevel.L3
|
||||
# L1 + L3 两次调用,跳过了 L2 的 LLM 调用
|
||||
assert len(ollama.prompts) == 2
|
||||
# 大纲由标题树渲染,包含标题文本与层级缩进
|
||||
assert summary.l2_outline is not None
|
||||
assert "- 安装指南" in summary.l2_outline
|
||||
assert " - 环境准备" in summary.l2_outline
|
||||
assert " - 安装步骤" in summary.l2_outline
|
||||
assert summary.l2_outline != "LLM 生成结果"
|
||||
# L3 的 prompt 中注入了标题树大纲
|
||||
assert "安装指南" in ollama.prompts[1]
|
||||
|
||||
|
||||
class TestPlainDocument:
|
||||
"""无结构文档:L2 回退 LLM 生成"""
|
||||
|
||||
async def test_l2_falls_back_to_llm(self):
|
||||
ollama = FakeOllama()
|
||||
summarizer = Summarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
text = "这是一段没有标题结构的正文内容," * 40 # 640 字符,不触发 2.5 回退
|
||||
|
||||
summary = await summarizer.summarize(text, title="")
|
||||
|
||||
assert summary.level == SummaryLevel.L3
|
||||
# L1 + L2 + L3 三次调用
|
||||
assert len(ollama.prompts) == 3
|
||||
assert summary.l2_outline == "LLM 生成结果"
|
||||
|
||||
|
||||
class TestFallbackUnaffected:
|
||||
"""2.5 级回退逻辑不受标题树改造影响"""
|
||||
|
||||
async def test_short_structured_text_still_fallback(self):
|
||||
ollama = FakeOllama()
|
||||
summarizer = Summarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
text = "# 标题一\n简短内容。\n# 标题二\n简短内容。"
|
||||
|
||||
summary = await summarizer.summarize(text, title="")
|
||||
|
||||
assert summary.level == SummaryLevel.L2_HALF
|
||||
assert summary.l2_outline is None
|
||||
# L1 + L2.5 两次调用
|
||||
assert len(ollama.prompts) == 2
|
||||
Reference in New Issue
Block a user