Initial commit: QMDSearch 分层信息检索服务

- FastAPI + Qdrant + Redis + Ollama 技术栈
- L1→L2→L3→chunk 四层分层检索(dense + sparse RRF 融合)
- 文档三级总结与 2.5 级回退
- query 解析路由与分类
- /admin 管理页面
This commit is contained in:
2026-07-29 21:24:40 +08:00
commit 51dc8dc4f6
83 changed files with 10794 additions and 0 deletions
+74
View File
@@ -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
View File
@@ -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 与 messagepartial_summary(已产出总结)保留
- [x] 状态写 Redisdone 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 附 resultfailed 附 error.stage/message
- [x] 查询不存在 task_id 返回 code=1004
- [x] 既有入库相关测试全部适配新契约且通过
## 管理页面
- [x] 入库提交后展示 task_id 与状态,2s 轮询直至 done/failed
- [x] done 展示 category/置信度/tags/总结层级/chunks_countfailed 展示 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` 由同步返回 IngestionResult200)改为立即返回 `{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`(默认 604800s7 天)
- 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`(完整 IngestionResultdocument_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_countfailed 展示 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: 配置 + IngestTaskManagerconfig 新增 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-kl1/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/chunkspayload 含 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/QdrantRedis 宕机检索仍可用
## 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 RecallL1/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 + sparseBM25)双向量生成、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 总结生成 sparseBM25)向量。dense 向量提供方可配置(openai / local),sparse 向量使用本地分词 + 特征哈希 BM25 实现,不依赖外部模型下载。
#### Scenario: 四层集合写入
- **WHEN** 文档完成总结与分类
- **THEN** L1/L2/L3 各层节点与原文 chunk 分别写入对应集合,每条数据携带 `doc_id``category``tags``section_path`chunk/L2/L3 层)payloadchunk 同时携带 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 RateLLM-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` 增加 ChunkModeldoc_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`denseopenai | local 双 provider 统一异步接口,批量编码)+ `app/core/sparse.py`(本地分词 + 特征哈希 BM25 稀疏向量,输出 Qdrant sparse vector 格式,不依赖外部模型)
- [x] SubTask 2.1: EmbeddingService 抽象 + openai provider 实现 + local providerOllama embedding 接口)实现
- [x] SubTask 2.2: SparseEncoder(中文分词 + BM25 权重 + 特征哈希到固定维度),单测验证输出稀疏格式合法
- [x] Task 3: Qdrant 服务封装 `app/services/qdrant.py`:启动时初始化 4 个集合(doc_l1 / doc_l2 / doc_l3 / chunkschunks 与 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@10Hallucination 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 分页返回 itemsdoc_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.pyok/error/ApiError 结构与 code
- [x] ollama.pygenerate 正常/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.pyHTTP 行为)、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=0data 含 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.pymock Ollama:断言抽取/支持判定/解析失败返回 0.0)、response.pyok/error/ApiError 结构)、ollama.pymock httpxgenerate/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_detailL1+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.pymock Ollama:正常判定/部分断言不支持/解析失败返回 0.0)、tests/test_response.py、tests/test_ollama_client.pymock httpxgenerate/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 完成后)
+160
View File
@@ -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 + sparseRRF 融合)。
0. **query 解析路由**: 用 Ollama 小模型将 query 解析为结构化结果(rewrite、关键词、命中类目+置信度);高置信且类目数不超上限时按类目过滤,低置信/解析失败/类目过多则全库兜底(不丢召回)
1. **L1 - 文档定位层**: 在文档 L1 总结集合中检索,产出候选文档集合;无命中时直接全库 chunk 兜底(fallback=True
2. **L2 - 章节大纲层**: 在候选文档内检索 L2 章节大纲,按 section 剪枝定位;未命中的文档(含 2.5 级文档)落入 L3 b 路(仅按 doc 过滤)
3. **L3 - 小节内容层**: 两路查询(aL2 命中文档按 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 附 resultfailed 附 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 保留 24hfailed 保留 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 和各级总结分别生成 embeddingdense + 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
View File
@@ -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"]
View File
View File
+26
View File
@@ -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)
View File
+111
View File
@@ -0,0 +1,111 @@
"""文档 APIPOST /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_atdone 附 resultfailed 附 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())})
+75
View File
@@ -0,0 +1,75 @@
"""知识分类 APIGET /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
+58
View File
@@ -0,0 +1,58 @@
"""检索 APIPOST /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
+60
View File
@@ -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()
View File
+121
View File
@@ -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 为祖先标题链(含自身标题),用 " / " 连接;
首个标题前的引导正文归入无前缀 sectionsection_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
+139
View File
@@ -0,0 +1,139 @@
"""文档分类器
入库链路第二步:基于 L1 总结,用 Ollama 小模型将文档判定为 taxonomy 中的
主类目 + 附加标签:
- LLM 输出解析失败 / 类目名不在 taxonomy → 归 uncategorizedconfidence=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]
+99
View File
@@ -0,0 +1,99 @@
"""Dense 向量嵌入服务
提供统一的 EmbeddingService 接口,支持两种 provider
- openaiOpenAI 兼容 APIAsyncOpenAI
- 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,
)
+94
View File
@@ -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
+172
View File
@@ -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_donefailed 用 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)
+306
View File
@@ -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. 批量 embeddingL1 + 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)
+214
View File
@@ -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
+24
View File
@@ -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]:
"""标准 RRFReciprocal 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]
+169
View File
@@ -0,0 +1,169 @@
"""分层检索引擎
L1(文档总结)→ L2(章节大纲)→ L3(小节定位)→ chunks(原文)逐层收窄:
- L1 无候选文档:直接全库 chunk 兜底,fallback=True
- L2 无命中:全部候选文档回退到 L3 的 doc 级查询(b 路)
- 2.5 级文档无 L2 节点,天然落入 L3 b 路(仅按 doc 过滤)
- L3 两路(aL2 命中文档按 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 payloaddoc_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
+68
View File
@@ -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]
+138
View File
@@ -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
View File
@@ -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")
View File
+51
View File
@@ -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="主类目分类置信度")
+68
View File
@@ -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
+30
View File
@@ -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="是否走了全库兜底路径")
View File
+62
View File
@@ -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
+351
View File
@@ -0,0 +1,351 @@
"""Qdrant 向量数据库服务封装
分层摘要索引 RAG 的存储层:
- doc_l1:文档级总结(dense + sparse
- doc_l2 / doc_l3:大纲节点(dense
- chunks:原文分块(dense + sparse
提供集合初始化(幂等)、按层 upsert、dense / hybridRRF 融合)检索接口。
"""
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_pathsmust 条件 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)
+75
View File
@@ -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
+692
View File
@@ -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>
View File
+78
View File
@@ -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
+41
View File
@@ -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"]
View File
+93
View File
@@ -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"])
+106
View File
@@ -0,0 +1,106 @@
"""离线评测指标:纯函数实现,不依赖 Qdrant / Ollama,可独立单测
- entity_recall:实体保留率(摘要质量)
- routing_f1L1 路由层 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 querynegative 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
+73
View File
@@ -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 默认 1536sparse_enabled 默认开启。按照《运维规范》第三条要求:修改 retrieval_top_k 或 l1_doc_top_n 后必须重启服务,并执行回归冒烟用例。\n# 故障排查\n若 Qdrant 连接超时,先检查容器网络 qmd-net 是否互通;Ollama 拉取模型失败时重试 ollama pull qwen2.5:1.5bRedis 不可用时会自动降级为直连检索,不影响主流程。",
"golden_category": "技术文档"
},
{
"id": "doc002",
"title": "Embedding 服务接入指南",
"text": "# 接入方式\nEmbeddingService 统一接口支持 openai 与 local 两种 provider。openai provider 默认模型 text-embedding-3-small,走 OpenAI 兼容 APIlocal 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}
]
}
+353
View File
@@ -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())
+165
View File
@@ -0,0 +1,165 @@
"""真实 Ollama 现场冒烟脚本
用法:uv run python scripts/smoke_live.py
链路:内存 Qdrant + 真实本地 Ollamadeepseek-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 '(None2.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()))
View File
+329
View File
@@ -0,0 +1,329 @@
"""管理闭环集成测试(Spec Task 7:集成验证)
两部分覆盖,均不修改 app/ 与 scripts/ 实现代码:
1. TestAdminClosedLoop:真实内存 Qdrantlocation=":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. statsdocuments_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. 删除文档 Adeleted_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})"
+125
View 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]:
"""TestClientlifespan 中的 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 → 200content-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
+130
View File
@@ -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 == "第二章"
+133
View File
@@ -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 调用 LLMprompt 包含全部类目名与 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 响应 → 归 uncategorizedtags 为空,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 → 归 uncategorizedconfidence=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,候选类目名保留进 tagsconfidence 保留原值"""
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
+171
View File
@@ -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]:
"""TestClientlifespan 中的 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
+74
View File
@@ -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]:
"""TestClientlifespan 中的 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 202code=0data 含 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"
+328
View File
@@ -0,0 +1,328 @@
"""端到端集成测试
真实 Ingester + Retriever + 内存 Qdrantlocation=":memory:")串联分层 RAG 全链路,
仅替换两个外部边界:
- FakeOllama:按 prompt 内容返回确定性的 L1/L2/L3 总结、分类 JSON、query 解析 JSON
- DeterministicEmbedding:稳定哈希生成固定维度向量(无语义,仅保证流程可跑)
本机无需 Docker / Ollama / RedisCI 可直接运行。
"""
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 ≥1chunk 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 payloaddoc_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):
"""正常 queryhits 非空,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):
"""空库 searchhits 为空、不报错、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
+150
View File
@@ -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/embedpayload 与 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)
+94
View File
@@ -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=1query2: 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([]) == {}
+104
View File
@@ -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 - 验证"
+262
View File
@@ -0,0 +1,262 @@
"""异步入库集成验证
在真实内存 Qdrant + FakeOllama 环境下验证异步入库全链路闭环:
POST 202 拿 task_id → wait_done 终态 → 任务查询 → 文档详情 → 检索命中 → 删除;
并覆盖两条边界路径:
- 失败路径:分类阶段抛错 → failed + error.stage/partial_summaryRedis 失败镜像 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 调用的假 Redisget_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 / Retrieverlifespan 建集合改空操作"""
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
+160
View File
@@ -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]:
"""TestClientlifespan 中的 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 /documentsHTTP 202data.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 == []
+200
View File
@@ -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:任务 failederror.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:
"""普通异常:任务 failederror.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_donefailed 用 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"
+273
View File
@@ -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
# L1category/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 写入失败:抛 IngestionErrorstage=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 == []
+99
View File
@@ -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]
+38
View File
@@ -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]:
"""TestClientlifespan 中的 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
+175
View File
@@ -0,0 +1,175 @@
"""知识统计 API 测试(GET /api/v1/knowledge/stats
两层覆盖:
1. TestClient + monkeypatch 注入假 Qdrant 服务:响应结构与异常路径
2. 真实内存 Qdrantlocation=":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]:
"""TestClientlifespan 中的 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 篇 L12 篇同类目 + 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 个 chunkdoc-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
+134
View File
@@ -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]
+126
View File
@@ -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.Responseraise_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
+236
View File
@@ -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"
+197
View File
@@ -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
+238
View File
@@ -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=Truerewrite 回退为原 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"
+79
View File
@@ -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")]
# l1a rank1b rank5l2b rank5a 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) == []
+247
View File
@@ -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,不调用 Ollamadecide_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"
+42
View File
@@ -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 == "请求参数校验失败"
+350
View File
@@ -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:
"""构造测试用 ScoredPointpayload 模拟 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="测试查询"))
# L1categories 过滤、无 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
# L2doc_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 路 d22.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
# chunkdoc_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="测试查询"))
# L1categories 与 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"}
+131
View File
@@ -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
+57
View File
@@ -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
+75
View File
@@ -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
Generated
+1044
View File
File diff suppressed because it is too large Load Diff