From fdb664e546a18825915c4056181b911f183f06b6 Mon Sep 17 00:00:00 2001 From: kplam Date: Thu, 30 Jul 2026 14:25:12 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E7=BB=9F=E4=B8=80=E4=BB=A3?= =?UTF-8?q?=E7=A0=81=E6=A0=BC=E5=BC=8F=EF=BC=8C=E8=B0=83=E6=95=B4=E5=A4=9A?= =?UTF-8?q?=E8=A1=8C=E4=BB=A3=E7=A0=81=E6=8D=A2=E8=A1=8C=E9=A3=8E=E6=A0=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 对多个文件进行代码格式化调整，将长行参数拆分为多行书写，提升代码可读性，包括： - 调整函数定义、调用的多行换行格式 - 优化列表、元组、字典的多行排版 - 新增README.md项目说明文档 --- README.md | 309 ++++++++++++++++++++++++++++++ app/api/v1/document.py | 57 ++++-- app/config.py | 8 +- app/core/file_parser.py | 8 +- app/core/ingest_tasks.py | 56 ++++-- app/main.py | 4 +- app/models/document.py | 9 +- tests/test_document_upload_api.py | 20 +- tests/test_file_parser.py | 25 ++- tests/test_ingest_tasks.py | 44 ++++- 10 files changed, 491 insertions(+), 49 deletions(-) create mode 100644 README.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..071581f --- /dev/null +++ b/README.md @@ -0,0 +1,309 @@ +# QMDSearch + +面向 AI Agent 的分层信息检索服务，支持多层级知识库检索、向量语义搜索和结构化数据查询。基于 Docker 部署在 NAS 上，为 AI Agent 提供高效、精准的信息检索能力。 + +## 特性 + +- **分层检索架构**：L1 文档定位 → L2 章节大纲 → L3 小节定位 → chunk 原文，逐层收窄，精准召回 +- **Hybrid 检索**：dense 向量 + sparse 稀疏向量 RRF 融合，兼顾语义与关键词匹配 +- **文档三级总结**：Ollama 本地模型自动生成文档 L1/L2/L3 总结，支持知识分类 +- **多格式文件入库**：支持 `.txt` `.md` `.html` `.htm` `.pdf` `.docx`，PDF 扫描件自动 OCR 降级 +- **文本去重**：Sha256 文本去重，重复文档自动复用已有结果 +- **JWT 认证**：Token 鉴权，支持注册开关与默认管理员 +- **管理后台**：`/admin` 单页面，含概览、文档管理、检索测试台、类目管理 +- **Docker 部署**：开箱即用的 Docker Compose 编排，适合 NAS 环境 + +## 技术栈 + +| 组件 | 技术 | +|------|------| +| 语言 | Python 3.12+ | +| Web 框架 | FastAPI + Uvicorn | +| 向量数据库 | Qdrant（Docker） | +| 缓存 | Redis（Docker） | +| 嵌入模型 | OpenAI / Ollama 本地模型（可切换） | +| 本地推理 | Ollama + qwen2.5:1.5b（文档总结） | +| OCR | rapidocr-onnxruntime + pypdfium2（扫描件 PDF） | +| 部署 | Docker Compose on NAS | + +## 项目结构 + +``` +QMDSearch/ +├── app/ # 应用主目录 +│ ├── main.py # FastAPI 入口（含 /admin 管理页面） +│ ├── config.py # 配置管理（pydantic-settings / 环境变量） +│ ├── api/ # API 路由层 +│ │ ├── response.py # 统一响应格式 +│ │ └── v1/ # API v1 +│ │ ├── search.py # 检索接口 +│ │ ├── document.py # 文档入库/管理/上传 +│ │ ├── knowledge.py # 知识库接口 +│ │ └── auth.py # 认证接口 +│ ├── core/ # 核心业务逻辑 +│ │ ├── retriever.py # 分层检索引擎 +│ │ ├── query_parser.py # Query 解析与分类路由 +│ │ ├── embeddings.py # 向量嵌入 +│ │ ├── sparse.py # 稀疏向量编码 +│ │ ├── ranker.py # RRF 融合与结果截断 +│ │ ├── summarizer.py # 文档三级总结 +│ │ ├── headings.py # 原生标题树解析 +│ │ ├── classifier.py # 文档分类 +│ │ ├── chunker.py # 标题树感知 chunk 切分 +│ │ ├── ingestion.py # 文档入库流水线 +│ │ ├── ingest_tasks.py # 异步任务管理（含文本去重） +│ │ ├── file_parser.py # 多格式文件解析（含 PDF OCR） +│ │ └── auth.py # JWT 认证 +│ ├── models/ # Pydantic 数据模型 +│ ├── services/ # 外部服务客户端 +│ │ ├── qdrant.py # Qdrant 客户端 +│ │ ├── redis.py # Redis 缓存客户端 +│ │ └── ollama.py # Ollama 客户端 +│ └── static/ # 静态资源 +│ └── admin.html # 管理后台（单文件） +├── tests/ # 测试 +├── scripts/ # 运维与评测脚本 +│ ├── eval/ # 回归评测 +│ └── smoke_live.py # 在线冒烟 +├── docker-compose.yml # Docker 编排 +├── Dockerfile # 应用镜像 +├── .env.example # 环境变量模板 +└── pyproject.toml # 项目配置 +``` + +## 分层检索架构 + +``` +Query → 解析路由 → L1（文档总结，候选文档集） + ↓ + L2（章节大纲，按 section 剪枝） + ↓ + L3（小节定位，两路 RRF 融合） + ↓ + Chunk（原文 hybrid 检索，dense + sparse） + ↓ + Top-K 结果 +``` + +### 检索链路详解 + +1. **Query 解析路由**：Ollama 小模型将 query 解析为结构化结果（rewrite、关键词、命中类目+置信度）。高置信且类目数不超上限时按类目过滤，否则全库兜底。 +2. **L1 - 文档定位层**：在文档 L1 总结集合中检索，产出候选文档集合。无命中时直接全库 chunk 兜底。 +3. **L2 - 章节大纲层**：在候选文档内检索 L2 章节大纲，按 section 剪枝定位。未命中的文档落入 L3 b 路。 +4. **L3 - 小节内容层**：两路查询（a：L2 命中文档按 section_path 过滤；b：其余文档仅按 doc_id 过滤）后 RRF 融合。 +5. **Chunk 层**：在 L3 收窄的范围内对原文 chunk 做 hybrid 检索（dense 向量 + sparse 稀疏向量 RRF 融合），返回最终 Top-K。 + +## 快速开始 + +### 前置要求 + +- Docker & Docker Compose +- OpenAI API Key（或本地 Ollama 模型） + +### 1. 克隆项目 + +```bash +git clone +cd QMDSearch +``` + +### 2. 配置环境变量 + +```bash +cp .env.example .env +# 编辑 .env，至少配置 OPENAI_API_KEY +``` + +### 3. 启动服务 + +```bash +docker compose up -d +``` + +### 4. 拉取 Ollama 模型（首次启动后） + +```bash +docker exec qmdsearch-ollama ollama pull qwen2.5:1.5b +``` + +### 5. 验证服务 + +```bash +curl http://localhost:8000/api/v1/health +# {"status": "ok"} +``` + +管理后台：浏览器访问 `http://localhost:8000/admin` + +## API 文档 + +| 方法 | 路径 | 说明 | 认证 | +|------|------|------|------| +| `GET` | `/api/v1/health` | 健康检查 | 否 | +| `POST` | `/api/v1/auth/register` | 用户注册 | 否（可关闭） | +| `POST` | `/api/v1/auth/login` | 用户登录，返回 JWT token | 否 | +| `GET` | `/api/v1/auth/me` | 获取当前用户信息 | 是 | +| `POST` | `/api/v1/search` | 分层检索 | 是 | +| `POST` | `/api/v1/documents` | 文档入库（JSON 文本，202 异步入库） | 是 | +| `POST` | `/api/v1/documents/upload` | 文件上传入库（multipart，202 异步） | 是 | +| `GET` | `/api/v1/documents/tasks/{task_id}` | 入库任务状态查询 | 是 | +| `GET` | `/api/v1/documents` | 文档列表（分页） | 是 | +| `GET` | `/api/v1/documents/{doc_id}` | 文档详情 | 是 | +| `DELETE` | `/api/v1/documents/{doc_id}` | 删除文档（幂等） | 是 | +| `GET` | `/api/v1/knowledge/categories` | 知识分类类目集 | 是 | +| `GET` | `/api/v1/knowledge/stats` | 统计（四层点数 + 类目分布） | 是 | +| `GET` | `/admin` | 管理后台 | 否 | + +### 统一响应格式 + +```json +{ + "code": 0, + "data": { ... }, + "message": "ok" +} +``` + +错误码：`0` 成功，`1xxx` 客户端错误，`2xxx` 服务端错误。 + +### 检索示例 + +```bash +curl -X POST http://localhost:8000/api/v1/search \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{"query": "如何配置 Redis 缓存", "top_k": 5}' +``` + +### 文件上传示例 + +```bash +curl -X POST http://localhost:8000/api/v1/documents/upload \ + -H "Authorization: Bearer " \ + -F "file=@document.pdf" +``` + +## 文档入库流程 + +``` +文件上传 → 文本提取 → 三级总结（Ollama）→ 分类判定（L1 总结）→ 向量化 → 写入 Qdrant + ↓ + 不足三级 → 2.5 级回退 +``` + +### 三级总结 + +| 层级 | 名称 | 说明 | +|------|------|------| +| L1 | 总结 | 对整篇文档的一句话高度概括 | +| L2 | 大纲 | 文档主要章节和关键主题 | +| L3 | 内容大纲 | 每个章节的详细内容摘要 | + +**2.5 级回退**：短文本（< 500 字符）自动降级为二级总结，跳过大纲层。 + +### 文件格式支持 + +| 格式 | 解析方式 | +|------|---------| +| `.txt` `.md` | UTF-8 解码 | +| `.html` `.htm` | HTML 标签剥离 | +| `.pdf` | pypdf 提取文本层；扫描件自动 OCR 降级 | +| `.docx` | python-docx 段落提取 | + +### PDF OCR 降级 + +图片型/扫描件 PDF 在文本层为空时，自动触发 OCR： +1. pypdfium2 渲染每页为图片 +2. rapidocr-onnxruntime 识别文字 +3. 拼接返回结果 + +可通过环境变量 `PDF_OCR_ENABLED=false` 关闭，或通过 `PDF_OCR_MAX_PAGES` / `PDF_OCR_DPI` 控制行为。 + +### 文本去重 + +入库时自动计算 `sha256(doc.text)`，相同文本重复提交不重跑流水线，直接复用已有文档 ID，结果中 `deduplicated=true` 标记。 + +## 配置项 + +所有配置通过环境变量注入，完整列表见 `.env.example`。主要配置项： + +| 变量 | 说明 | 默认值 | +|------|------|--------| +| `EMBEDDING_PROVIDER` | 嵌入模型提供商（openai / local） | `openai` | +| `OPENAI_API_KEY` | OpenAI API Key | - | +| `EMBEDDING_MODEL` | 嵌入模型名称 | `text-embedding-3-small` | +| `OLLAMA_MODEL` | Ollama 本地模型 | `qwen2.5:1.5b` | +| `RETRIEVAL_TOP_K` | 检索召回数 | `20` | +| `RETRIEVAL_FINAL_K` | 最终返回数 | `5` | +| `INGEST_MAX_CONCURRENCY` | 入库并发上限 | `2` | +| `JWT_SECRET_KEY` | JWT 签名密钥（生产必填） | 自动生成（仅开发） | +| `AUTH_REGISTER_ENABLED` | 是否开放注册 | `true` | +| `UPLOAD_MAX_SIZE_MB` | 上传文件大小上限 | `20` | +| `PDF_OCR_ENABLED` | 是否启用 PDF OCR | `true` | +| `PDF_OCR_MAX_PAGES` | OCR 最大页数 | `30` | +| `PDF_OCR_DPI` | OCR 渲染 DPI | `200` | + +## 开发 + +### 本地开发 + +```bash +# 安装依赖 +uv sync --extra dev + +# 启动开发服务器 +uv run uvicorn app.main:app --reload --host 0.0.0.0 --port 8000 + +# 运行测试 +uv run pytest + +# 代码检查 +uv run ruff check app/ tests/ +``` + +### 构建镜像 + +```bash +docker compose build +``` + +### 目录结构约定 + +- 类型注解必须（Python 3.12+ 语法） +- 异步优先（async/await） +- 配置通过环境变量注入（pydantic-settings） +- 日志使用 structlog 结构化日志 +- 测试使用 pytest + pytest-asyncio + +## 部署 + +### NAS 部署 + +```bash +# 1. 创建数据目录 +mkdir -p /path/to/nas/data/{qdrant,redis,ollama} + +# 2. 配置 .env，设置 NAS_DATA_DIR 和 API Key +NAS_DATA_DIR=/path/to/nas/data +OPENAI_API_KEY=sk-xxx + +# 3. 启动 +docker compose up -d + +# 4. 查看日志 +docker compose logs -f app +``` + +### 端口映射 + +| 服务 | 端口 | +|------|------| +| API | 8000 | +| Qdrant | 6333 | +| Qdrant Dashboard | 6334 | +| Redis | 6379 | +| Ollama | 11434 | + +## License + +MIT \ No newline at end of file diff --git a/app/api/v1/document.py b/app/api/v1/document.py index 01c927e..0ddf3e4 100644 --- a/app/api/v1/document.py +++ b/app/api/v1/document.py @@ -57,30 +57,44 @@ def _get_task_manager() -> IngestTaskManager: except Exception: logger.warning("Redis 缓存不可用，入库任务状态仅保留在内存", exc_info=True) redis = None - _task_manager = IngestTaskManager(ingester=_get_ingester(), redis=redis, settings=Settings()) + _task_manager = IngestTaskManager( + ingester=_get_ingester(), redis=redis, settings=Settings() + ) return _task_manager def _allowed_extensions() -> set[str]: """解析 settings.upload_allowed_extensions 逗号分隔字符串为扩展名集合（全小写、含点号）""" - return {ext.strip().lower() for ext in settings.upload_allowed_extensions.split(",") if ext.strip()} + return { + ext.strip().lower() + for ext in settings.upload_allowed_extensions.split(",") + if ext.strip() + } @router.post("/documents") -async def ingest_document(doc: DocumentInput, user: AuthUser = Depends(get_current_user)) -> JSONResponse: +async def ingest_document( + doc: DocumentInput, user: AuthUser = Depends(get_current_user) +) -> 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"})) + return JSONResponse( + status_code=202, content=ok({"task_id": task_id, "status": "pending"}) + ) @router.post("/documents/upload") async def upload_document( - file: UploadFile = File(..., description="上传的文件（.txt/.md/.html/.htm/.pdf/.docx）"), + file: UploadFile = File( + ..., description="上传的文件（.txt/.md/.html/.htm/.pdf/.docx）" + ), title: str = Form(default="", description="可选标题，默认取原文件名去扩展"), source: str = Form(default="", description="可选来源标识，默认 file:{原文件名}"), - metadata: str = Form(default="", description='可选元数据 JSON 字符串，如 \'{"author":"x"}\''), + metadata: str = Form( + default="", description='可选元数据 JSON 字符串，如 \'{"author":"x"}\'' + ), user: AuthUser = Depends(get_current_user), ) -> JSONResponse: """文件上传入库入口：校验 → 提取文本 → 落盘 → 提交异步入库流水线 @@ -129,9 +143,16 @@ async def upload_document( "original_size_bytes": str(len(content)), } ) - logger.info("上传文件已落盘", doc_id=doc_id, saved_path=saved_path, size=len(content)) + logger.info( + "上传文件已落盘", doc_id=doc_id, saved_path=saved_path, size=len(content) + ) except Exception: - logger.warning("上传文件落盘失败，仅做文本入库", doc_id=doc_id, filename=original_filename, exc_info=True) + logger.warning( + "上传文件落盘失败，仅做文本入库", + doc_id=doc_id, + filename=original_filename, + exc_info=True, + ) # 5. 合并用户传入的 metadata（落盘元数据优先级更高，不与用户键冲突） if metadata.strip(): @@ -150,7 +171,9 @@ async def upload_document( source = f"file:{original_filename}" # 7. 提交入库流水线 - doc_input = DocumentInput(text=text, title=title, source=source, metadata=metadata_dict) + doc_input = DocumentInput( + text=text, title=title, source=source, metadata=metadata_dict + ) task_id = await _get_task_manager().submit(doc_input) return JSONResponse( @@ -160,7 +183,9 @@ async def upload_document( @router.get("/documents/tasks/{task_id}") -async def get_ingest_task(task_id: str, user: AuthUser = Depends(get_current_user)) -> dict[str, Any]: +async def get_ingest_task( + task_id: str, user: AuthUser = Depends(get_current_user) +) -> dict[str, Any]: """查询入库任务状态：含 task_id/status/created_at/updated_at，done 附 result，failed 附 error""" task = await _get_task_manager().get(task_id) if task is None: @@ -184,7 +209,9 @@ async def list_documents( @router.get("/documents/{doc_id}") -async def get_document(doc_id: str, user: AuthUser = Depends(get_current_user)) -> dict[str, Any]: +async def get_document( + doc_id: str, user: AuthUser = Depends(get_current_user) +) -> dict[str, Any]: """获取文档详情：L1 记录 + L2/L3 节点 + chunks 数量""" try: detail = await _get_qdrant().get_doc_detail(doc_id) @@ -197,11 +224,15 @@ async def get_document(doc_id: str, user: AuthUser = Depends(get_current_user)) @router.delete("/documents/{doc_id}") -async def delete_document(doc_id: str, user: AuthUser = Depends(require_admin)) -> dict[str, Any]: +async def delete_document( + doc_id: str, user: AuthUser = Depends(require_admin) +) -> 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())}) + return ok( + {"doc_id": doc_id, "deleted": deleted, "deleted_total": sum(deleted.values())} + ) diff --git a/app/config.py b/app/config.py index bbc765e..b8ad1a3 100644 --- a/app/config.py +++ b/app/config.py @@ -65,10 +65,14 @@ class Settings(BaseSettings): # 文件上传 upload_dir: str = "./uploads" # 原始文件保存目录（相对路径以工作目录为基） upload_max_size_mb: int = 20 # 单文件大小上限（MB） - upload_allowed_extensions: str = ".txt,.md,.html,.htm,.pdf,.docx" # 允许上传的扩展名（逗号分隔） + upload_allowed_extensions: str = ( + ".txt,.md,.html,.htm,.pdf,.docx" # 允许上传的扩展名（逗号分隔） + ) # PDF OCR（图片型/扫描件降级，pypdf extract_text 为空时触发） - pdf_ocr_enabled: bool = True # 是否启用 OCR 降级（关闭则扫描件按"无法提取文本"拒绝入库） + pdf_ocr_enabled: bool = ( + True # 是否启用 OCR 降级（关闭则扫描件按"无法提取文本"拒绝入库） + ) pdf_ocr_max_pages: int = 30 # 单文件 OCR 页数上限，超过仅前 N 页 pdf_ocr_dpi: int = 200 # 渲染 DPI（越高越准但越慢，72~300 合理） diff --git a/app/core/file_parser.py b/app/core/file_parser.py index 796c7e4..1d71ef0 100644 --- a/app/core/file_parser.py +++ b/app/core/file_parser.py @@ -102,7 +102,9 @@ def _parse_pdf(content: bytes) -> str: if not settings.pdf_ocr_enabled: return "" - ocr_text = _ocr_pdf(content, max_pages=settings.pdf_ocr_max_pages, dpi=settings.pdf_ocr_dpi) + ocr_text = _ocr_pdf( + content, max_pages=settings.pdf_ocr_max_pages, dpi=settings.pdf_ocr_dpi + ) return ocr_text @@ -143,7 +145,9 @@ def _ocr_pdf(content: bytes, max_pages: int, dpi: int) -> str: result, _ = _ocr_engine(pil_image) if result: # result: [[box, text, score], ...]，按行拼接 - lines = [item[1] for item in result if item and len(item) >= 2 and item[1]] + lines = [ + item[1] for item in result if item and len(item) >= 2 and item[1] + ] if lines: page_texts.append("\n".join(lines)) pdf.close() diff --git a/app/core/ingest_tasks.py b/app/core/ingest_tasks.py index efad77b..a0c9c57 100644 --- a/app/core/ingest_tasks.py +++ b/app/core/ingest_tasks.py @@ -43,7 +43,9 @@ class IngestTaskStatus(StrEnum): # 终态集合 -TERMINAL_STATUSES: frozenset[str] = frozenset({IngestTaskStatus.DONE, IngestTaskStatus.FAILED}) +TERMINAL_STATUSES: frozenset[str] = frozenset( + {IngestTaskStatus.DONE, IngestTaskStatus.FAILED} +) def _utc_now_iso() -> str: @@ -54,7 +56,9 @@ def _utc_now_iso() -> str: class IngestTaskManager: """入库异步任务管理器：登记、后台执行、状态查询与 Redis 持久镜像""" - def __init__(self, ingester: Ingester, redis: RedisCache | None, settings: Settings) -> None: + def __init__( + self, ingester: Ingester, redis: RedisCache | None, settings: Settings + ) -> None: self._ingester = ingester self._redis = redis self._settings = settings @@ -123,7 +127,11 @@ class IngestTaskManager: 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) + 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]: @@ -147,7 +155,9 @@ class IngestTaskManager: """后台执行入库：并发限流 + 阶段状态推进 + 结果/错误落账 + 去重记录写入""" async with self._semaphore: try: - result = await self._ingester.ingest(doc, progress_cb=lambda stage: self._on_progress(task_id, stage)) + result = await self._ingester.ingest( + doc, progress_cb=lambda stage: self._on_progress(task_id, stage) + ) except IngestionError as exc: # 入库已知失败：透传阶段与已产出的部分总结 self._finish_failed( @@ -155,11 +165,18 @@ class IngestTaskManager: { "stage": exc.stage, "message": str(exc), - "partial_summary": exc.summary.model_dump(mode="json") if exc.summary is not None else None, + "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}) + self._finish_failed( + task_id, + {"stage": "unknown", "message": str(exc), "partial_summary": None}, + ) else: result_dict = result.model_dump(mode="json") self._tasks[task_id].update( @@ -178,7 +195,9 @@ class IngestTaskManager: try: return await self._redis.get_json(f"{DEDUP_KEY_PREFIX}{text_hash}") except Exception: - logger.warning("去重记录查询失败，降级为未命中", text_hash=text_hash, exc_info=True) + logger.warning( + "去重记录查询失败，降级为未命中", text_hash=text_hash, exc_info=True + ) return None async def _record_dedup(self, text_hash: str, result_dict: dict[str, Any]) -> None: @@ -196,9 +215,16 @@ class IngestTaskManager: 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._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"]) + 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 使用）""" @@ -213,7 +239,9 @@ class IngestTaskManager: """将当前任务状态快照异步镜像到 Redis（同步上下文也可调用）""" if self._redis is None: return - mirror = asyncio.create_task(self._mirror_to_redis(task_id, dict(self._tasks[task_id]))) + 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) @@ -225,6 +253,10 @@ class IngestTaskManager: else self._settings.ingest_task_ttl_done ) try: - await self._redis.set_json(f"{REDIS_KEY_PREFIX}{task_id}", snapshot, ttl=ttl) + 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) + logger.warning( + "入库任务状态镜像 Redis 失败", task_id=task_id, exc_info=True + ) diff --git a/app/main.py b/app/main.py index 28e045d..998a835 100644 --- a/app/main.py +++ b/app/main.py @@ -62,7 +62,9 @@ async def api_error_handler(request: Request, exc: ApiError) -> JSONResponse: @app.exception_handler(RequestValidationError) -async def validation_error_handler(request: Request, exc: RequestValidationError) -> JSONResponse: +async def validation_error_handler( + request: Request, exc: RequestValidationError +) -> JSONResponse: """请求参数校验失败 → code 1001""" errors = exc.errors() message = f"请求参数校验失败: {errors[0]['msg']}" if errors else "请求参数校验失败" diff --git a/app/models/document.py b/app/models/document.py index f117fd0..4933e6d 100644 --- a/app/models/document.py +++ b/app/models/document.py @@ -16,7 +16,9 @@ class DocumentSummary(BaseModel): """文档三级总结结果""" l1_summary: str = Field(description="L1 总结：一句话高度概括") - l2_outline: str | None = Field(default=None, description="L2 大纲：主要章节和关键主题") + l2_outline: str | None = Field( + default=None, description="L2 大纲：主要章节和关键主题" + ) l3_content_outline: str = Field(description="L3/L2.5 内容大纲：详细内容摘要") level: SummaryLevel = Field(description="实际使用的总结层级") @@ -49,4 +51,7 @@ class IngestionResult(BaseModel): chunks_count: int = Field(default=0, description="写入的 chunk 数量") tags: list[str] = Field(default_factory=list, description="附加分类标签") category_confidence: float = Field(default=0.0, description="主类目分类置信度") - deduplicated: bool = Field(default=False, description="是否命中去重复用旧文档（True 时 document_id 为既有文档 ID）") + deduplicated: bool = Field( + default=False, + description="是否命中去重复用旧文档（True 时 document_id 为既有文档 ID）", + ) diff --git a/tests/test_document_upload_api.py b/tests/test_document_upload_api.py index 87989f6..054253d 100644 --- a/tests/test_document_upload_api.py +++ b/tests/test_document_upload_api.py @@ -34,7 +34,9 @@ def client(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[TestClie return None monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections) - monkeypatch.setattr(document_module.settings, "upload_dir", str(tmp_path / "uploads")) + monkeypatch.setattr( + document_module.settings, "upload_dir", str(tmp_path / "uploads") + ) with TestClient(app) as test_client: yield test_client @@ -45,7 +47,13 @@ def _inject_manager(monkeypatch: pytest.MonkeyPatch, manager: FakeManager) -> No def _make_minimal_pdf(text: str = "Hello PDF World") -> bytes: content_stream = f"BT /F1 24 Tf 100 700 Td ({text}) Tj ET".encode("latin-1") - content_obj = b"<< /Length " + str(len(content_stream)).encode() + b" >>\nstream\n" + content_stream + b"\nendstream" + content_obj = ( + b"<< /Length " + + str(len(content_stream)).encode() + + b" >>\nstream\n" + + content_stream + + b"\nendstream" + ) return ( b"%PDF-1.0\n" b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n" @@ -161,7 +169,13 @@ def test_upload_docx_extracts_text( content = _make_docx(["第一段落", "第二段落"]) resp = client.post( "/api/v1/documents/upload", - files={"file": ("doc.docx", content, "application/vnd.openxmlformats-officedocument.wordprocessingml.document")}, + files={ + "file": ( + "doc.docx", + content, + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ) + }, ) assert resp.status_code == 202 diff --git a/tests/test_file_parser.py b/tests/test_file_parser.py index 6235757..8f7cb4a 100644 --- a/tests/test_file_parser.py +++ b/tests/test_file_parser.py @@ -14,7 +14,13 @@ from app.core.file_parser import parse_file, supported_extensions def _make_minimal_pdf(text: str = "Hello PDF World") -> bytes: """构造一个含一页文本的最小 PDF（pypdf 可读出文本）""" content_stream = f"BT /F1 24 Tf 100 700 Td ({text}) Tj ET".encode("latin-1") - content_obj = b"<< /Length " + str(len(content_stream)).encode() + b" >>\nstream\n" + content_stream + b"\nendstream" + content_obj = ( + b"<< /Length " + + str(len(content_stream)).encode() + + b" >>\nstream\n" + + content_stream + + b"\nendstream" + ) return ( b"%PDF-1.0\n" b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n" @@ -36,6 +42,7 @@ def _make_minimal_pdf(text: str = "Hello PDF World") -> bytes: def _make_docx(text_lines: list[str]) -> bytes: from docx import Document # type: ignore[import-untyped] + document = Document() for line in text_lines: document.add_paragraph(line) @@ -222,7 +229,9 @@ def _patch_pdf_ocr_deps(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", fake_module) -def test_parse_pdf_ocr_fallback_when_text_layer_empty(monkeypatch: pytest.MonkeyPatch) -> None: +def test_parse_pdf_ocr_fallback_when_text_layer_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: """扫描件 PDF（文本层全空）触发 OCR 降级，返回识别文本""" _patch_pdf_ocr_deps(monkeypatch) monkeypatch.setattr(settings, "pdf_ocr_enabled", True) @@ -256,7 +265,9 @@ def test_parse_pdf_ocr_respects_max_pages(monkeypatch: pytest.MonkeyPatch) -> No assert "OCR文本第2页" not in result -def test_parse_pdf_ocr_returns_empty_when_dependency_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: +def test_parse_pdf_ocr_returns_empty_when_dependency_unavailable( + monkeypatch: pytest.MonkeyPatch, +) -> None: """rapidocr 导入失败：降级返回空文本，且把 _ocr_unavailable 置 True 避免重试""" monkeypatch.setattr("pypdf.PdfReader", _FakePdfReader) # 故意让 rapidocr_onnxruntime 提供一个非类的 RapidOCR，构造时抛错 @@ -273,7 +284,9 @@ def test_parse_pdf_ocr_returns_empty_when_dependency_unavailable(monkeypatch: py assert fp_module._ocr_unavailable is True -def test_parse_pdf_ocr_runtime_exception_falls_back_to_empty(monkeypatch: pytest.MonkeyPatch) -> None: +def test_parse_pdf_ocr_runtime_exception_falls_back_to_empty( + monkeypatch: pytest.MonkeyPatch, +) -> None: """OCR 运行时抛错：仅告警，降级返回空文本（不抛出 ValueError）""" monkeypatch.setattr("pypdf.PdfReader", _FakePdfReader) @@ -291,7 +304,9 @@ def test_parse_pdf_ocr_runtime_exception_falls_back_to_empty(monkeypatch: pytest assert parse_file("scan.pdf", b"fake pdf bytes") == "" -def test_parse_pdf_text_layer_present_skips_ocr(monkeypatch: pytest.MonkeyPatch) -> None: +def test_parse_pdf_text_layer_present_skips_ocr( + monkeypatch: pytest.MonkeyPatch, +) -> None: """文本层非空：直接返回文本，OCR 引擎不会被实例化""" call_count = 0 diff --git a/tests/test_ingest_tasks.py b/tests/test_ingest_tasks.py index ddbdf9e..667cef0 100644 --- a/tests/test_ingest_tasks.py +++ b/tests/test_ingest_tasks.py @@ -7,15 +7,28 @@ from datetime import datetime from typing import Any from app.config import Settings -from app.core.ingest_tasks import DEDUP_KEY_PREFIX, REDIS_KEY_PREFIX, IngestTaskManager, IngestTaskStatus +from app.core.ingest_tasks import ( + DEDUP_KEY_PREFIX, + REDIS_KEY_PREFIX, + IngestTaskManager, + IngestTaskStatus, +) from app.core.ingestion import IngestionError -from app.models.document import DocumentInput, DocumentSummary, IngestionResult, SummaryLevel +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 + l1_summary="一句话总结", + l2_outline=None, + l3_content_outline="内容大纲", + level=SummaryLevel.L3, ) @@ -42,7 +55,9 @@ class FakeIngester: self.gate: asyncio.Event | None = None self.started = asyncio.Event() - async def ingest(self, doc: DocumentInput, progress_cb: Callable[[str], None] | None = None) -> IngestionResult: + 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: @@ -64,7 +79,9 @@ class FakeRedis: 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: + 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") @@ -105,8 +122,12 @@ async def test_submit_returns_immediately_and_completes() -> None: async def test_ingestion_error_marks_failed_with_stage_and_partial_summary() -> None: """IngestionError：任务 failed，error.stage 透传，partial_summary 保留""" - summary = DocumentSummary(l1_summary="L1", l2_outline=None, l3_content_outline="L3", level=SummaryLevel.L3) - ingester = FakeIngester(error=IngestionError("classify", "分类判定失败: boom", summary=summary)) + 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="正文")) @@ -179,10 +200,14 @@ async def test_redis_mirror_ttl_done_and_failed() -> None: assert done_writes[-1][0]["status"] == IngestTaskStatus.DONE redis2 = FakeRedis() - failing_manager = IngestTaskManager(FakeIngester(error=RuntimeError("boom")), redis2, settings) + 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)] + 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 @@ -275,6 +300,7 @@ async def test_dedup_skipped_when_redis_unavailable() -> None: async def test_dedup_lookup_failure_falls_back_to_normal_pipeline() -> None: """Redis get_json 抛错：去重查询降级为未命中，走原流水线""" + class _ExplodingRedis(FakeRedis): async def get_json(self, key: str) -> dict[str, Any] | None: raise RuntimeError("redis down")