Compare commits

..

20 Commits

Author SHA1 Message Date
Serendipity 2ddb6de92f feat: 第三优先级完善 — CI/CD、PyPI、pre-commit、CHANGELOG、贡献指南、基准测试、ADR
CI / Test (Python 3.13) (push) Has been cancelled
2026-07-11 20:13:36 +08:00
Serendipity 7e0b9553c8 chore: 配置 pre-commit 钩子(ruff + mypy + 通用检查) 2026-07-11 20:08:21 +08:00
Serendipity 4710a305be docs: README 补充 Web UI、.docx 支持说明 2026-07-11 20:04:11 +08:00
Serendipity 6b89488677 test: 补充测试覆盖率 — DB/CLI/Search/Ingest + 修复 chromadb 异常兼容 2026-07-11 20:02:24 +08:00
Serendipity 10c5d3d347 feat: 添加 Web 管理界面(Vue 3 SPA 单文件)+ /admin 挂载 2026-07-11 19:57:30 +08:00
Serendipity e71018a9a3 feat: 新增数据导出功能和 CLI export 命令(JSON/CSV) 2026-07-11 19:55:57 +08:00
Serendipity d9b03ab5f0 feat: 新增 DocxSplitter — 支持 .docx 文档入库 2026-07-11 19:55:03 +08:00
Serendipity 04f2f8a8a8 refactor: 修复 Embedder Protocol 标准写法和 PDF/EPUB Splitter 参数语义 2026-07-11 19:53:08 +08:00
Serendipity 2d0e8c4997 chore: ruff lint 修复 + 第一优先级全部完成 2026-07-11 19:50:35 +08:00
Serendipity 3e51e166c2 docs: 更新 README — 混合检索、增量入库、Docker、导出功能 2026-07-11 19:48:15 +08:00
Serendipity 1df4793acb feat: 添加 Dockerfile 和 docker-compose.yml 部署方案 2026-07-11 19:47:39 +08:00
Serendipity 9784f5f436 feat: API 安全加固 — 请求体大小限制、审计日志、健康检查免限速 2026-07-11 19:47:14 +08:00
Serendipity ba814c7a70 feat: CLI 添加 --incremental/--force/--mode 选项 2026-07-11 19:45:50 +08:00
Serendipity 82924ff5f3 feat: 添加 Cross-Encoder Reranker + 集成到 Searcher 2026-07-11 19:44:56 +08:00
Serendipity 1d58a55a73 feat: Searcher 集成 HybridRetriever + SearchConfig 配置 2026-07-11 19:43:49 +08:00
Serendipity cb470e516b feat: 添加 HybridRetriever — BM25+向量混合检索 2026-07-11 19:42:20 +08:00
Serendipity 7b3c7d5323 feat: DocumentIngestor 支持增量入库(FileTracker) 2026-07-11 19:41:21 +08:00
Serendipity 3e2c61402a feat: 添加 FileTracker — 基于 SHA256 的文件变更追踪 2026-07-11 19:39:12 +08:00
Serendipity 78ad21d311 chore: 添加 rank-bm25 依赖(混合检索用) 2026-07-11 19:38:33 +08:00
Serendipity f7801b45a7 fix: CLI 路径检查改用 is_safe_cli_path(允许绝对路径),更新文档 2026-07-10 15:45:22 +08:00
63 changed files with 6938 additions and 120 deletions
BIN
View File
Binary file not shown.
+16
View File
@@ -0,0 +1,16 @@
__pycache__/
*.pyc
.venv/
.pytest_cache/
.ruff_cache/
.mypy_cache/
.vscode/
.git/
.gitignore
.env
data/
*.egg-info/
dist/
build/
.coverage
coverage.xml
+48
View File
@@ -0,0 +1,48 @@
# CI
name: CI
on:
push:
branches: [main]
pull_request:
branches: [main]
jobs:
test:
name: Test (Python ${{ matrix.python-version }})
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
python-version: ["3.13"]
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: uv sync --extra dev --extra all
- name: Run lint (ruff)
run: uv run ruff check src/ tests/
- name: Run type check (mypy)
run: uv run mypy src/ --ignore-missing-imports
- name: Run tests with coverage
run: |
uv run pytest tests/ \
--cov=src \
--cov-report=term-missing \
--cov-report=xml \
-v
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:
files: ./coverage.xml
fail_ci_if_error: false
+28
View File
@@ -0,0 +1,28 @@
# Publish to PyPI
name: Publish to PyPI
on:
push:
tags:
- "v*"
jobs:
publish:
name: Build and publish
runs-on: ubuntu-latest
permissions:
id-token: write
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
- name: Build package
run: uv build
- name: Publish to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: dist/
+31
View File
@@ -0,0 +1,31 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.11.0
hooks:
- id: ruff
args: [--fix]
- id: ruff-format
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v5.0.0
hooks:
- id: trailing-whitespace
- id: end-of-file-fixer
- id: check-yaml
args: [--unsafe]
- id: check-json
- id: check-toml
- id: check-added-large-files
args: ['--maxkb=500']
- id: detect-private-key
- id: debug-statements
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v1.15.0
hooks:
- id: mypy
args: [--ignore-missing-imports, --no-strict-optional]
files: ^src/
additional_dependencies:
- types-pyyaml
- types-requests
+41
View File
@@ -0,0 +1,41 @@
# Changelog
所有值得注意的更改都将记录在此文件中。
格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)
版本号遵循 [语义化版本](https://semver.org/lang/zh-CN/)。
## [Unreleased]
### Added
- 混合检索(BM25 + 向量联合),HybridRetriever + SearchConfig 可配置
- 增量入库(SHA256 文件变更追踪),CLI `--incremental` / `--force` 选项
- Cross-Encoder 结果重排序(BAAI/bge-reranker-base
- Web 管理界面(Vue 3 SPA`/admin`
- .docx 文档支持(markitdown
- 数据导出功能(JSON/CSV),CLI `export` 命令
- Docker 部署方案 + docker-compose(含 GPU profile
- API 审计日志、请求体大小限制、健康检查免限速
- GitHub Actions CI 流水线 + PyPI 发布
- Pre-commit 钩子配置(ruff + mypy
- 性能基准测试(嵌入 + 检索)
### Changed
- Searcher 支持 hybrid/vector 检索模式切换
- Embedder Protocol 修复为标准写法
- PDF/EPUB Splitter 参数名统一为 `source`
### Fixed
- CORS `allow_credentials` 配置修复
- chromadb 不同版本异常类型兼容(NotFoundError
## [0.1.0] - 2026-07-05
### Added
- Markdown 文档解析与语义检索
- 多 Provider 嵌入支持(local/OpenAI/DashScope
- 多格式文档(.md/.txt/.pdf/.html/.epub
- GPU 自动检测加速
- FastAPI HTTP API + Typer CLI
- API Key 认证和速率限制
- 路径遍历安全防护
+10 -6
View File
@@ -4,14 +4,14 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
## 项目概述
Markdown 文档向量数据库 — 将 .md 文件分块 → 嵌入 → 存入 ChromaDB,通过 FastAPI HTTP 或 CLI 提供语义检索。
文档向量数据库 — 将 .md/.txt/.pdf/.html/.epub 文件分块 → 嵌入 → 存入 ChromaDB,通过 FastAPI HTTP 或 CLI 提供语义检索。
## 常用命令
```bash
```Shell
# --- 安装与测试 ---
uv sync # 安装依赖(lockfile 已锁定 CUDA torch
uv run pytest tests/ -v # 全部测试 (115+ 个)
uv run pytest tests/ -v # 全部测试 (120 个)
uv run pytest tests/test_api.py -v # 单个测试模块
uv run pytest tests/ -v -k "test_search" # 按名称过滤
@@ -36,6 +36,7 @@ uv run python scripts/ingest_obsidian.py
│ ├── db.py # VectorDB: 线程安全的 ChromaDB 封装
│ ├── embedder.py # 策略模式: LocalEmbedder / OpenAIEmbedder / DashscopeEmbedder
│ ├── ingest.py # DocumentIngestor: 按扩展名自动选择 Splitter
│ ├── security.py # 路径遍历防护 (is_safe_path + is_path_within_workspace)
│ ├── search.py # Searcher: 语义检索 + 源文件管理
│ └── splitters/ # 文档分块器包
│ ├── base.py # Splitter(Protocol) + BaseTextSplitter(ABC)
@@ -43,6 +44,7 @@ uv run python scripts/ingest_obsidian.py
│ ├── text.py # TextSplitter: 纯文本段落切分
│ ├── pdf.py # PDFSplitter: pymupdf 提取文字
│ ├── html.py # HTMLSplitter: bs4 去标签
│ ├── epub.py # EPUBSplitter: ebooklib 提取章节
│ └── registry.py # 扩展名 → Splitter 自动选择
server/ # FastAPI HTTP 层
@@ -55,7 +57,7 @@ cli/main.py # Typer CLI5 个命令 + --config 选项
**数据流**: 文件 → `get_splitter(path)` 自动选择 → `Splitter.split()``batch_embed()``ChromaDB collection.add()``Searcher.search()`
**支持的格式**: `.md` / `.txt` / `.pdf` / `.html` — 安装可选依赖: `uv sync --extra all`
**支持的格式**: `.md` / `.txt` / `.pdf` / `.html` / `.epub` — 安装可选依赖: `uv sync --extra all`
**依赖方向**: `config``db``embedder``ingest`/`search``server`/`cli`
@@ -90,7 +92,7 @@ cli/main.py # Typer CLI5 个命令 + --config 选项
- **API Key**: 环境变量 `MD_VECTOR_API_KEY``verify_api_key` 依赖注入到 ingest/search/delete 端点;未设置则跳过
- **速率限制**: `RateLimiter` 中间件,默认 60s 窗口内最多 30 请求
- **路径遍历防护**: `_is_safe_path()` 拒绝绝对路径`..` 穿越
- **路径遍历防护**: `is_path_within_workspace()` 拒绝绝对路径`..` 穿越和目录外访问
- **错误信息**: 500 返回通用消息,详细错误记入 `logger.exception`
### 配置系统 (config.py)
@@ -119,7 +121,7 @@ cli/main.py # Typer CLI5 个命令 + --config 选项
## 已入库知识库
| 集合名 | 来源 | 文件数 | chunks | 说明 |
|--------|------|--------|--------|------|
| ----------------- | -------------------------------------------------------------------------------- | ------ | ------ | ------------------------------ |
| `novel_taohou` | `D:\Code\doing_exercises\exercise\Novel\我有太后罩着,你们有什么\原有章节剧情` | 220 | 670 | 小说章节(GPU bge-small-v1.5 |
| `obsidian_blog` | `D:\Code\Obsidian` | 51 | 3,611 | 博客笔记(GPU bge-small-v1.5 |
| `default` | 测试文件 | 2 | ~30 | test-guide.md + stdin-doc.md |
@@ -127,6 +129,7 @@ cli/main.py # Typer CLI5 个命令 + --config 选项
搜索时务必用 `-C` 指定集合,否则只会搜到 default 中的测试数据。
**小说搜索示例**
```bash
uv run md-vector-db search "张莽和孙太后的关系" -k 3 -C novel_taohou
uv run md-vector-db search "抄家事件" -k 5 -C novel_taohou
@@ -149,6 +152,7 @@ index-strategy = "unsafe-best-match" # 允许跨源查找
**不要删除 `[tool.uv]` 配置**,否则 `uv sync` 会重新解析为 CPU 版 torch。
### MarkdownSplitter 边界情况
`_split_single_paragraph` 中,当段落分隔符(。!?等)距 chunk 起点 < overlap(100) 时,
`start` 会回退为负数,Python `str.rfind` 的负索引会绕回文本末尾,造成死循环。
此 bug 已被修复(`start = max(start + 1, next_start)`),但给超长段落测试时需留意类似问题。
+54
View File
@@ -0,0 +1,54 @@
# 贡献指南
## 开发环境
```bash
git clone git@github.com:LHY0125/md-vector-db.git
cd md-vector-db
uv sync --extra dev --extra all
uv run pre-commit install
```
## 开发流程
1. **Fork 仓库** → 创建功能分支 `feat/xxx`
2. **写测试**TDD: 先在 `tests/` 下写失败测试
3. **实现功能**: 最少代码让测试通过
4. **运行全部测试**: `uv run pytest tests/ --cov=src -v`
5. **确保覆盖率 ≥ 80%**: `uv run pytest tests/ --cov=src --cov-fail-under=80`
6. **Lint 检查**: `uv run ruff check src/ tests/`
7. **类型检查**: `uv run mypy src/ --ignore-missing-imports`
8. **提交**: 遵循约定式提交格式
9. **创建 PR**: 描述清楚变更内容和测试结果
## 提交消息格式
```
<类型>: <描述>
<可选正文>
```
类型: `feat`, `fix`, `refactor`, `docs`, `test`, `chore`, `perf`, `ci`
## 代码风格
- Python 3.13+,遵循 PEP 8
- 类型注解覆盖所有 public API
- 中文注释
- 小文件原则(<800 行),函数 <50 行
## 添加新文档格式支持
1.`src/core/splitters/` 下实现 Splitter Protocol
2.`src/core/splitters/registry.py` 注册扩展名
3.`pyproject.toml` 添加可选依赖
4.`tests/` 下添加测试
5. 更新 `README.md`
## 运行基准测试
```bash
uv sync --extra bench
uv run pytest tests/benchmarks/ -v --benchmark-only
```
+38
View File
@@ -0,0 +1,38 @@
# Dockerfile — md-vector-db 生产镜像
FROM python:3.13-slim-bookworm
LABEL org.opencontainers.image.title="md-vector-db"
LABEL org.opencontainers.image.description="Markdown 文档向量数据库"
# 系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
# 先复制依赖文件以利用 Docker 层缓存
COPY pyproject.toml uv.lock ./
# 安装 uv 并同步依赖(CPU 模式)
RUN pip install --no-cache-dir uv \
&& uv sync --frozen --no-dev \
&& uv cache clean
# 复制源码和配置
COPY config.yaml .env.example ./
COPY src/ ./src/
COPY scripts/ ./scripts/
# 创建数据目录
RUN mkdir -p /app/data
# 暴露端口
EXPOSE 8000
# 健康检查
HEALTHCHECK --interval=30s --timeout=10s --start-period=60s --retries=3 \
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/api/v1/health')" || exit 1
# 默认启动 HTTP 服务
CMD ["uv", "run", "md-vector-db", "serve", "--port", "8000"]
+62 -9
View File
@@ -1,18 +1,28 @@
# md-vector-db
[![CI](https://lhy-git.liuhangyv.top/Serendipity/md-vector-db/actions/workflows/ci.yml/badge.svg)](https://github.com/LHY0125/md-vector-db/actions/workflows/ci.yml)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
Markdown 文档向量数据库 — 将 Markdown 文件自动分块、嵌入、存入 ChromaDB,通过语义检索快速查找相关内容。提供 **CLI 命令行工具**和 **HTTP API** 两种使用方式供其他项目集成。
## 功能特性
- **文档入库**: 支持单文件、目录批量导入 Markdown 文档,自动按标题+段落智能分块
- **文档入库**: 支持单文件、目录批量导入多种格式文档,自动按标题+段落智能分块
- **语义检索**: 自然语言查询,返回最相关的文档片段及来源定位(文件名、章节标题)
- **多 Provider**: 本地模型 + 云端 APIOpenAI、阿里云 DashScope、硅基流动等 OpenAI 兼容服务)
- **GPU 加速**: 本地模型自动检测 CUDARTX 4060 实测:729 chunks 嵌入仅 1.8s
- **多集合**: 支持多项目数据隔离,不同知识库存入不同 ChromaDB collection
- **HTTP API**: FastAPI 提供 RESTful 接口,附带 Swagger 文档
- **安全**: 可选 API Key 认证、速率限制、路径遍历防护
- **多格式文档**: 支持 `.md` / `.txt` / `.pdf` / `.html`,按扩展名自动选择分块器,可通过 `Splitter` Protocol 扩展
- **多格式文档**: 支持 `.md` / `.txt` / `.pdf` / `.html` / `.epub`,按扩展名自动选择分块器,可通过 `Splitter` Protocol 扩展
- **去重**: 同一文件重复入库自动覆盖旧版本(基于路径 SHA256 哈希)
- **混合检索**: BM25 关键词 + 向量语义联合检索,加权融合排序,支持切换纯向量模式
- **增量入库**: 基于 SHA256 哈希自动跳过未变更文件,避免重复嵌入浪费 GPU
- **结果重排序**: 可选 Cross-Encoder 精确重排(`BAAI/bge-reranker-base`),提升检索精度
- **数据导出**: 支持 JSON/CSV 导出 collection 全量数据
- **Web 管理界面**: 内置 Vue 3 管理面板(`/admin`),可视化搜索、入库、查看集合
- **Docker 部署**: 提供 Dockerfile 和 docker-compose.yml,一键部署(含 GPU profile
- **.docx 支持**: 通过 markitdown 库支持 Word 文档入库
## 快速开始
@@ -52,7 +62,7 @@ embed:
# 单文件
uv run md-vector-db ingest docs/intro.md
# 批量导入目录(递归扫描所有 .md 文件
# 批量导入目录(递归扫描所有支持的文档格式
uv run md-vector-db ingest-dir ./md_docs/
# 指定集合(多项目数据隔离)
@@ -145,16 +155,42 @@ for r in resp.json()["results"]:
---
---
## Docker 部署
```bash
# 构建并启动(CPU 模式)
docker compose up -d
# GPU 模式(需 nvidia-container-toolkit
docker compose --profile gpu up -d
# 查看日志
docker compose logs -f
# 停止服务
docker compose down
# CLI 使用示例
docker exec -it md-vector-db uv run md-vector-db stats
docker exec -it md-vector-db uv run md-vector-db ingest /app/md_docs/doc.md
```
---
## CLI 命令参考
所有命令均支持 `--config/-c`(配置文件)、`--collection/-C`(集合名,默认 `default`)。
| 命令 | 说明 |
| --------------------------- | -------------------------------------------------------- |
| `ingest <文件路径>` | 入库单个 .md 文件,支持`-C` 指定集合 |
| `ingest-dir <目录路径>` | 递归入库目录下所有 .md 文件 |
| `search <查询> -k <数量>` | 语义检索,`-k` 默认 10、最大 100`--json` JSON 输出 |
| ----------------------------------- | -------------------------------------------------------- |
| `ingest <文件路径> --incremental` | 增量入库单文件,自动跳过未变更文件 |
| `ingest <文件路径> --force` | 强制重新入库(忽略增量检查) |
| `ingest-dir <目录路径>` | 递归入库目录下所有支持的文档格式 |
| `search <查询> -k <数量> --mode` | 语义检索,`--mode hybrid\|vector``--json` JSON 输出 |
| `stats` | 显示 chunks 总数、源文件列表 |
| `export -o <文件> -f <json\|csv>` | 导出 collection 数据 |
| `serve -p <端口>` | 启动 HTTP 服务(默认 8000 |
---
@@ -182,6 +218,12 @@ chunk:
server:
host: 0.0.0.0 # 服务监听地址
port: 8000 # 服务监听端口
search:
mode: hybrid # 检索模式: hybrid (BM25+向量) | vector (纯向量)
bm25_weight: 0.3 # BM25 权重 (0=纯向量, 1=纯BM25)
candidate_multiplier: 3 # 向量检索候选倍数
enable_rerank: false # 是否启用 Cross-Encoder 重排序
```
### .env 环境变量
@@ -224,6 +266,7 @@ md-vector-db/
│ │ ├── db.py # ChromaDB 封装(线程安全)
│ │ ├── embedder.py # 嵌入器(Local/OpenAI/Dashscope
│ │ ├── ingest.py # 混合分块 + 入库
│ │ ├── security.py # 路径遍历防护
│ │ ├── search.py # 语义检索
│ │ └── splitters/ # 文档分块器包
│ │ ├── base.py # Splitter(Protocol) + BaseTextSplitter(ABC)
@@ -244,13 +287,13 @@ md-vector-db/
├── scripts/
│ ├── serve.py # 快速启动脚本
│ └── ingest_obsidian.py # 批量入库 Obsidian 知识库
└── tests/ # 测试(115+ 个)
└── tests/ # 测试(120 个)
```
## 测试
```bash
uv run pytest tests/ -v # 全部测试 (115+ 个)
uv run pytest tests/ -v # 全部测试 (120 个)
uv run pytest tests/test_embedder.py -v # 嵌入器测试
uv run pytest tests/ -v -k "search" # 按名称过滤
uv run pytest tests/ -v --cov=src --cov-report=term-missing # 覆盖率
@@ -287,6 +330,16 @@ index-strategy = "unsafe-best-match" # 允许跨源查找
---
---
## 文档
- [API 参考](docs/api_reference.md)
- [架构决策记录](docs/architecture.md)
-
- [贡献指南](CONTRIBUTING.md)
- [变更日志](CHANGELOG.md)
## License
MIT
+6
View File
@@ -21,3 +21,9 @@ server:
port: 8000
# ssl_keyfile: "" # HTTPS 私钥路径(设置后启用 HTTPS)
# ssl_certfile: "" # HTTPS 证书路径(设置后启用 HTTPS)
search:
mode: hybrid # hybrid | vector
bm25_weight: 0.3 # BM25 权重 (0=纯向量, 1=纯BM25)
candidate_multiplier: 3 # 向量检索候选倍数
enable_rerank: false # Cross-Encoder 重排序
+66
View File
@@ -0,0 +1,66 @@
# docker-compose.yml — md-vector-db 本地开发与生产部署
version: "3.8"
services:
md-vector-db:
build:
context: .
dockerfile: Dockerfile
image: md-vector-db:latest
container_name: md-vector-db
restart: unless-stopped
ports:
- "${MD_VECTOR_PORT:-8000}:8000"
volumes:
# 持久化 ChromaDB 数据
- ./data:/app/data
# 挂载配置文件
- ./config.yaml:/app/config.yaml:ro
# 挂载待入库文档目录
- ${MD_VECTOR_DOCS_DIR:-./md_docs}:/app/md_docs:ro
environment:
- MD_VECTOR_CONFIG=/app/config.yaml
- MD_VECTOR_DB_DATA_DIR=/app/data
- MD_VECTOR_API_KEY=${MD_VECTOR_API_KEY:-}
- EMBED_API_KEY=${EMBED_API_KEY:-}
- CORS_ORIGINS=${CORS_ORIGINS:-http://localhost:3000}
- MAX_REQUEST_BODY_SIZE=${MAX_REQUEST_BODY_SIZE:-10485760}
env_file:
- .env
healthcheck:
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/api/v1/health')"]
interval: 30s
timeout: 10s
retries: 3
start_period: 60s
# 可选:GPU 版本(需 nvidia-container-toolkit
md-vector-db-gpu:
profiles: ["gpu"]
build:
context: .
dockerfile: Dockerfile
image: md-vector-db:latest
container_name: md-vector-db-gpu
restart: unless-stopped
ports:
- "${MD_VECTOR_PORT:-8000}:8000"
volumes:
- ./data:/app/data
- ./config.yaml:/app/config.yaml:ro
- ${MD_VECTOR_DOCS_DIR:-./md_docs}:/app/md_docs:ro
environment:
- MD_VECTOR_CONFIG=/app/config.yaml
- MD_VECTOR_DB_DATA_DIR=/app/data
- MD_VECTOR_API_KEY=${MD_VECTOR_API_KEY:-}
- EMBED_API_KEY=${EMBED_API_KEY:-}
- CORS_ORIGINS=${CORS_ORIGINS:-http://localhost:3000}
env_file:
- .env
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
+68
View File
@@ -0,0 +1,68 @@
# md-vector-db API 参考文档
## REST API
- **Base URL**: `http://localhost:8000/api/v1`
- **认证**: `X-API-Key` Header(取决于 `MD_VECTOR_API_KEY` 环境变量)
- **速率限制**: 60s 窗口内最多 30 请求
### GET /health — 健康检查
无需认证。
```json
{
"status": "ok",
"checks": {
"chromadb": {"status": "ok", "count": 1240},
"embedder": {"status": "ok", "dimension": 512}
}
}
```
### GET /collections — 列出集合
需 API Key。返回 `{"collections": [{"name": "...", "count": N}, ...]}`
### POST /search — 语义检索
| 字段 | 类型 | 必填 | 约束 |
|------|------|------|------|
| query | string | ✅ | 1-2000 字符 |
| top_k | int | ❌ | 1-100,默认 10 |
| collection | string | ❌ | ≤128 字符 |
返回 `{"results": [...], "collection": "..."}`,每条结果含 `id`, `content`, `source_file`, `section_title`, `heading_level`, `score`
### POST /ingest — 入库文档
| 字段 | 类型 | 必填 | 约束 |
|------|------|------|------|
| file_path | string | 二选一 | 安全路径 |
| content | string | 二选一 | ≤500KB |
| file_name | string | 推荐 | 1-255 字符 |
| collection | string | ❌ | ≤128 字符 |
返回 `{"status": "ok", "chunks": N, "file": "...", "collection": "..."}`
### DELETE /documents/{file_name} — 删除文档
查询参数: `collection` (可选)。
| 状态码 | 含义 |
|--------|------|
| 200 | 删除成功 |
| 400 | file_name 不合法 |
| 401 | API Key 无效 |
| 404 | 文档不存在 |
## CLI 命令
```bash
md-vector-db ingest <文件> --incremental --force
md-vector-db ingest-dir <目录> -C <集合>
md-vector-db search "<查询>" -k 10 --mode hybrid --json
md-vector-db stats -C <集合>
md-vector-db export -o output.json -f json
md-vector-db serve -p 8000
```
+43
View File
@@ -0,0 +1,43 @@
# 架构决策记录 (ADR)
## ADR-1: 为什么选择 ChromaDB
**日期**: 2026-07-05 | **状态**: 已采纳
**背景**: 需要一个嵌入式向量数据库来持久化文档嵌入。
**候选**: ChromaDB(嵌入式 SQLite)、Qdrant(独立服务)、FAISS(纯内存)、Milvus(生产级集群)
**决策**: ChromaDB。零运维成本远超吞吐量考量。内置 Collection 概念映射多知识库场景。
**代价**: 高并发下逊于 Qdrant/Milvus。未来可透明迁移(Embedder/Searcher 已隔离 ChromaDB 依赖)。
---
## ADR-2: 为什么默认 bge-small-zh-v1.5
**日期**: 2026-07-05 | **状态**: 已采纳
**候选**: bge-small-zh-v1.5 (512维/23M)、bge-large-zh-v1.5 (1024维/324M)、text2vec-large-chinese、m3e-base
**决策**: bge-small-zh-v1.5。RTX 4060 上 729 chunks 嵌入仅 1.8s,日常精度足够。高精度场景可切换 large 模型或 OpenAI API。
---
## ADR-3: 为什么采用 Protocol 而非 ABC
**日期**: 2026-07-05 | **状态**: 已采纳
**候选**: typing.Protocol(结构化子类型)、abc.ABC(名义子类型)、Callable(丢失类型信息)
**决策**: Protocol。外部模块无需依赖本项目源码即可实现 Splitter/Embedder,对插件化友好。
---
## ADR-4: 为什么使用 rank-bm25 而非集成搜索引擎?
**日期**: 2026-07-11 | **状态**: 已采纳
**候选**: rank-bm25(纯 Python BM25)、Elasticsearch(外部服务)、Whoosh(纯 Python 全文搜索)
**决策**: rank-bm25。零运维、轻量、与现有 ChromaDB 架构匹配。在向量候选上做 BM25 重打分(而非全文索引所有文档),兼顾性能和精度。
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+31 -2
View File
@@ -1,7 +1,21 @@
[project]
name = "md-vector-db"
version = "0.1.0"
description = "Markdown 文档向量数据库,支持语义检索"
description = "Markdown 文档向量数据库语义检索、混合检索、REST API"
readme = "README.md"
license = {text = "MIT"}
authors = [
{name = "刘航宇", email = "3364451258@qq.com"},
]
keywords = ["vector-database", "semantic-search", "rag", "chromadb", "markdown"]
classifiers = [
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3.13",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Text Processing :: Markup",
]
requires-python = ">=3.13"
dependencies = [
"chromadb>=0.5.0",
@@ -12,8 +26,15 @@ dependencies = [
"markdown-it-py>=3.0.0",
"typer>=0.12.0",
"python-dotenv>=1.2.2",
"rank-bm25>=0.2.2",
]
[project.urls]
Homepage = "https://github.com/LHY0125/md-vector-db"
Documentation = "https://github.com/LHY0125/md-vector-db#readme"
Repository = "https://github.com/LHY0125/md-vector-db"
Issues = "https://github.com/LHY0125/md-vector-db/issues"
[project.scripts]
md-vector-db = "src.cli.main:app"
@@ -22,7 +43,9 @@ dev = ["pytest>=8.0", "httpx>=0.27.0", "pytest-cov>=5.0", "ruff>=0.8.0", "mypy>=
pdf = ["pymupdf>=1.24.0"]
html = ["beautifulsoup4>=4.12.0"]
epub = ["ebooklib>=0.18"]
all = ["md-vector-db[pdf,html,epub]", "requests>=2.31.0", "openai>=1.0.0"]
docx = ["markitdown>=0.1.0"]
bench = ["pytest-benchmark>=5.0"]
all = ["md-vector-db[pdf,html,epub,docx]", "requests>=2.31.0", "openai>=1.0.0"]
[build-system]
requires = ["hatchling"]
@@ -39,6 +62,11 @@ index-strategy = "unsafe-best-match"
line-length = 100
target-version = "py313"
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
skip-magic-trailing-comma = false
[tool.ruff.lint]
select = ["E", "F", "I", "N", "W"]
@@ -49,3 +77,4 @@ ignore_missing_imports = true
[tool.pytest.ini_options]
testpaths = ["tests"]
pythonpath = ["src"]
norecursedirs = ["tests/benchmarks"]
+47
View File
@@ -0,0 +1,47 @@
"""版本号管理 — 更新 pyproject.toml 中的 version."""
import re
import sys
from pathlib import Path
def bump(part: str) -> str:
"""递增版本号.
Args:
part: "major" | "minor" | "patch"
Returns:
新版本号字符串
"""
pyproject = Path(__file__).parent.parent / "pyproject.toml"
content = pyproject.read_text(encoding="utf-8")
match = re.search(r'version\s*=\s*"(\d+)\.(\d+)\.(\d+)"', content)
if not match:
print("错误: 未找到 version 字段", file=sys.stderr)
sys.exit(1)
major, minor, patch = int(match[1]), int(match[2]), int(match[3])
if part == "major":
major += 1
minor = 0
patch = 0
elif part == "minor":
minor += 1
patch = 0
elif part == "patch":
patch += 1
else:
print(f"错误: 未知的版本部分 '{part}',可选: major/minor/patch", file=sys.stderr)
sys.exit(1)
new_version = f"{major}.{minor}.{patch}"
new_content = content.replace(match[0], f'version = "{new_version}"')
pyproject.write_text(new_content, encoding="utf-8")
print(f"版本: {match[1]}.{match[2]}.{match[3]}{new_version}")
return new_version
if __name__ == "__main__":
if len(sys.argv) != 2:
print("用法: python scripts/bump_version.py <major|minor|patch>")
sys.exit(1)
bump(sys.argv[1])
+44 -10
View File
@@ -1,10 +1,10 @@
"""命令行工具入口 — 可作为 MCP tool 直接调用."""
from pathlib import Path
import sys
import glob as _glob
import io
import json
import os
import glob as _glob
import sys
from pathlib import Path
from typing import Annotated
import typer
@@ -17,8 +17,8 @@ if sys.stdout.encoding != "utf-8":
sys.path.insert(0, str(Path(__file__).parent.parent))
from src.core.config import DEFAULT_CONFIG_PATH
from src.core.security import is_path_within_workspace
from src.server.deps import get_state, get_default_collection
from src.core.security import is_safe_cli_path
from src.server.deps import get_default_collection, get_state
app = typer.Typer(
name="md-vector-db",
@@ -63,6 +63,12 @@ def ingest(
typer.Argument(help="Markdown 文件路径 (可多个, 或 - 从标准输入读取)"),
] = None,
name: Annotated[str | None, typer.Option("--name", help="标准输入模式下的虚拟文件名")] = None,
incremental: Annotated[
bool, typer.Option("--incremental", help="增量模式:跳过未变更文件")
] = False,
force: Annotated[
bool, typer.Option("--force", help="强制重新入库(忽略增量检查)")
] = False,
config: ConfigOpt = DEFAULT_CONFIG_PATH,
collection: CollectionOpt = None,
):
@@ -82,7 +88,7 @@ def ingest(
if file_paths:
total = 0
for fp in file_paths:
if not is_path_within_workspace(fp):
if not is_safe_cli_path(fp):
typer.echo(f"[SKIP] 不安全的路径: {fp}", err=True)
continue
# 支持通配符 (shell 展开或 Python glob)
@@ -90,14 +96,14 @@ def ingest(
if "*" in fp or "?" in fp:
matches = _glob.glob(fp, recursive=True)
for m in matches:
if not is_path_within_workspace(m):
if not is_safe_cli_path(m):
typer.echo(f"[SKIP] 不安全的路径: {m}", err=True)
continue
c = ingestor.ingest_file(m)
c = ingestor.ingest_file(m, incremental=incremental, force=force)
typer.echo(f" {m}: {c} chunks")
total += c
elif p.is_file():
c = ingestor.ingest_file(fp)
c = ingestor.ingest_file(fp, incremental=incremental, force=force)
typer.echo(f" {fp}: {c} chunks")
total += c
else:
@@ -117,7 +123,7 @@ def ingest_dir(
collection: CollectionOpt = None,
):
_init_config(config)
if not is_path_within_workspace(dir_path):
if not is_safe_cli_path(dir_path):
typer.echo(f"错误: 不安全的路径 — {dir_path}", err=True)
raise typer.Exit(code=1)
state = get_state()
@@ -137,12 +143,15 @@ def search(
query: Annotated[str, typer.Argument(help="搜索关键词或自然语言查询")],
top_k: Annotated[int, typer.Option("--top-k", "-k", help="返回结果数量 (1-100)")] = 10,
json_output: Annotated[bool, typer.Option("--json", help="以 JSON 格式输出")] = False,
mode: Annotated[str, typer.Option("--mode", help="检索模式: hybrid | vector")] = "",
config: ConfigOpt = DEFAULT_CONFIG_PATH,
collection: CollectionOpt = None,
):
_init_config(config)
state = get_state()
searcher = state.get_searcher(_resolve_collection(collection))
if mode:
searcher._search_config.mode = mode
results = searcher.search(query, top_k=top_k)
if json_output:
@@ -213,5 +222,30 @@ def stats(
typer.echo(f" - {s}")
@app.command(help="导出 collection 数据为 JSON 或 CSV.")
def export(
output: Annotated[str, typer.Option("--output", "-o", help="输出文件路径")],
fmt: Annotated[
str, typer.Option("--format", "-f", help="导出格式: json | csv")
] = "json",
config: ConfigOpt = DEFAULT_CONFIG_PATH,
collection: CollectionOpt = None,
):
"""导出 collection 数据."""
_init_config(config)
state = get_state()
searcher = state.get_searcher(_resolve_collection(collection))
if fmt == "json":
searcher.export_json(file_path=output)
elif fmt == "csv":
searcher.export_csv(file_path=output)
else:
typer.echo(f"错误: 不支持的格式 '{fmt}',可选: json, csv", err=True)
raise typer.Exit(code=1)
typer.echo(f"[OK] 已导出到: {output}")
if __name__ == "__main__":
app()
+13 -2
View File
@@ -5,9 +5,8 @@ import os
from dataclasses import dataclass, field
from pathlib import Path
from dotenv import load_dotenv
import yaml
from dotenv import load_dotenv
# 加载 .env 文件 (若存在)
load_dotenv()
@@ -65,6 +64,16 @@ class ServerConfig:
ssl_certfile: str = "" # HTTPS 证书文件路径 (空则使用 HTTP)
@dataclass
class SearchConfig:
"""检索配置."""
mode: str = "hybrid" # hybrid | vector
bm25_weight: float = 0.3 # BM25 权重 (0=纯向量, 1=纯BM25)
candidate_multiplier: int = 3 # 向量检索候选倍数
enable_rerank: bool = False # 是否启用 Cross-Encoder 重排序
@dataclass
class AppConfig:
"""应用总配置."""
@@ -73,6 +82,7 @@ class AppConfig:
embed: EmbedConfig = field(default_factory=EmbedConfig)
chunk: ChunkConfig = field(default_factory=ChunkConfig)
server: ServerConfig = field(default_factory=ServerConfig)
search: SearchConfig = field(default_factory=SearchConfig)
@classmethod
def from_dict(cls, data: dict) -> "AppConfig":
@@ -82,6 +92,7 @@ class AppConfig:
embed=EmbedConfig(**data.get("embed", {})),
chunk=ChunkConfig(**data.get("chunk", {})),
server=ServerConfig(**data.get("server", {})),
search=SearchConfig(**data.get("search", {})),
)
+3
View File
@@ -45,6 +45,9 @@ class VectorDB:
self.client.delete_collection(name=name)
except ValueError:
pass # collection 不存在则忽略
except Exception:
# chromadb 不同版本可能抛出 NotFoundError 等
pass
def delete_by_source(self, collection_name: str, file_name: str) -> bool:
"""按 source_file 删除文档 (线程安全)."""
+16 -3
View File
@@ -17,9 +17,9 @@
vectors = batch_embed(embedder, long_text_list)
"""
import logging
import os
import threading
import logging
from typing import Protocol
from src.core.config import EmbedConfig
@@ -48,9 +48,22 @@ _PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
# -- 接口 --
class Embedder(Protocol):
"""嵌入器接口."""
@property
def dimension(self) -> int: ...
def embed(self, texts: list[str]) -> list[list[float]]: ...
def dimension(self) -> int:
"""返回嵌入向量的维度."""
...
def embed(self, texts: list[str]) -> list[list[float]]:
"""对文本列表进行嵌入.
Args:
texts: 待嵌入的文本列表
Returns:
嵌入向量列表,每个向量为 float 列表
"""
...
# -- 本地模型 --
+117
View File
@@ -0,0 +1,117 @@
"""文件变更追踪 — 基于 SHA256 + mtime 判断文件是否需要重新入库."""
from __future__ import annotations
import hashlib
import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from threading import Lock
from typing import TypedDict
logger = logging.getLogger("md-vector-db")
class FileRecord(TypedDict):
"""追踪记录."""
hash: str
mtime: float
size: int
ingested_at: str
class FileTracker:
"""文件入库追踪器 — JSON 文件持久化.
通过比对文件内容的 SHA256 哈希判断文件是否变更。
使用 threading.Lock 保护 JSON 文件的并发读写。
"""
def __init__(self, db_path: str = "./ingest_tracker.json"):
self._db_path = Path(db_path)
self._lock = Lock()
self._records: dict[str, FileRecord] = {}
self._load()
# -- 公开 API --
def is_stale(self, file_path: str) -> bool:
"""检查文件是否需要重新入库.
Returns:
True: 文件不存在 / 无记录 / 内容已变更
False: 文件未变更且记录存在
"""
path = Path(file_path)
if not path.exists():
return True
record = self.get_record(file_path)
if record is None:
return True
current_hash = self.compute_hash(file_path)
return current_hash != record["hash"]
def get_record(self, file_path: str) -> FileRecord | None:
"""获取文件的追踪记录(无记录返回 None)."""
return self._records.get(self._abs_key(file_path))
def mark_ingested(self, file_path: str) -> None:
"""标记文件已入库(创建或更新追踪记录)."""
path = Path(file_path)
if not path.exists():
logger.warning("标记已入库时文件不存在: %s", file_path)
return
stat = path.stat()
record: FileRecord = {
"hash": self.compute_hash(file_path),
"mtime": stat.st_mtime,
"size": stat.st_size,
"ingested_at": self._now_iso(),
}
with self._lock:
self._records[self._abs_key(file_path)] = record
self._save()
def remove_record(self, file_path: str) -> None:
"""移除文件的追踪记录."""
key = self._abs_key(file_path)
with self._lock:
if key in self._records:
del self._records[key]
self._save()
@staticmethod
def compute_hash(file_path: str) -> str:
"""计算文件 SHA256 哈希(分块读取,适合大文件)."""
sha = hashlib.sha256()
with open(file_path, "rb") as f:
while chunk := f.read(8192):
sha.update(chunk)
return sha.hexdigest()
# -- 内部 --
def _abs_key(self, file_path: str) -> str:
"""生成标准化 key — 使用绝对路径."""
return str(Path(file_path).resolve())
@staticmethod
def _now_iso() -> str:
return datetime.now(timezone.utc).isoformat()
def _load(self) -> None:
if self._db_path.exists():
try:
with open(self._db_path, "r", encoding="utf-8") as f:
self._records = json.load(f)
except (json.JSONDecodeError, OSError) as e:
logger.warning("tracker 文件损坏,重置为空: %s", e)
self._records = {}
def _save(self) -> None:
try:
with open(self._db_path, "w", encoding="utf-8") as f:
json.dump(self._records, f, ensure_ascii=False, indent=2)
except OSError as e:
logger.error("无法写入 tracker 文件: %s", e)
+27 -8
View File
@@ -1,13 +1,13 @@
"""Markdown 文档解析与入库模块."""
import logging
import re
from pathlib import Path
from src.core.config import ChunkConfig
from src.core.db import VectorDB
from src.core.embedder import Embedder, batch_embed
from src.core.splitters.markdown import MarkdownSplitter # 兼容旧 import 路径
from src.core.file_tracker import FileTracker
from src.core.splitters.base import Splitter # 兼容旧 import 路径
from src.core.splitters.markdown import MarkdownSplitter # 兼容旧 import 路径
from src.core.splitters.registry import SUPPORTED_SUFFIXES, get_splitter
logger = logging.getLogger("md-vector-db")
@@ -23,28 +23,41 @@ class DocumentIngestor:
collection_name: str,
splitter: Splitter | None = None,
chunk_config: ChunkConfig | None = None,
file_tracker: FileTracker | None = None,
):
self.db = db
self.embedder = embedder
self.collection_name = collection_name
self.splitter = splitter or MarkdownSplitter()
self.chunk_config = chunk_config or ChunkConfig()
self.file_tracker = file_tracker # None = 不使用增量功能
@property
def collection(self):
return self.db.get_or_create_collection(self.collection_name)
def ingest_file(self, file_path: str) -> int:
def ingest_file(self, file_path: str, incremental: bool = False, force: bool = False) -> int:
"""入库单个文件, 返回 chunk 数量.
根据文件扩展名自动选择 Splitter.md→MarkdownSplitter, .txt→TextSplitter, .pdf→PDFSplitter 等)。
根据扩展名自动选择 Splitter.md/.txt/.pdf/.html/.epub 等)。
使用文件路径的 SHA256 前 12 位 + 文件名作为唯一标识。
Args:
file_path: 文件路径
incremental: 启用增量模式(需 file_tracker 已注入)
force: 强制重新入库(忽略增量检查)
"""
import hashlib
path = Path(file_path).resolve()
path_hash = hashlib.sha256(str(path).encode()).hexdigest()[:12]
file_name = f"{path_hash}_{path.name}"
# 增量模式:检查是否需要重新入库
if incremental and not force and self.file_tracker is not None:
if not self.file_tracker.is_stale(str(path)):
logger.debug("跳过未变更文件: %s", path.name)
return 0
splitter = self.splitter or get_splitter(
file_path,
max_size=self.chunk_config.max_size,
@@ -54,11 +67,17 @@ class DocumentIngestor:
# PDF/EPUB 二进制文件特殊处理:splitter 内部读取文件
suffix = path.suffix.lower()
if suffix in (".pdf", ".epub"):
chunks = splitter.split(str(path), source_file=file_name)
return self._add_chunks(chunks, file_name)
chunks = splitter.split(source=str(path), source_file=file_name)
result = self._add_chunks(chunks, file_name)
else:
content = path.read_text(encoding="utf-8")
return self._ingest_with_splitter(content, file_name, splitter)
result = self._ingest_with_splitter(content, file_name, splitter)
# 入库成功后更新 tracker
if self.file_tracker is not None:
self.file_tracker.mark_ingested(str(path))
return result
def _ingest_with_splitter(self, content: str, file_name: str, splitter) -> int:
"""分块 + 嵌入 + 入库(文本文件通用路径)."""
+69
View File
@@ -0,0 +1,69 @@
"""Cross-Encoder 重排序器."""
from __future__ import annotations
import logging
logger = logging.getLogger("md-vector-db")
class Reranker:
"""使用 Cross-Encoder 模型对检索结果重排序.
默认模型: BAAI/bge-reranker-base(中文友好)
首次调用时懒加载模型。
"""
_DEFAULT_MODEL = "BAAI/bge-reranker-base"
def __init__(self, model_name: str | None = None):
self._model_name = model_name or self._DEFAULT_MODEL
self._model = None
def rerank(
self, query: str, candidates: list[dict], top_k: int = 10
) -> list[dict]:
"""对候选列表重排序.
Args:
query: 原始查询
candidates: 候选结果列表(需含 "content" 字段)
top_k: 返回数量
Returns:
按 cross-encoder 分数降序的结果列表
"""
if not candidates:
return []
self._ensure_model()
# 构建 (query, document) 对
pairs = [(query, c["content"]) for c in candidates]
try:
scores = self._model.predict(pairs, show_progress_bar=False)
except Exception as e:
logger.error("Cross-Encoder 重排序失败: %s", e)
# 降级:保留原始顺序
return candidates[:top_k]
# 附加 rerank_score
for i, c in enumerate(candidates):
c["rerank_score"] = round(float(scores[i]), 4)
# 按 rerank_score 降序排列
candidates.sort(key=lambda x: x.get("rerank_score", 0), reverse=True)
# 返回 top_k,将 rerank_score 作为最终 score
result = candidates[:top_k]
for r in result:
r["score"] = r.get("rerank_score", r.get("score", 0))
return result
def _ensure_model(self) -> None:
"""懒加载 Cross-Encoder 模型."""
if self._model is not None:
return
from sentence_transformers import CrossEncoder
logger.info("加载 Cross-Encoder 模型: %s", self._model_name)
self._model = CrossEncoder(self._model_name)
+187
View File
@@ -0,0 +1,187 @@
"""混合检索器 — BM25 关键词 + 向量语义联合检索."""
from __future__ import annotations
import logging
import re
from typing import TYPE_CHECKING, TypedDict
from rank_bm25 import BM25Okapi
from src.core.db import VectorDB
from src.core.embedder import Embedder
if TYPE_CHECKING:
from src.core.reranker import Reranker
logger = logging.getLogger("md-vector-db")
class SearchResult(TypedDict):
"""检索结果类型."""
id: str
content: str
source_file: str
section_title: str
heading_level: int
chunk_index: int
score: float
bm25_score: float
vector_score: float
class HybridRetriever:
"""BM25 + 向量混合检索器.
Architecture:
1. 向量检索取得 top_k * candidate_multiplier 候选
2. BM25 对候选打分
3. 加权融合排序(默认 0.7 向量 + 0.3 BM25
4. 返回 top_k 结果
支持按 source_file 元数据过滤。
"""
def __init__(
self,
db: VectorDB,
embedder: Embedder,
collection_name: str,
bm25_weight: float = 0.3,
vector_candidate_multiplier: int = 3,
reranker: "Reranker | None" = None,
):
self._db = db
self._embedder = embedder
self._collection_name = collection_name
self._bm25_weight = bm25_weight
self._vector_multiplier = vector_candidate_multiplier
self._reranker = reranker
self._bm25_index: BM25Okapi | None = None
self._bm25_docs: list[str] = []
self._bm25_ids: list[str] = []
@property
def _collection(self):
return self._db.get_or_create_collection(self._collection_name)
def search(
self,
query: str,
top_k: int = 10,
source_file: str | None = None,
) -> list[dict]:
"""混合检索.
Args:
query: 查询文本
top_k: 返回结果数量
source_file: 可选,按源文件过滤
Returns:
按混合分数降序排列的结果列表
"""
# 1. 向量检索:多取候选
vector_candidates = self._vector_search(
query, top_k * self._vector_multiplier, source_file
)
if not vector_candidates:
return []
# 2. BM25 打分
bm25_scored = self._bm25_rerank(query, vector_candidates)
# 3. 分数融合
fused = self._fuse_scores(bm25_scored, self._bm25_weight)
# 4. 排序取 top_k
fused.sort(key=lambda x: x["score"], reverse=True)
# 4.5 可选:Cross-Encoder 重排序
if self._reranker is not None and len(fused) > 1:
fused = self._reranker.rerank(query, fused, top_k=top_k)
return fused[:top_k]
# -- 内部方法 --
def _vector_search(
self, query: str, n: int, source_file: str | None
) -> list[dict]:
"""向量检索取得候选."""
embeddings = self._embedder.embed([query])
if not embeddings:
return []
query_embedding = embeddings[0]
where_filter = {"source_file": source_file} if source_file else None
results = self._collection.query(
query_embeddings=[query_embedding],
n_results=n,
where=where_filter,
include=["documents", "metadatas", "distances"],
)
candidates = []
if results["ids"] and results["ids"][0]:
for i, doc_id in enumerate(results["ids"][0]):
metadata = results["metadatas"][0][i] if results["metadatas"] else {}
distance = results["distances"][0][i] if results["distances"] else 0.0
vector_score = max(0.0, round(1.0 - distance, 4))
candidates.append({
"id": doc_id,
"content": results["documents"][0][i] if results["documents"] else "",
"source_file": metadata.get("source_file", ""),
"section_title": metadata.get("section_title", ""),
"heading_level": metadata.get("heading_level", 0),
"chunk_index": metadata.get("chunk_index", 0),
"vector_score": vector_score,
})
return candidates
def _bm25_rerank(self, query: str, candidates: list[dict]) -> list[dict]:
"""用 BM25 对候选列表重新打分."""
if not candidates:
return candidates
tokenized_query = self._tokenize(query)
tokenized_candidates = [self._tokenize(c["content"]) for c in candidates]
bm25 = BM25Okapi(tokenized_candidates)
scores = bm25.get_scores(tokenized_query)
# 归一化 BM25 分数到 [0, 1]
max_score = max(scores) if max(scores) > 0 else 1.0
for i, c in enumerate(candidates):
c["bm25_score"] = round(scores[i] / max_score, 4)
return candidates
def _fuse_scores(self, items: list[dict], bm25_weight: float) -> list[dict]:
"""加权融合向量分和 BM25 分."""
vector_weight = 1.0 - bm25_weight
for item in items:
bm25_s = item.get("bm25_score", 0.0)
vec_s = item.get("vector_score", 0.0)
item["score"] = round(vec_s * vector_weight + bm25_s * bm25_weight, 4)
return items
def _ensure_bm25_index(self) -> None:
"""确保 BM25 索引已构建(从 collection 所有文档构建)."""
if self._bm25_index is not None:
return
all_data = self._collection.get(include=["documents", "metadatas"])
if all_data and all_data["ids"]:
self._bm25_ids = all_data["ids"]
self._bm25_docs = all_data["documents"] or []
tokenized = [self._tokenize(d) for d in self._bm25_docs]
self._bm25_index = BM25Okapi(tokenized) if tokenized else None
else:
self._bm25_ids = []
self._bm25_docs = []
self._bm25_index = None
@staticmethod
def _tokenize(text: str) -> list[str]:
"""简易中文+英文分词(按中文单字 + 英文单词拆分).
注意: 这是基础实现。生产环境建议集成 jieba 分词。
"""
tokens = []
for match in re.finditer(r"[a-zA-Z0-9]+|[一-鿿]|[^\s]", text):
tokens.append(match.group().lower())
return tokens
+100 -3
View File
@@ -1,30 +1,67 @@
"""语义检索模块."""
import csv
import io
import json as json_lib
import logging
from src.core.config import SearchConfig
from src.core.db import VectorDB
from src.core.embedder import Embedder
from src.core.retriever import HybridRetriever
logger = logging.getLogger("md-vector-db")
class Searcher:
"""向量检索器."""
"""向量检索器 — 支持纯向量或 HybridRetriever 混合检索."""
def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str):
def __init__(
self,
db: VectorDB,
embedder: Embedder,
collection_name: str,
search_config: SearchConfig | None = None,
):
self.db = db
self.embedder = embedder
self.collection_name = collection_name
self._search_config = search_config or SearchConfig()
self._hybrid: HybridRetriever | None = None
@property
def collection(self):
return self.db.get_or_create_collection(self.collection_name)
def _get_hybrid(self) -> HybridRetriever:
if self._hybrid is None:
reranker = None
if self._search_config.enable_rerank:
from src.core.reranker import Reranker
reranker = Reranker()
self._hybrid = HybridRetriever(
self.db,
self.embedder,
self.collection_name,
bm25_weight=self._search_config.bm25_weight,
vector_candidate_multiplier=self._search_config.candidate_multiplier,
reranker=reranker,
)
return self._hybrid
def search(
self,
query: str,
top_k: int = 10,
source_file: str | None = None,
) -> list[dict]:
"""语义检索, 返回格式化结果列表."""
"""语义检索, 返回格式化结果列表.
根据 search.mode 配置自动选择 hybrid 或 vector 模式.
"""
if self._search_config.mode == "hybrid":
return self._get_hybrid().search(query, top_k=top_k, source_file=source_file)
# 纯向量模式(原有逻辑)
embeddings = self.embedder.embed([query])
if not embeddings:
raise RuntimeError("嵌入器返回空结果, 无法进行检索")
@@ -83,3 +120,63 @@ class Searcher:
def delete_by_source(self, file_name: str) -> bool:
"""按文件名删除文档 (委托 VectorDB)."""
return self.db.delete_by_source(self.collection_name, file_name)
def export_json(self, file_path: str | None = None) -> str:
"""导出 collection 所有 chunks 为 JSON.
Args:
file_path: 可选,写入文件路径。不传则返回 JSON 字符串。
Returns:
JSON 字符串
"""
all_data = self.collection.get(include=["documents", "metadatas"])
records = []
if all_data and all_data["ids"]:
for i, doc_id in enumerate(all_data["ids"]):
meta = all_data["metadatas"][i] if all_data["metadatas"] else {}
records.append({
"id": doc_id,
"content": all_data["documents"][i] if all_data["documents"] else "",
"source_file": meta.get("source_file", ""),
"section_title": meta.get("section_title", ""),
"heading_level": meta.get("heading_level", 0),
"chunk_index": meta.get("chunk_index", i),
})
json_str = json_lib.dumps(records, ensure_ascii=False, indent=2)
if file_path:
with open(file_path, "w", encoding="utf-8") as f:
f.write(json_str)
return json_str
def export_csv(self, file_path: str | None = None) -> str:
"""导出 collection 所有 chunks 为 CSV.
Args:
file_path: 可选,写入文件路径。不传则返回 CSV 字符串。
Returns:
CSV 字符串
"""
all_data = self.collection.get(include=["documents", "metadatas"])
output = io.StringIO()
writer = csv.writer(output)
writer.writerow([
"id", "content", "source_file", "section_title", "heading_level", "chunk_index",
])
if all_data and all_data["ids"]:
for i, doc_id in enumerate(all_data["ids"]):
meta = all_data["metadatas"][i] if all_data["metadatas"] else {}
writer.writerow([
doc_id,
all_data["documents"][i] if all_data["documents"] else "",
meta.get("source_file", ""),
meta.get("section_title", ""),
meta.get("heading_level", 0),
meta.get("chunk_index", i),
])
csv_str = output.getvalue()
if file_path:
with open(file_path, "w", encoding="utf-8", newline="") as f:
f.write(csv_str)
return csv_str
+12
View File
@@ -46,3 +46,15 @@ def is_path_within_workspace(path_str: str) -> bool:
except ValueError:
return False
return common == cwd
def is_safe_cli_path(path_str: str) -> bool:
"""CLI 路径安全检查 — 仅拒绝 .. 穿越组件,允许绝对路径和任意目录。
CLI 是本地工具,用户有权限访问系统中任意路径。
与 is_path_within_workspaceAPI 用,绑定当前目录)相比更宽松。
"""
parts = path_str.replace("\\", "/").split("/")
if ".." in parts:
return False
return True
+9 -7
View File
@@ -1,12 +1,13 @@
"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML / EPUB."""
"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML / EPUB / Docx."""
from src.core.splitters.base import Splitter, BaseTextSplitter
from src.core.splitters.markdown import MarkdownSplitter
from src.core.splitters.text import TextSplitter
from src.core.splitters.pdf import PDFSplitter
from src.core.splitters.html import HTMLSplitter
from src.core.splitters.base import BaseTextSplitter, Splitter
from src.core.splitters.docx import DocxSplitter
from src.core.splitters.epub import EPUBSplitter
from src.core.splitters.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES
from src.core.splitters.html import HTMLSplitter
from src.core.splitters.markdown import MarkdownSplitter
from src.core.splitters.pdf import PDFSplitter
from src.core.splitters.registry import SUPPORTED_SUFFIXES, get_splitter, register_splitter
from src.core.splitters.text import TextSplitter
__all__ = [
"Splitter",
@@ -16,6 +17,7 @@ __all__ = [
"PDFSplitter",
"HTMLSplitter",
"EPUBSplitter",
"DocxSplitter",
"get_splitter",
"register_splitter",
"SUPPORTED_SUFFIXES",
+58
View File
@@ -0,0 +1,58 @@
""".docx Word 文档分块器 — 使用 markitdown 提取文字后委托 TextSplitter."""
from __future__ import annotations
import logging
from pathlib import Path
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
class DocxSplitter:
"""Docx 分块器:markitdown 提取文字 → TextSplitter 分块.
实现 Splitter Protocol,内部组合 TextSplitter 实例。
split() 的 source 参数可接收文件路径或纯文本。
"""
def __init__(self, max_size: int = 1000, overlap: int = 100):
self._text_splitter = TextSplitter(max_size=max_size, overlap=overlap)
def split(self, source: str, source_file: str = "") -> list[dict]:
"""从 .docx 文件或纯文本提取文字并分块.
Args:
source: .docx 文件路径或纯文本内容
source_file: 来源文件名
Returns:
分块后的 chunk 列表
"""
path = Path(source)
if path.suffix.lower() == ".docx":
text = self._extract_from_docx(str(path))
else:
text = source
if not text.strip():
return []
return self._text_splitter.split(text, source_file=source_file)
def _extract_from_docx(self, file_path: str) -> str:
"""使用 markitdown 从 .docx 文件中提取文字."""
try:
from markitdown import MarkItDown
except ImportError:
raise ImportError(
"Docx 支持需要 markitdown 库. 请执行: uv sync --extra docx"
)
try:
md = MarkItDown()
result = md.convert(file_path)
return result.text_content
except Exception as e:
logger.error("Docx 解析失败: %s%s", file_path, e)
raise ValueError(f"Docx 解析失败: {e}") from e
+5 -4
View File
@@ -1,5 +1,6 @@
"""EPUB 电子书分块器 — 使用 ebooklib 提取文字后委托 TextSplitter."""
import logging
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
@@ -9,17 +10,17 @@ class EPUBSplitter:
"""EPUB 分块器:ebooklib 提取各章节文字 → TextSplitter 分块.
实现 Splitter Protocol,内部组合 TextSplitter 实例。
split() 的 text 参数实际接收 EPUB 文件路径(非文本内容)。
split() 的 source 参数接收 EPUB 文件路径(非文本内容)。
"""
def __init__(self, max_size: int = 1000, overlap: int = 100):
self._text_splitter = TextSplitter(max_size=max_size, overlap=overlap)
def split(self, text: str, source_file: str = "") -> list[dict]:
def split(self, source: str, source_file: str = "") -> list[dict]:
"""从 EPUB 文件提取各章节文字并分块.
Args:
text: EPUB 文件路径(非文本内容,由 ingest_file 传入)
source: EPUB 文件路径
source_file: 来源文件名
"""
try:
@@ -30,7 +31,7 @@ class EPUBSplitter:
"EPUB 支持需要 ebooklib 库. 请执行: uv sync --extra epub"
)
epub_path = text
epub_path = source
try:
book = epub.read_epub(epub_path)
except Exception as e:
+1
View File
@@ -1,5 +1,6 @@
"""HTML 文档分块器 — 使用 BeautifulSoup 去标签后委托 TextSplitter."""
import logging
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
+1
View File
@@ -1,5 +1,6 @@
"""Markdown 文档分块器."""
import re
from src.core.splitters.base import BaseTextSplitter
+5 -4
View File
@@ -1,5 +1,6 @@
"""PDF 文档分块器 — 使用 pymupdf 提取文字后委托 TextSplitter."""
import logging
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
@@ -9,17 +10,17 @@ class PDFSplitter:
"""PDF 分块器:pymupdf 提取文字 → TextSplitter 分块.
实现 Splitter Protocol,内部组合 TextSplitter 实例。
注意: split() 的 text 参数实际接收 PDF 文件路径(非文本内容)。
split() 的 source 参数接收 PDF 文件路径(非文本内容)。
"""
def __init__(self, max_size: int = 1000, overlap: int = 100):
self._text_splitter = TextSplitter(max_size=max_size, overlap=overlap)
def split(self, text: str, source_file: str = "") -> list[dict]:
def split(self, source: str, source_file: str = "") -> list[dict]:
"""从 PDF 文件提取文字并分块.
Args:
text: PDF 文件路径(非文本内容!由 ingest_file 传入)
source: PDF 文件路径
source_file: 来源文件名
"""
try:
@@ -29,7 +30,7 @@ class PDFSplitter:
"PDF 支持需要 pymupdf 库. 请执行: uv sync --extra pdf"
)
pdf_path = text # text 参数实际是文件路径
pdf_path = source
extracted_pages = []
try:
with fitz.open(pdf_path) as doc:
+6
View File
@@ -1,5 +1,6 @@
"""Splitter 注册表 — 按文件扩展名自动选择分块器."""
from pathlib import Path
from src.core.splitters.base import Splitter
# 扩展名 → Splitter 类名映射
@@ -11,6 +12,7 @@ _DEFAULT_MAP: dict[str, str] = {
".html": "html",
".htm": "html",
".epub": "epub",
".docx": "docx",
}
# 所有支持的扩展名集合(供外部遍历文件使用)
@@ -73,6 +75,10 @@ def get_splitter(
from src.core.splitters.epub import EPUBSplitter
return EPUBSplitter(max_size=max_size, overlap=overlap)
if kind == "docx":
from src.core.splitters.docx import DocxSplitter
return DocxSplitter(max_size=max_size, overlap=overlap)
# 回退
from src.core.splitters.text import TextSplitter
return TextSplitter(max_size=max_size, overlap=overlap)
+42 -7
View File
@@ -1,19 +1,19 @@
"""FastAPI 服务层."""
import os
import uuid
import logging
import os
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Depends, Request, Query
from fastapi import Depends, FastAPI, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse
from pydantic import BaseModel, Field, model_validator
from src.core.security import is_path_within_workspace
from src.server.auth import verify_api_key, rate_limiter
from src.server.deps import get_state, AppState
from src.server.auth import rate_limiter, verify_api_key
from src.server.deps import AppState, get_state
logger = logging.getLogger("md-vector-db")
@@ -61,8 +61,15 @@ async def lifespan(app: FastAPI):
yield
from fastapi.staticfiles import StaticFiles
app = FastAPI(title="md-vector-db", version="0.1.0", lifespan=lifespan)
# 挂载 Web 管理界面
web_dir = Path(__file__).parent.parent / "web"
if web_dir.exists():
app.mount("/admin", StaticFiles(directory=str(web_dir), html=True), name="admin")
# CORS 中间件
app.add_middleware(
CORSMiddleware,
@@ -84,8 +91,36 @@ async def security_headers_middleware(request: Request, call_next):
return response
@app.middleware("http")
async def body_size_limit_middleware(request: Request, call_next):
"""限制请求体大小(防止内存耗尽攻击)."""
content_length = request.headers.get("content-length")
max_size = int(os.environ.get("MAX_REQUEST_BODY_SIZE", str(10 * 1024 * 1024))) # 默认 10MB
if content_length and int(content_length) > max_size:
raise HTTPException(status_code=413, detail="请求体过大")
return await call_next(request)
@app.middleware("http")
async def audit_log_middleware(request: Request, call_next):
"""记录所有 API 请求的审计日志."""
start = time.time()
response = await call_next(request)
duration_ms = (time.time() - start) * 1000
logger.info(
"audit: %s %s%d (%.1fms) [%s]",
request.method, request.url.path,
response.status_code, duration_ms,
request.client.host if request.client else "unknown",
)
return response
@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
# 健康检查和根路径不需要速率限制
if request.url.path in ("/api/v1/health", "/"):
return await call_next(request)
await rate_limiter(request)
response = await call_next(request)
return response
+3 -3
View File
@@ -1,9 +1,9 @@
"""API 认证与安全中间件."""
import hmac
import os
import time
import threading
import logging
import os
import threading
import time
from collections import defaultdict
from fastapi import Header, HTTPException, Request
+5 -2
View File
@@ -1,7 +1,7 @@
"""FastAPI 依赖注入 — 集中管理应用状态, 替代模块级全局变量."""
import logging
import os
import threading
import logging
from src.core.config import load_config
from src.core.db import VectorDB
@@ -38,7 +38,10 @@ class AppState:
name = collection or self.default_collection
with self._cache_lock:
if name not in self._searchers:
self._searchers[name] = Searcher(self.db, self.embedder, name)
self._searchers[name] = Searcher(
self.db, self.embedder, name,
search_config=self.config.search,
)
return self._searchers[name]
def get_ingestor(self, collection: str | None = None) -> DocumentIngestor:
+147
View File
@@ -0,0 +1,147 @@
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>md-vector-db 管理面板</title>
<script src="https://unpkg.com/vue@3/dist/vue.global.prod.js"></script>
<style>
:root { --bg: #f8f9fa; --card-bg: #fff; --text: #212529; --muted: #6c757d; --border: #dee2e6; --primary: #0d6efd; --success: #198754; --danger: #dc3545; }
* { margin:0; padding:0; box-sizing:border-box; }
body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; background: var(--bg); color: var(--text); line-height: 1.6; }
.container { max-width: 960px; margin: 0 auto; padding: 24px 16px; }
h1 { font-size: 1.5rem; margin-bottom: 24px; }
.card { background: var(--card-bg); border: 1px solid var(--border); border-radius: 8px; padding: 20px; margin-bottom: 16px; }
.card h2 { font-size: 1.1rem; margin-bottom: 12px; }
.form-group { margin-bottom: 12px; }
label { display: block; font-size: .875rem; color: var(--muted); margin-bottom: 4px; }
input, textarea { width: 100%; padding: 8px 12px; border: 1px solid var(--border); border-radius: 4px; font-size: .9rem; font-family: inherit; }
textarea { min-height: 100px; resize: vertical; }
.btn { display: inline-block; padding: 8px 16px; border: none; border-radius: 4px; font-size: .875rem; cursor: pointer; }
.btn-primary { background: var(--primary); color: #fff; }
.btn-sm { padding: 4px 10px; font-size: .8rem; }
.result-item { border: 1px solid var(--border); border-radius: 6px; padding: 12px; margin-bottom: 8px; }
.result-item .score { font-weight: 600; color: var(--primary); }
.result-item .source { font-size: .8rem; color: var(--muted); }
.result-item .section { font-size: .85rem; color: var(--success); margin-bottom: 4px; }
.badge { display: inline-block; padding: 2px 8px; border-radius: 12px; font-size: .75rem; font-weight: 600; }
.badge-ok { background: #d1e7dd; color: #0f5132; }
.badge-err { background: #f8d7da; color: #842029; }
.tabs { display: flex; gap: 4px; margin-bottom: 16px; }
.tab { padding: 8px 16px; border: 1px solid var(--border); border-radius: 6px 6px 0 0; background: var(--bg); cursor: pointer; font-size: .875rem; }
.tab.active { background: var(--card-bg); border-bottom-color: var(--card-bg); font-weight: 600; }
.toast { position: fixed; top: 16px; right: 16px; padding: 12px 20px; border-radius: 6px; color: #fff; font-size: .875rem; z-index: 999; }
.toast-success { background: var(--success); }
.toast-error { background: var(--danger); }
table { width: 100%; border-collapse: collapse; }
th { text-align: left; padding: 8px; border-bottom: 2px solid var(--border); }
td { padding: 8px; border-bottom: 1px solid var(--border); }
</style>
</head>
<body>
<div id="app">
<div class="container">
<h1>📚 md-vector-db 管理面板</h1>
<div class="tabs">
<div :class="['tab', { active: activeTab === 'search' }]" @click="activeTab = 'search'">🔍 搜索</div>
<div :class="['tab', { active: activeTab === 'ingest' }]" @click="activeTab = 'ingest'">📥 入库</div>
<div :class="['tab', { active: activeTab === 'collections' }]" @click="activeTab = 'collections'">📂 集合</div>
</div>
<div v-if="toast" :class="['toast', 'toast-' + toast.type]">{{ toast.message }}</div>
<!-- 搜索 -->
<div v-if="activeTab === 'search'" class="card">
<h2>语义检索</h2>
<div class="form-group"><label>搜索查询</label><input v-model="searchQuery" @keyup.enter="doSearch" placeholder="输入关键词或自然语言..."></div>
<div style="display:flex;gap:12px">
<div class="form-group" style="flex:1"><label>返回条数</label><input v-model.number="topK" type="number" min="1" max="100"></div>
<div class="form-group" style="flex:1"><label>集合名</label><input v-model="collectionName" placeholder="default"></div>
</div>
<button class="btn btn-primary" @click="doSearch" :disabled="searching">{{ searching ? '搜索中...' : '搜索' }}</button>
<div v-if="searchResults.length" style="margin-top:16px">
<h3>{{ searchResults.length }} 条结果</h3>
<div v-for="(r,i) in searchResults" :key="i" class="result-item">
<div class="score">相似度: {{ (r.score || 0).toFixed(4) }}</div>
<div class="source">📄 {{ r.source_file }}</div>
<div v-if="r.section_title" class="section">📑 {{ r.section_title }}</div>
<pre style="white-space:pre-wrap;font-size:.85rem;margin-top:8px">{{ (r.content||'').substring(0,300) }}{{ (r.content||'').length > 300 ? '...' : '' }}</pre>
</div>
</div>
</div>
<!-- 入库 -->
<div v-if="activeTab === 'ingest'" class="card">
<h2>入库文档</h2>
<div class="form-group"><label>内容</label><textarea v-model="ingestContent" placeholder="粘贴 Markdown/文本..."></textarea></div>
<div class="form-group"><label>文件名</label><input v-model="ingestFileName" placeholder="document.md"></div>
<div class="form-group"><label>目标集合</label><input v-model="ingestCollection" placeholder="default"></div>
<button class="btn btn-primary" @click="doIngest" :disabled="ingesting">{{ ingesting ? '入库中...' : '入库' }}</button>
<p v-if="ingestResult" style="margin-top:12px">{{ ingestResult }}</p>
</div>
<!-- 集合 -->
<div v-if="activeTab === 'collections'" class="card">
<h2>集合列表</h2>
<button class="btn btn-primary btn-sm" @click="loadCollections" :disabled="loadingColl">{{ loadingColl ? '加载中...' : '刷新' }}</button>
<table v-if="collections.length" style="margin-top:12px">
<thead><tr><th>名称</th><th style="text-align:right">Chunks</th></tr></thead>
<tbody><tr v-for="c in collections" :key="c.name"><td>{{ c.name }}</td><td style="text-align:right">{{ c.count }}</td></tr></tbody>
</table>
</div>
<!-- 健康状态 -->
<div class="card" v-if="health">
<h2>服务状态 <span :class="['badge', health.status==='ok' ? 'badge-ok' : 'badge-err']">{{ health.status }}</span></h2>
<pre style="font-size:.8rem">{{ JSON.stringify(health.checks, null, 2) }}</pre>
</div>
</div>
</div>
<script>
const { createApp } = Vue
createApp({
data() {
return {
activeTab: 'search', searchQuery: '', topK: 10, collectionName: '',
searchResults: [], searching: false,
ingestContent: '', ingestFileName: '', ingestCollection: '', ingesting: false, ingestResult: null,
collections: [], loadingColl: false, health: null, toast: null, apiBase: window.location.origin,
}
},
mounted() { this.checkHealth() },
methods: {
async api(method, path, body) {
const headers = { 'Content-Type': 'application/json' }
const res = await fetch(`${this.apiBase}${path}`, { method, headers, body: body ? JSON.stringify(body) : undefined })
if (!res.ok) { const err = await res.text(); throw new Error(`HTTP ${res.status}: ${err}`) }
return res.json()
},
showToast(msg, type='success') { this.toast = { message: msg, type }; setTimeout(() => { this.toast = null }, 3000) },
async checkHealth() { try { this.health = await this.api('GET', '/api/v1/health') } catch(e) { this.health = { status: 'error', checks: {} } } },
async doSearch() {
if (!this.searchQuery.trim()) return; this.searching = true; this.searchResults = []
try {
const body = { query: this.searchQuery, top_k: this.topK }; if (this.collectionName) body.collection = this.collectionName
const data = await this.api('POST', '/api/v1/search', body); this.searchResults = data.results || []
} catch(e) { this.showToast(`搜索失败: ${e.message}`, 'error') } finally { this.searching = false }
},
async doIngest() {
if (!this.ingestContent.trim()) return; this.ingesting = true; this.ingestResult = null
try {
const body = { content: this.ingestContent, file_name: this.ingestFileName || 'untitled.md' }; if (this.ingestCollection) body.collection = this.ingestCollection
const data = await this.api('POST', '/api/v1/ingest', body)
this.ingestResult = `✅ 入库成功: ${data.chunks} chunks → ${data.collection}`; this.showToast(this.ingestResult)
} catch(e) { this.ingestResult = `❌ 入库失败: ${e.message}`; this.showToast(this.ingestResult, 'error') } finally { this.ingesting = false }
},
async loadCollections() {
this.loadingColl = true
try { const data = await this.api('GET', '/api/v1/collections'); this.collections = data.collections || [] }
catch(e) { this.showToast(`加载失败: ${e.message}`, 'error') } finally { this.loadingColl = false }
},
}
}).mount('#app')
</script>
</body>
</html>
+41
View File
@@ -0,0 +1,41 @@
"""基准测试共享 fixture."""
import pytest
from src.core.config import ChunkConfig, EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
from src.core.search import Searcher
@pytest.fixture(scope="module")
def benchmark_db(tmp_path_factory):
"""模块级共享 ChromaDB 实例."""
persist_dir = tmp_path_factory.mktemp("bench_data")
return VectorDB(persist_dir=str(persist_dir))
@pytest.fixture(scope="module")
def benchmark_embedder():
"""模块级共享 LocalEmbedder."""
config = EmbedConfig(mode="local", local_model="BAAI/bge-small-zh-v1.5")
return create_embedder(config)
@pytest.fixture(scope="module")
def benchmark_searcher(benchmark_db, benchmark_embedder):
"""预填充数据的 Searcher."""
ingestor = DocumentIngestor(
benchmark_db, benchmark_embedder, "bench_collection",
chunk_config=ChunkConfig(max_size=1000, overlap=100),
)
for i in range(100):
content = (
f"# 文档{i}\n\n"
+ "\n\n".join(
f"{j}段用于基准测试。关键词: Python, Rust, GPU, 向量数据库。"
for j in range(5)
)
)
ingestor.ingest_content(content, f"bench_{i}.md")
return Searcher(benchmark_db, benchmark_embedder, "bench_collection")
+23
View File
@@ -0,0 +1,23 @@
"""嵌入性能基准测试."""
import pytest
def test_bench_embed_single(benchmark, benchmark_embedder):
"""单文本嵌入耗时."""
text = "这是一段测试文本,用于测量嵌入速度。"
benchmark(benchmark_embedder.embed, [text])
def test_bench_embed_batch_32(benchmark, benchmark_embedder):
"""批量 32 文本嵌入耗时."""
texts = [f"测试文本第{i}条,模拟真实文档内容。" for i in range(32)]
benchmark(benchmark_embedder.embed, texts)
def test_bench_embed_batch_100(benchmark, benchmark_embedder):
"""批量 100 文本嵌入耗时(GPU 优势显著)."""
texts = [
f"测试文本第{i}条。Python 通用编程语言,用于数据科学和 AI。"
for i in range(100)
]
benchmark(benchmark_embedder.embed, texts)
+17
View File
@@ -0,0 +1,17 @@
"""检索性能基准测试."""
import pytest
def test_bench_search_top5(benchmark, benchmark_searcher):
"""top_k=5 检索耗时."""
benchmark(benchmark_searcher.search, "Python 向量数据库", top_k=5)
def test_bench_search_top20(benchmark, benchmark_searcher):
"""top_k=20 检索耗时."""
benchmark(benchmark_searcher.search, "Rust 编程语言 GPU", top_k=20)
def test_bench_search_cold_start(benchmark, benchmark_searcher):
"""冷启动检索耗时."""
benchmark(benchmark_searcher.search, "GPU 加速 深度学习 嵌入", top_k=10)
+13
View File
@@ -138,3 +138,16 @@ class TestDeleteEndpoint:
results = search_resp.json()["results"]
sources = [r["source_file"] for r in results]
assert "delete-test.md" not in sources
def test_admin_ui_served(client):
"""管理界面可访问."""
response = client.get("/admin")
assert response.status_code == 200
def test_health_skips_rate_limit(client):
"""健康检查多次访问不触发速率限制."""
for _ in range(5):
resp = client.get("/api/v1/health")
assert resp.status_code == 200
+2 -2
View File
@@ -1,11 +1,11 @@
"""认证与速率限制测试."""
import time
import threading
import time
import pytest
from fastapi import HTTPException
from src.server.auth import verify_api_key, RateLimiter
from src.server.auth import RateLimiter, verify_api_key
class TestVerifyApiKey:
+19
View File
@@ -1,5 +1,6 @@
"""CLI 命令测试."""
from typer.testing import CliRunner
from src.cli.main import app
runner = CliRunner()
@@ -60,3 +61,21 @@ class TestCLIJSONOutput:
if result.stdout.strip():
data = json.loads(result.stdout)
assert isinstance(data, dict)
def test_ingest_help():
"""ingest --help 正常输出."""
result = runner.invoke(app, ["ingest", "--help"])
assert result.exit_code == 0
def test_search_help():
"""search --help 正常输出."""
result = runner.invoke(app, ["search", "--help"])
assert result.exit_code == 0
def test_export_help():
"""export --help 正常输出."""
result = runner.invoke(app, ["export", "--help"])
assert result.exit_code == 0
+1 -4
View File
@@ -1,11 +1,8 @@
"""配置加载模块测试."""
import os
import tempfile
from pathlib import Path
import pytest
from src.core.config import AppConfig, EmbedConfig, ChunkConfig, load_config
from src.core.config import AppConfig, EmbedConfig, load_config
class TestEmbedConfig:
+30 -1
View File
@@ -1,5 +1,4 @@
"""数据库层测试."""
import gc
import tempfile
from pathlib import Path
@@ -46,3 +45,33 @@ class TestVectorDB:
# 再次获取会创建新的
col = db.get_or_create_collection("tmp_col")
assert col.count() == 0
def test_write_guard_context_manager(tmp_path):
"""write_guard 上下文管理器正常获取和释放锁."""
from src.core.db import VectorDB
vdb = VectorDB(persist_dir=str(tmp_path))
with vdb.write_guard():
pass
def test_close(tmp_path):
"""close 正常执行不抛异常."""
from src.core.db import VectorDB
vdb = VectorDB(persist_dir=str(tmp_path))
vdb.close()
def test_delete_collection_nonexistent(tmp_path):
"""删除不存在的 collection 不抛异常."""
from src.core.db import VectorDB
vdb = VectorDB(persist_dir=str(tmp_path))
vdb.delete_collection("nonexistent-collection-12345")
def test_delete_by_source_no_match(tmp_path):
"""删除不存在的 source 返回 False."""
from src.core.db import VectorDB
vdb = VectorDB(persist_dir=str(tmp_path))
result = vdb.delete_by_source("test_col", "no-such-file.md")
assert result is False
-3
View File
@@ -1,10 +1,7 @@
"""AppState 和依赖注入测试."""
import os
import tempfile
import pytest
from src.core.config import load_config
from src.server.deps import AppState
+6 -2
View File
@@ -3,8 +3,12 @@ import pytest
from src.core.config import EmbedConfig
from src.core.embedder import (
LocalEmbedder, OpenAIEmbedder, DashscopeEmbedder,
create_embedder, batch_embed, SUPPORTED_PROVIDERS,
SUPPORTED_PROVIDERS,
DashscopeEmbedder,
LocalEmbedder,
OpenAIEmbedder,
batch_embed,
create_embedder,
)
+106
View File
@@ -0,0 +1,106 @@
"""导出功能测试."""
import csv
import io
import json
import pytest
from src.core.search import Searcher
class FakeExportCollection:
def count(self):
return 2
def get(self, include=None):
return {
"ids": ["doc_0", "doc_1"],
"documents": ["内容A。\n\n段落B。", "内容C。"],
"metadatas": [
{"source_file": "a.md", "section_title": "标题A", "heading_level": 1, "chunk_index": 0},
{"source_file": "b.md", "section_title": "", "heading_level": 0, "chunk_index": 0},
],
}
class FakeExportDB:
def get_or_create_collection(self, name):
return FakeExportCollection()
def list_collections(self):
return []
class FakeExportEmbedder:
@property
def dimension(self):
return 4
def embed(self, texts):
return [[0.1, 0.2, 0.3, 0.4] for _ in texts]
@pytest.fixture
def searcher():
db = FakeExportDB()
embedder = FakeExportEmbedder()
return Searcher(db, embedder, "test_export")
def test_export_json_returns_valid_json(searcher):
"""export_json 返回合法的 JSON 字符串."""
output = searcher.export_json()
data = json.loads(output)
assert isinstance(data, list)
assert len(data) == 2
def test_export_json_contains_all_fields(searcher):
"""导出包含所有必要字段."""
output = searcher.export_json()
data = json.loads(output)
first = data[0]
assert "id" in first
assert "content" in first
assert "source_file" in first
assert "section_title" in first
def test_export_csv_returns_valid_csv(searcher):
"""export_csv 返回合法的 CSV 字符串."""
output = searcher.export_csv()
reader = csv.DictReader(io.StringIO(output))
rows = list(reader)
assert len(rows) == 2
def test_export_csv_has_header(searcher):
"""CSV 包含表头."""
output = searcher.export_csv()
reader = csv.DictReader(io.StringIO(output))
assert reader.fieldnames is not None
assert "content" in reader.fieldnames
assert "source_file" in reader.fieldnames
def test_export_empty_collection():
"""空 collection 导出空列表."""
class EmptyCollection:
def count(self):
return 0
def get(self, include=None):
return {"ids": [], "documents": [], "metadatas": []}
class EmptyDB:
def get_or_create_collection(self, name):
return EmptyCollection()
def list_collections(self):
return []
s = Searcher(EmptyDB(), FakeExportEmbedder(), "empty")
json_out = s.export_json()
assert json.loads(json_out) == []
csv_out = s.export_csv()
lines = csv_out.strip().split("\n")
assert len(lines) == 1 # 仅表头
+104
View File
@@ -0,0 +1,104 @@
"""文件变更追踪器测试."""
import hashlib
from pathlib import Path
from src.core.file_tracker import FileTracker
def test_compute_file_hash(tmp_path: Path):
"""计算文件 SHA256 哈希."""
file = tmp_path / "test.md"
file.write_text("hello world", encoding="utf-8")
result = FileTracker.compute_hash(str(file))
expected = hashlib.sha256(b"hello world").hexdigest()
assert result == expected
def test_compute_hash_deterministic(tmp_path: Path):
"""同一文件内容产生相同哈希."""
file = tmp_path / "a.md"
file.write_text("same content", encoding="utf-8")
assert FileTracker.compute_hash(str(file)) == FileTracker.compute_hash(str(file))
def test_hash_changes_with_content(tmp_path: Path):
"""内容变更导致哈希不同."""
file = tmp_path / "b.md"
file.write_text("v1", encoding="utf-8")
h1 = FileTracker.compute_hash(str(file))
file.write_text("v2", encoding="utf-8")
h2 = FileTracker.compute_hash(str(file))
assert h1 != h2
def test_is_stale_new_file(tmp_path: Path):
"""新文件(无记录)视为过期."""
file = tmp_path / "new.md"
file.write_text("content", encoding="utf-8")
tracker = FileTracker(tmp_path / "tracker.json")
record = tracker.get_record(str(file))
assert record is None
assert tracker.is_stale(str(file)) is True
def test_is_stale_unchanged_file(tmp_path: Path):
"""未修改文件视为未过期."""
file = tmp_path / "unchanged.md"
file.write_text("stable", encoding="utf-8")
tracker = FileTracker(tmp_path / "tracker.json")
tracker.mark_ingested(str(file))
assert tracker.is_stale(str(file)) is False
def test_is_stale_modified_file(tmp_path: Path):
"""修改后文件视为过期."""
file = tmp_path / "modified.md"
file.write_text("v1", encoding="utf-8")
tracker = FileTracker(tmp_path / "tracker.json")
tracker.mark_ingested(str(file))
file.write_text("v2", encoding="utf-8")
assert tracker.is_stale(str(file)) is True
def test_mark_ingested_updates_record(tmp_path: Path):
"""mark_ingested 创建/更新记录."""
file = tmp_path / "x.md"
file.write_text("hello", encoding="utf-8")
tracker = FileTracker(tmp_path / "tracker.json")
tracker.mark_ingested(str(file))
record = tracker.get_record(str(file))
assert record is not None
assert record["hash"] == FileTracker.compute_hash(str(file))
assert "ingested_at" in record
def test_persistence_across_instances(tmp_path: Path):
"""tracker 数据持久化到 JSON,跨实例可读."""
file = tmp_path / "p.md"
file.write_text("persist me", encoding="utf-8")
db_path = tmp_path / "tracker.json"
t1 = FileTracker(db_path)
t1.mark_ingested(str(file))
t2 = FileTracker(db_path)
assert t2.is_stale(str(file)) is False
def test_file_deleted_considered_stale(tmp_path: Path):
"""文件被删除后视为过期."""
file = tmp_path / "tmp.md"
file.write_text("temp", encoding="utf-8")
tracker = FileTracker(tmp_path / "tracker.json")
tracker.mark_ingested(str(file))
file.unlink()
assert tracker.is_stale(str(file)) is True
def test_binary_file_hash(tmp_path: Path):
"""二进制文件也能正确计算哈希(如 PDF)."""
file = tmp_path / "doc.pdf"
file.write_bytes(b"\x00\x01\x02\x03")
h = FileTracker.compute_hash(str(file))
assert len(h) == 64
assert h == hashlib.sha256(b"\x00\x01\x02\x03").hexdigest()
+108 -8
View File
@@ -4,8 +4,8 @@ from pathlib import Path
import pytest
from src.core.splitters import MarkdownSplitter
from src.core.ingest import DocumentIngestor
from src.core.splitters import MarkdownSplitter
class TestMarkdownSplitter:
@@ -122,7 +122,6 @@ class TestIngestorIntegration:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
db = VectorDB(persist_dir=str(tmp_path))
embedder = create_embedder(EmbedConfig(mode="local"))
@@ -137,13 +136,12 @@ class TestIngestorIntegration:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
db = VectorDB(persist_dir=str(tmp_path))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(db, embedder, "test_dedup")
c1 = ingestor.ingest_content("# A", "dup.md")
ingestor.ingest_content("# A", "dup.md")
c2 = ingestor.ingest_content("# B", "dup.md")
assert ingestor.collection.count() == c2
@@ -156,7 +154,6 @@ class TestIngestFile:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
md_file = tmp_path / "hello.md"
md_file.write_text("# 测试\n这是测试内容。", encoding="utf-8")
@@ -174,7 +171,6 @@ class TestIngestFile:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
txt_file = tmp_path / "notes.txt"
txt_file.write_text("这是一段纯文本内容。\n\n第二段内容在这里。", encoding="utf-8")
@@ -195,7 +191,6 @@ class TestIngestDirectory:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
(tmp_path / "a.md").write_text("# A\n内容 A", encoding="utf-8")
(tmp_path / "b.txt").write_text("内容 B", encoding="utf-8")
@@ -213,7 +208,6 @@ class TestIngestDirectory:
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
db = VectorDB(persist_dir=str(tmp_path / "db"))
embedder = create_embedder(EmbedConfig(mode="local"))
@@ -221,3 +215,109 @@ class TestIngestDirectory:
results = ingestor.ingest_directory(str(tmp_path))
assert results == {}
class TestIncrementalIngest:
"""增量入库测试."""
def test_ingest_file_incremental_skips_unchanged(self, tmp_path):
"""增量模式: 未修改的文件跳过入库."""
from src.core.config import ChunkConfig, EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.file_tracker import FileTracker
file = tmp_path / "stable.md"
file.write_text("# 稳定文档\n\n内容不变。", encoding="utf-8")
tracker_path = str(tmp_path / "tracker.json")
db = VectorDB(persist_dir=str(tmp_path / "db"))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(
db, embedder, "test_incr",
chunk_config=ChunkConfig(max_size=1000, overlap=100),
file_tracker=FileTracker(tracker_path),
)
count1 = ingestor.ingest_file(str(file), incremental=True)
assert count1 > 0
count2 = ingestor.ingest_file(str(file), incremental=True)
assert count2 == 0 # 跳过
def test_ingest_file_incremental_reingests_modified(self, tmp_path):
"""增量模式: 修改后的文件重新入库."""
from src.core.config import ChunkConfig, EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.file_tracker import FileTracker
file = tmp_path / "changing.md"
file.write_text("# v1\n\n初始版本的内容段落。", encoding="utf-8")
tracker_path = str(tmp_path / "tracker2.json")
db = VectorDB(persist_dir=str(tmp_path / "db2"))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(
db, embedder, "test_incr2",
chunk_config=ChunkConfig(max_size=1000, overlap=100),
file_tracker=FileTracker(tracker_path),
)
count1 = ingestor.ingest_file(str(file), incremental=True)
assert count1 > 0
file.write_text("# v2\n\n新增段落,内容完全不同了。", encoding="utf-8")
count2 = ingestor.ingest_file(str(file), incremental=True)
assert count2 > 0
def test_ingest_file_force_mode_always_reingests(self, tmp_path):
"""force=True 时始终重新入库(忽略 tracker)."""
from src.core.config import ChunkConfig, EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.file_tracker import FileTracker
file = tmp_path / "force.md"
file.write_text("# force test\n\n这是强制入库测试的内容。", encoding="utf-8")
tracker_path = str(tmp_path / "tracker3.json")
db = VectorDB(persist_dir=str(tmp_path / "db3"))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(
db, embedder, "test_force",
chunk_config=ChunkConfig(max_size=1000, overlap=100),
file_tracker=FileTracker(tracker_path),
)
count1 = ingestor.ingest_file(str(file), incremental=True)
count2 = ingestor.ingest_file(str(file), incremental=True, force=True)
assert count1 > 0
assert count2 > 0 # force 模式重新入库
def test_ingest_content_default_splitter(tmp_path):
"""未指定 splitter 时用 MarkdownSplitter."""
from src.core.config import EmbedConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
db = VectorDB(persist_dir=str(tmp_path / "db"))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(db, embedder, "test_content")
count = ingestor.ingest_content("# 测试\n\n一些内容。", "test.md")
assert count >= 1
def test_ingest_directory_recursive(tmp_path):
"""ingest_directory 递归处理子目录."""
from src.core.config import EmbedConfig, ChunkConfig
from src.core.db import VectorDB
from src.core.embedder import create_embedder
from src.core.ingest import DocumentIngestor
(tmp_path / "sub").mkdir()
(tmp_path / "a.md").write_text("# A\n\n内容A。", encoding="utf-8")
(tmp_path / "sub" / "b.md").write_text("# B\n\n内容B。", encoding="utf-8")
db = VectorDB(persist_dir=str(tmp_path / "db_r"))
embedder = create_embedder(EmbedConfig(mode="local"))
ingestor = DocumentIngestor(
db, embedder, "test_recurse",
chunk_config=ChunkConfig(max_size=1000, overlap=100),
)
results = ingestor.ingest_directory(str(tmp_path))
assert len(results) >= 2
assert all(v > 0 for v in results.values())
+87
View File
@@ -0,0 +1,87 @@
"""重排序器测试."""
from src.core.reranker import Reranker
class FakeCrossEncoder:
"""模拟 Cross-Encoder 模型."""
def predict(self, pairs, **kwargs):
# 包含"重要"的 pair 分数高
scores = []
for pair in pairs:
score = 5.0 if "重要" in pair[1] else 1.0
scores.append(score)
return scores
def test_reranker_returns_same_count():
"""重排序不改变结果数量."""
reranker = Reranker(model_name="test-model")
reranker._model = FakeCrossEncoder()
candidates = [
{"content": "普通文档", "score": 0.8},
{"content": "重要文档", "score": 0.6},
{"content": "另一个普通", "score": 0.7},
]
result = reranker.rerank("查询", candidates, top_k=3)
assert len(result) == 3
def test_reranker_promotes_relevant():
"""重排序将更相关的内容提前."""
reranker = Reranker(model_name="test-model")
reranker._model = FakeCrossEncoder()
candidates = [
{"content": "普通 A", "score": 0.9},
{"content": "重要内容在这里", "score": 0.5},
{"content": "普通 B", "score": 0.7},
]
result = reranker.rerank("查询", candidates, top_k=3)
assert "重要" in result[0]["content"]
def test_reranker_truncates_to_top_k():
"""rerank 截断到指定的 top_k."""
reranker = Reranker(model_name="test-model")
reranker._model = FakeCrossEncoder()
candidates = [
{"content": f"文档{i}", "score": 0.9 - i * 0.1}
for i in range(20)
]
result = reranker.rerank("查询", candidates, top_k=5)
assert len(result) == 5
def test_reranker_empty_input():
"""空输入返回空列表."""
reranker = Reranker(model_name="test-model")
result = reranker.rerank("查询", [], top_k=5)
assert result == []
def test_reranker_preserves_metadata():
"""重排序保留文档元数据."""
reranker = Reranker(model_name="test-model")
reranker._model = FakeCrossEncoder()
candidates = [
{
"content": "带元数据的文档",
"score": 0.5,
"source_file": "meta.md",
"section_title": "第一章",
}
]
result = reranker.rerank("查询", candidates, top_k=1)
assert result[0]["source_file"] == "meta.md"
assert result[0]["section_title"] == "第一章"
def test_reranker_score_replaced_with_rerank():
"""重排序后 score 更新为 rerank_score."""
reranker = Reranker(model_name="test-model")
reranker._model = FakeCrossEncoder()
candidates = [{"content": "测试", "score": 0.5}]
result = reranker.rerank("查询", candidates, top_k=1)
assert "rerank_score" in result[0]
assert result[0]["score"] == result[0]["rerank_score"]
+161
View File
@@ -0,0 +1,161 @@
"""混合检索器测试."""
import hashlib
import pytest
from src.core.retriever import HybridRetriever
class FakeEmbedder:
"""模拟嵌入器 — 返回伪向量."""
@property
def dimension(self) -> int:
return 4
def embed(self, texts: list[str]) -> list[list[float]]:
result = []
for t in texts:
h = hashlib.md5(t.encode()).digest()
vec = [float(b) / 255.0 for b in h[:4]]
result.append(vec)
return result
class FakeCollection:
"""模拟 ChromaDB collection."""
def __init__(self):
self._docs: list[dict] = []
def add(self, ids, embeddings, documents, metadatas):
for i, doc_id in enumerate(ids):
self._docs.append({
"id": doc_id,
"embedding": embeddings[i] if embeddings else [],
"document": documents[i],
"metadata": metadatas[i] if metadatas else {},
})
def query(self, query_embeddings, n_results, where=None, include=None):
# 返回所有已入库文档(模拟向量检索)
n = min(n_results, len(self._docs))
if n == 0:
return {"ids": [[]], "documents": [[]], "metadatas": [[]], "distances": [[]]}
ids_list = [d["id"] for d in self._docs[:n]]
docs_list = [d["document"] for d in self._docs[:n]]
metas_list = [d["metadata"] for d in self._docs[:n]]
dists_list = [0.2 + i * 0.05 for i in range(n)] # 伪距离
return {
"ids": [ids_list],
"documents": [docs_list],
"metadatas": [metas_list],
"distances": [dists_list],
}
def get(self, include=None):
return {
"ids": [d["id"] for d in self._docs],
"documents": [d["document"] for d in self._docs],
"metadatas": [d["metadata"] for d in self._docs],
}
def count(self) -> int:
return len(self._docs)
def delete(self, ids):
self._docs = [d for d in self._docs if d["id"] not in ids]
class FakeDB:
"""模拟 VectorDB."""
def __init__(self):
self._collections: dict[str, FakeCollection] = {}
def get_or_create_collection(self, name: str):
if name not in self._collections:
self._collections[name] = FakeCollection()
return self._collections[name]
@pytest.fixture
def retriever():
db = FakeDB()
embedder = FakeEmbedder()
return HybridRetriever(db, embedder, "test", bm25_weight=0.3)
def test_search_returns_list(retriever):
"""基本语义检索返回列表."""
results = retriever.search("测试查询", top_k=5)
assert isinstance(results, list)
def test_search_with_source_filter(retriever):
"""按 source_file 过滤."""
results = retriever.search("查询", top_k=5, source_file="doc.md")
assert isinstance(results, list)
def test_search_top_k_bounds(retriever):
"""top_k 在合理范围内."""
for k in [1, 10, 50]:
results = retriever.search("test", top_k=k)
assert len(results) <= k
def test_bm25_index_built_from_collection(retriever):
"""BM25 索引从 collection 文档构建."""
coll = retriever._db.get_or_create_collection("test")
coll.add(
ids=["doc_0", "doc_1", "doc_2"],
embeddings=[[0.1] * 4, [0.2] * 4, [0.3] * 4],
documents=["Python 是一门编程语言", "Java 也是编程语言", "今天天气很好"],
metadatas=[
{"source_file": "a.md"},
{"source_file": "b.md"},
{"source_file": "c.md"},
],
)
retriever._bm25_index = None # 强制重建
results = retriever.search("编程语言", top_k=2)
assert len(results) == 2
assert any("编程语言" in r["content"] for r in results)
def test_hybrid_score_fusion(retriever):
"""混合分数融合:向量分 + BM25 分加权."""
coll = retriever._db.get_or_create_collection("test")
coll.add(
ids=["d0", "d1"],
embeddings=[[1.0] * 4, [0.5] * 4],
documents=["Docker 容器化部署指南", "Python 数据分析入门"],
metadatas=[{"source_file": "x.md"}, {"source_file": "y.md"}],
)
retriever._bm25_index = None
results = retriever.search("Docker 部署", top_k=2)
assert len(results) >= 1
assert "Docker" in results[0]["content"]
def test_empty_collection_returns_empty(retriever):
"""空 collection 返回空列表."""
results = retriever.search("查询", top_k=5)
assert results == []
def test_metadata_in_results(retriever):
"""结果中包含完整元数据."""
coll = retriever._db.get_or_create_collection("test")
coll.add(
ids=["meta_test"],
embeddings=[[0.5] * 4],
documents=["带元数据的文档"],
metadatas=[{"source_file": "meta.md", "section_title": "第一章", "heading_level": 1}],
)
retriever._bm25_index = None
results = retriever.search("元数据", top_k=1)
assert len(results) == 1
assert results[0]["source_file"] == "meta.md"
assert results[0]["section_title"] == "第一章"
+22
View File
@@ -122,3 +122,25 @@ class TestSearcher:
searcher = Searcher(db, embedder, "empty_coll")
sources = searcher.list_sources()
assert sources == []
def test_list_sources_empty():
"""空 collection 的 list_sources 返回空列表."""
from src.core.search import Searcher
class EmptyColl:
def count(self): return 0
def get(self, **kwargs):
return {"ids": [], "documents": [], "metadatas": []}
class EmptyDB:
def get_or_create_collection(self, name): return EmptyColl()
def list_collections(self): return []
class FakeEmb:
@property
def dimension(self): return 4
def embed(self, texts): return [[0.0] * 4]
s = Searcher(EmptyDB(), FakeEmb(), "empty")
assert s.list_sources() == []
-1
View File
@@ -1,5 +1,4 @@
"""is_safe_path 路径遍历防护测试."""
import pytest
from src.core.security import is_safe_path
+7 -2
View File
@@ -1,6 +1,11 @@
"""Splitter 注册表和 TextSplitter 测试."""
import pytest
from src.core.splitters import TextSplitter, MarkdownSplitter, get_splitter, register_splitter, SUPPORTED_SUFFIXES
from src.core.splitters import (
SUPPORTED_SUFFIXES,
MarkdownSplitter,
TextSplitter,
get_splitter,
register_splitter,
)
class TestTextSplitter:
+53
View File
@@ -0,0 +1,53 @@
"""DocxSplitter 测试."""
import pytest
pytest.importorskip("markitdown", reason="需要 markitdown 库")
from src.core.splitters.docx import DocxSplitter
def test_docx_splitter_creates():
"""创建 DocxSplitter 实例."""
s = DocxSplitter(max_size=500, overlap=50)
assert s is not None
def test_docx_splitter_empty_text():
"""空纯文本返回空列表."""
s = DocxSplitter()
result = s.split(" ", source_file="empty.docx")
assert result == []
def test_docx_splitter_basic_text():
"""基本文本文档分块."""
s = DocxSplitter(max_size=200, overlap=20)
text = "段落A。\n\n段落B。\n\n段落C。"
result = s.split(text, source_file="test.docx")
assert len(result) >= 1
assert all("content" in r for r in result)
assert all(r["source_file"] == "test.docx" for r in result)
def test_docx_splitter_long_text():
"""长文本分多块."""
s = DocxSplitter(max_size=100, overlap=10)
text = "这是一段非常长的文本。\n\n" * 50
result = s.split(text, source_file="long.docx")
assert len(result) >= 5
def test_docx_splitter_source_file():
"""source_file 正确传递到每个 chunk."""
s = DocxSplitter(max_size=500, overlap=50)
result = s.split("测试内容。", source_file="myfile.docx")
assert all(r["source_file"] == "myfile.docx" for r in result)
def test_docx_splitter_chunk_index():
"""chunk_index 从 0 递增."""
s = DocxSplitter(max_size=100, overlap=10)
text = "chunk A。\n\n" * 20
result = s.split(text, source_file="index.docx")
indices = [r["chunk_index"] for r in result]
assert indices == list(range(len(result)))
+2 -1
View File
@@ -28,9 +28,10 @@ class TestEPUBSplitter:
def test_split_simple_epub(self, tmp_path):
"""用 ebooklib 创建一个简单 EPUB 并测试分块."""
from src.core.splitters.epub import EPUBSplitter
from ebooklib import epub
from src.core.splitters.epub import EPUBSplitter
epub_path = tmp_path / "test.epub"
book = epub.EpubBook()
+1
View File
@@ -1,5 +1,6 @@
"""MarkdownSplitter 边界测试."""
import pytest
from src.core.splitters import MarkdownSplitter
+5 -3
View File
@@ -1,6 +1,6 @@
"""PDFSplitter 测试."""
import pytest
from pathlib import Path
pymupdf = pytest.importorskip("fitz", reason="pymupdf 未安装")
@@ -10,9 +10,10 @@ class TestPDFSplitter:
def test_split_simple_pdf(self, tmp_path):
"""用 pymupdf 创建一个简单 PDF 并测试分块."""
from src.core.splitters.pdf import PDFSplitter
import fitz
from src.core.splitters.pdf import PDFSplitter
pdf_path = tmp_path / "test.pdf"
doc = fitz.open()
# 插入纯 ASCII 文本避免 CJK 字体编码问题
@@ -29,9 +30,10 @@ class TestPDFSplitter:
def test_empty_pdf(self, tmp_path):
"""空 PDF(有页但无文字)返回空列表."""
from src.core.splitters.pdf import PDFSplitter
import fitz
from src.core.splitters.pdf import PDFSplitter
pdf_path = tmp_path / "empty.pdf"
doc = fitz.open()
doc.new_page() # pymupdf 必须有至少一页才能保存
Generated
+149 -14
View File
@@ -2,9 +2,12 @@ version = 1
revision = 3
requires-python = ">=3.13"
resolution-markers = [
"python_full_version >= '3.15'",
"python_full_version == '3.14.*'",
"python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'linux'",
"python_full_version >= '3.15' and sys_platform == 'win32'",
"python_full_version >= '3.15' and sys_platform != 'win32'",
"python_full_version == '3.14.*' and sys_platform == 'win32'",
"python_full_version == '3.14.*' and sys_platform != 'win32'",
"python_full_version < '3.14' and sys_platform == 'win32'",
"python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'",
"(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux')",
]
@@ -466,6 +469,15 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/52/30/21b2ad45959cd50e909e02ebac1e30b4ceb7162e91c11d4c570223a458b7/coverage-7.15.0-py3-none-any.whl", hash = "sha256:56da6a4cbe8f7e9e80bd072ca9cefe67d7106a440a7ec06519ec6507ac94ad19", size = 212632, upload-time = "2026-07-02T13:10:48.641Z" },
]
[[package]]
name = "defusedxml"
version = "0.7.1"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0f/d5/c66da9b79e5bdb124974bfe172b4daf3c984ebd9c2a06e2b8a4dc7331c72/defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl", hash = "sha256:a352e7e428770286cc899e2542b6cdaedb2b4953ff269a210103ec58f6198a61" },
]
[[package]]
name = "distro"
version = "1.9.0"
@@ -1030,6 +1042,50 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/7f/2c/0f1e93c636720e8a3eb59af2bfda99d98b55891e1c53bc30c2e0e865f01b/lxml-6.1.1-cp314-cp314t-win_arm64.whl", hash = "sha256:58bb955caba94e467d2a96da17660d2d704e0675894cba21ab8a775b8621fd1c", size = 3817223, upload-time = "2026-05-19T19:22:56.823Z" },
]
[[package]]
name = "magika"
version = "0.6.2"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
resolution-markers = [
"python_full_version >= '3.15' and sys_platform == 'win32'",
"python_full_version == '3.14.*' and sys_platform == 'win32'",
"python_full_version < '3.14' and sys_platform == 'win32'",
]
dependencies = [
{ name = "click", marker = "sys_platform == 'win32'" },
{ name = "numpy", marker = "sys_platform == 'win32'" },
{ name = "onnxruntime", marker = "sys_platform == 'win32'" },
{ name = "python-dotenv", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fe/b6/8fdd991142ad3e037179a494b153f463024e5a211ef3ad948b955c26b4de/magika-0.6.2.tar.gz", hash = "sha256:37eb6ae8020f6e68f231bc06052c0a0cbe8e6fa27492db345e8dc867dbceb067", size = 3036634, upload-time = "2025-05-02T14:54:18.88Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/c2/07/4f7748f34279f2852068256992377474f9700b6fbad6735d6be58605178f/magika-0.6.2-py3-none-any.whl", hash = "sha256:5ef72fbc07723029b3684ef81454bc224ac5f60986aa0fc5a28f4456eebcb5b2", size = 2967609, upload-time = "2025-05-02T14:54:09.696Z" },
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/b0/1f/28e412d0ccedc068fbccdae6a6233faaa97ec3e5e2ffd242e49655b10064/magika-0.6.2-py3-none-win_amd64.whl", hash = "sha256:711f427a633e0182737dcc2074748004842f870643585813503ff2553b973b9f", size = 12385740, upload-time = "2025-05-02T14:54:14.096Z" },
]
[[package]]
name = "magika"
version = "0.6.3"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
resolution-markers = [
"python_full_version >= '3.15' and sys_platform != 'win32'",
"python_full_version == '3.14.*' and sys_platform != 'win32'",
"python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'",
"(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux')",
]
dependencies = [
{ name = "click", marker = "sys_platform != 'win32'" },
{ name = "numpy", marker = "sys_platform != 'win32'" },
{ name = "onnxruntime", marker = "sys_platform != 'win32'" },
{ name = "python-dotenv", marker = "sys_platform != 'win32'" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a3/f3/3d1dcdd7b9c41d589f5cff252d32ed91cdf86ba84391cfc81d9d8773571d/magika-0.6.3.tar.gz", hash = "sha256:7cc52aa7359af861957043e2bf7265ed4741067251c104532765cd668c0c0cb1", size = 3042784, upload-time = "2025-10-30T15:22:34.499Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/a2/e4/35c323beb3280482c94299d61626116856ac2d4ec16ecef50afc4fdd4291/magika-0.6.3-py3-none-any.whl", hash = "sha256:eda443d08006ee495e02083b32e51b98cb3696ab595a7d13900d8e2ef506ec9d", size = 2969474, upload-time = "2025-10-30T15:22:25.298Z" },
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/25/8f/132b0d7cd51c02c39fd52658a5896276c30c8cc2fd453270b19db8c40f7e/magika-0.6.3-py3-none-macosx_11_0_arm64.whl", hash = "sha256:86901e64b05dde5faff408c9b8245495b2e1fd4c226e3393d3d2a3fee65c504b", size = 13358841, upload-time = "2025-10-30T15:22:27.413Z" },
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/c4/03/5ed859be502903a68b7b393b17ae0283bf34195cfcca79ce2dc25b9290e7/magika-0.6.3-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:3d9661eedbdf445ac9567e97e7ceefb93545d77a6a32858139ea966b5806fb64", size = 15367335, upload-time = "2025-10-30T15:22:29.907Z" },
]
[[package]]
name = "markdown-it-py"
version = "4.2.0"
@@ -1042,6 +1098,37 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" },
]
[[package]]
name = "markdownify"
version = "1.2.3"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "beautifulsoup4" },
{ name = "six" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/92/ab/d1297139c0e2ceb151ae564c8c4f57ac0155d8f1f8b4cbd5d6523c82ea36/markdownify-1.2.3.tar.gz", hash = "sha256:1a176f05522c8a2cb1dd3ab9d307dcdadbed5c26ae717855bfc42b3b6d38d937", size = 18852, upload-time = "2026-06-30T20:27:39.06Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/04/10/fa543d484e8b1199243fe20eedd02cc5af050edebce98a7293a5773df592/markdownify-1.2.3-py3-none-any.whl", hash = "sha256:a189a0bedfd14009030fde5f85bb6f77c56897cb839b5c25315dd7d4e3e290ba", size = 15732, upload-time = "2026-06-30T20:27:38.094Z" },
]
[[package]]
name = "markitdown"
version = "0.1.6"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "beautifulsoup4" },
{ name = "charset-normalizer" },
{ name = "defusedxml" },
{ name = "magika", version = "0.6.2", source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }, marker = "sys_platform == 'win32'" },
{ name = "magika", version = "0.6.3", source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }, marker = "sys_platform != 'win32'" },
{ name = "markdownify" },
{ name = "requests" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/b7/91fe0e2df07107ab701a15c8ad3213135707e4d6206ae9bd8f457a7ad86a/markitdown-0.1.6.tar.gz", hash = "sha256:e5bdbaffd971b29598c7c39ef0e9afce2f08c0751fbfa4e4257678ebaf8cfc7e", size = 50795, upload-time = "2026-05-26T22:43:59.318Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/9e/30/8031f183ee86ea8ac4e7ea1296bab4cca1bee2fd036a26df69764eb7ca74/markitdown-0.1.6-py3-none-any.whl", hash = "sha256:07b2d5bf87b5c53e13a9f2fdc440df8ccc85e26f40c1e557781727b700049775", size = 70032, upload-time = "2026-05-26T22:44:03.209Z" },
]
[[package]]
name = "markupsafe"
version = "3.0.3"
@@ -1104,6 +1191,7 @@ dependencies = [
{ name = "markdown-it-py" },
{ name = "python-dotenv" },
{ name = "pyyaml" },
{ name = "rank-bm25" },
{ name = "sentence-transformers" },
{ name = "typer" },
{ name = "uvicorn", extra = ["standard"] },
@@ -1113,10 +1201,14 @@ dependencies = [
all = [
{ name = "beautifulsoup4" },
{ name = "ebooklib" },
{ name = "markitdown" },
{ name = "openai" },
{ name = "pymupdf" },
{ name = "requests" },
]
bench = [
{ name = "pytest-benchmark" },
]
dev = [
{ name = "httpx" },
{ name = "mypy" },
@@ -1124,6 +1216,9 @@ dev = [
{ name = "pytest-cov" },
{ name = "ruff" },
]
docx = [
{ name = "markitdown" },
]
epub = [
{ name = "ebooklib" },
]
@@ -1142,21 +1237,24 @@ requires-dist = [
{ name = "fastapi", specifier = ">=0.115.0" },
{ name = "httpx", marker = "extra == 'dev'", specifier = ">=0.27.0" },
{ name = "markdown-it-py", specifier = ">=3.0.0" },
{ name = "md-vector-db", extras = ["pdf", "html", "epub"], marker = "extra == 'all'" },
{ name = "markitdown", marker = "extra == 'docx'", specifier = ">=0.1.0" },
{ name = "md-vector-db", extras = ["pdf", "html", "epub", "docx"], marker = "extra == 'all'" },
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.13" },
{ name = "openai", marker = "extra == 'all'", specifier = ">=1.0.0" },
{ name = "pymupdf", marker = "extra == 'pdf'", specifier = ">=1.24.0" },
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
{ name = "pytest-benchmark", marker = "extra == 'bench'", specifier = ">=5.0" },
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5.0" },
{ name = "python-dotenv", specifier = ">=1.2.2" },
{ name = "pyyaml", specifier = ">=6.0" },
{ name = "rank-bm25", specifier = ">=0.2.2" },
{ name = "requests", marker = "extra == 'all'", specifier = ">=2.31.0" },
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.8.0" },
{ name = "sentence-transformers", specifier = ">=3.0.0" },
{ name = "typer", specifier = ">=0.12.0" },
{ name = "uvicorn", extras = ["standard"], specifier = ">=0.30.0" },
]
provides-extras = ["dev", "pdf", "html", "epub", "all"]
provides-extras = ["dev", "pdf", "html", "epub", "docx", "bench", "all"]
[[package]]
name = "mdurl"
@@ -1465,7 +1563,7 @@ name = "nvidia-cudnn-cu12"
version = "9.1.0.70"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "nvidia-cublas-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-cublas-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
]
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741, upload-time = "2024-04-22T15:24:15.253Z" },
@@ -1476,7 +1574,7 @@ name = "nvidia-cufft-cu12"
version = "11.2.1.3"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
]
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117, upload-time = "2024-04-03T20:57:40.402Z" },
@@ -1495,9 +1593,9 @@ name = "nvidia-cusolver-cu12"
version = "11.6.1.9"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "nvidia-cublas-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-cusparse-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-cublas-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-cusparse-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
]
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057, upload-time = "2024-04-03T20:58:28.735Z" },
@@ -1508,7 +1606,7 @@ name = "nvidia-cusparse-cu12"
version = "12.3.1.170"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12", marker = "python_full_version >= '3.14' or sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or sys_platform == 'darwin' or sys_platform == 'linux'" },
]
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763, upload-time = "2024-04-03T20:58:59.995Z" },
@@ -1848,6 +1946,15 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" },
]
[[package]]
name = "py-cpuinfo"
version = "9.0.0"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/37/a8/d832f7293ebb21690860d2e01d8115e5ff6f2ae8bbdc953f0eb0fa4bd2c7/py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690", size = 104716, upload-time = "2022-10-25T20:38:06.303Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" },
]
[[package]]
name = "pybase64"
version = "1.4.3"
@@ -2087,6 +2194,19 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" },
]
[[package]]
name = "pytest-benchmark"
version = "5.2.3"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "py-cpuinfo" },
{ name = "pytest" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/34/9f732b76456d64faffbef6232f1f9dbec7a7c4999ff46282fa418bd1af66/pytest_benchmark-5.2.3.tar.gz", hash = "sha256:deb7317998a23c650fd4ff76e1230066a76cb45dcece0aca5607143c619e7779", size = 341340, upload-time = "2025-11-09T18:48:43.215Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/33/29/e756e715a48959f1c0045342088d7ca9762a2f509b945f362a316e9412b7/pytest_benchmark-5.2.3-py3-none-any.whl", hash = "sha256:bc839726ad20e99aaa0d11a127445457b4219bdb9e80a1afc4b51da7f96b0803", size = 45255, upload-time = "2025-11-09T18:48:39.765Z" },
]
[[package]]
name = "pytest-cov"
version = "7.1.0"
@@ -2158,6 +2278,18 @@ wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
]
[[package]]
name = "rank-bm25"
version = "0.2.2"
source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" }
dependencies = [
{ name = "numpy" },
]
sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fc/0a/f9579384aa017d8b4c15613f86954b92a95a93d641cc849182467cf0bb3b/rank_bm25-0.2.2.tar.gz", hash = "sha256:096ccef76f8188563419aaf384a02f0ea459503fdf77901378d4fd9d87e5e51d", size = 8347, upload-time = "2022-02-16T12:10:52.196Z" }
wheels = [
{ url = "https://pypi.tuna.tsinghua.edu.cn/packages/2a/21/f691fb2613100a62b3fa91e9988c991e9ca5b89ea31c0d3152a3210344f9/rank_bm25-0.2.2-py3-none-any.whl", hash = "sha256:7bd4a95571adadfc271746fa146a4bcfd89c0cf731e49c3d1ad863290adbe8ae", size = 8584, upload-time = "2022-02-16T12:10:50.626Z" },
]
[[package]]
name = "referencing"
version = "0.37.0"
@@ -2662,9 +2794,12 @@ name = "torch"
version = "2.6.0+cu124"
source = { registry = "D:/settings/Language/Python/库" }
resolution-markers = [
"python_full_version >= '3.15'",
"python_full_version == '3.14.*'",
"python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'linux'",
"python_full_version >= '3.15' and sys_platform == 'win32'",
"python_full_version >= '3.15' and sys_platform != 'win32'",
"python_full_version == '3.14.*' and sys_platform == 'win32'",
"python_full_version == '3.14.*' and sys_platform != 'win32'",
"python_full_version < '3.14' and sys_platform == 'win32'",
"python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'",
]
dependencies = [
{ name = "filelock", marker = "(python_full_version >= '3.14' and sys_platform == 'darwin') or (python_full_version >= '3.14' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" },