Compare commits
20 Commits
ec3898ecb7
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 2ddb6de92f | |||
| 7e0b9553c8 | |||
| 4710a305be | |||
| 6b89488677 | |||
| 10c5d3d347 | |||
| e71018a9a3 | |||
| d9b03ab5f0 | |||
| 04f2f8a8a8 | |||
| 2d0e8c4997 | |||
| 3e51e166c2 | |||
| 1df4793acb | |||
| 9784f5f436 | |||
| ba814c7a70 | |||
| 82924ff5f3 | |||
| 1d58a55a73 | |||
| cb470e516b | |||
| 7b3c7d5323 | |||
| 3e2c61402a | |||
| 78ad21d311 | |||
| f7801b45a7 |
@@ -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
|
||||
@@ -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
|
||||
@@ -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/
|
||||
@@ -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
|
||||
@@ -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 认证和速率限制
|
||||
- 路径遍历安全防护
|
||||
@@ -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 CLI,5 个命令 + --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 CLI,5 个命令 + --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 CLI,5 个命令 + --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 CLI,5 个命令 + --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)`),但给超长段落测试时需留意类似问题。
|
||||
|
||||
@@ -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
@@ -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"]
|
||||
@@ -1,18 +1,28 @@
|
||||
# md-vector-db
|
||||
|
||||
[](https://github.com/LHY0125/md-vector-db/actions/workflows/ci.yml)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
Markdown 文档向量数据库 — 将 Markdown 文件自动分块、嵌入、存入 ChromaDB,通过语义检索快速查找相关内容。提供 **CLI 命令行工具**和 **HTTP API** 两种使用方式供其他项目集成。
|
||||
|
||||
## 功能特性
|
||||
|
||||
- **文档入库**: 支持单文件、目录批量导入 Markdown 文档,自动按标题+段落智能分块
|
||||
- **文档入库**: 支持单文件、目录批量导入多种格式文档,自动按标题+段落智能分块
|
||||
- **语义检索**: 自然语言查询,返回最相关的文档片段及来源定位(文件名、章节标题)
|
||||
- **多 Provider**: 本地模型 + 云端 API(OpenAI、阿里云 DashScope、硅基流动等 OpenAI 兼容服务)
|
||||
- **GPU 加速**: 本地模型自动检测 CUDA(RTX 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
|
||||
|
||||
@@ -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 重排序
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
```
|
||||
@@ -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
@@ -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"]
|
||||
|
||||
@@ -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
@@ -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
@@ -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", {})),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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 列表
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
# -- 本地模型 --
|
||||
|
||||
@@ -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
@@ -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:
|
||||
"""分块 + 嵌入 + 入库(文本文件通用路径)."""
|
||||
|
||||
@@ -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)
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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_workspace(API 用,绑定当前目录)相比更宽松。
|
||||
"""
|
||||
parts = path_str.replace("\\", "/").split("/")
|
||||
if ".." in parts:
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
@@ -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,5 +1,6 @@
|
||||
"""HTML 文档分块器 — 使用 BeautifulSoup 去标签后委托 TextSplitter."""
|
||||
import logging
|
||||
|
||||
from src.core.splitters.text import TextSplitter
|
||||
|
||||
logger = logging.getLogger("md-vector-db")
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Markdown 文档分块器."""
|
||||
import re
|
||||
|
||||
from src.core.splitters.base import BaseTextSplitter
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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>
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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,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
@@ -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
|
||||
|
||||
@@ -1,10 +1,7 @@
|
||||
"""AppState 和依赖注入测试."""
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
from src.core.config import load_config
|
||||
from src.server.deps import AppState
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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 # 仅表头
|
||||
@@ -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
@@ -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())
|
||||
|
||||
@@ -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"]
|
||||
@@ -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"] == "第一章"
|
||||
@@ -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,5 +1,4 @@
|
||||
"""is_safe_path 路径遍历防护测试."""
|
||||
import pytest
|
||||
from src.core.security import is_safe_path
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)))
|
||||
@@ -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,5 +1,6 @@
|
||||
"""MarkdownSplitter 边界测试."""
|
||||
import pytest
|
||||
|
||||
from src.core.splitters import MarkdownSplitter
|
||||
|
||||
|
||||
|
||||
@@ -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 必须有至少一页才能保存
|
||||
|
||||
@@ -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')" },
|
||||
|
||||
Reference in New Issue
Block a user