333 lines
12 KiB
Markdown
333 lines
12 KiB
Markdown
# md-vector-db
|
||
|
||
Markdown 文档向量数据库 — 将 Markdown 文件自动分块、嵌入、存入 ChromaDB，通过语义检索快速查找相关内容。提供 **CLI 命令行工具**和 **HTTP API** 两种使用方式供其他项目集成。
|
||
|
||
## 功能特性
|
||
|
||
- **文档入库**: 支持单文件、目录批量导入多种格式文档，自动按标题+段落智能分块
|
||
- **语义检索**: 自然语言查询，返回最相关的文档片段及来源定位（文件名、章节标题）
|
||
- **多 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` / `.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 文档入库
|
||
|
||
## 快速开始
|
||
|
||
### 1. 安装
|
||
|
||
```bash
|
||
git clone git@lhy-git.liuhangyv.top:Serendipity/md-vector-db.git
|
||
cd md-vector-db
|
||
uv sync
|
||
|
||
# 安装 PDF + HTML 支持（可选）
|
||
uv sync --extra all
|
||
```
|
||
|
||
### 2. 配置
|
||
|
||
```bash
|
||
# 复制环境变量模板（API 密钥等）
|
||
cp .env.example .env
|
||
```
|
||
|
||
编辑 `config.yaml` 选择嵌入模式:
|
||
|
||
```yaml
|
||
embed:
|
||
mode: local # 本地模型（默认，无需联网）
|
||
# 或使用云端 API：
|
||
# mode: api
|
||
# provider: openai # openai | dashscope
|
||
# api_base: https://api.siliconflow.cn/v1 # 可选：覆盖默认 API 地址
|
||
# model: BAAI/bge-large-zh-v1.5 # 可选：覆盖默认模型名
|
||
```
|
||
|
||
### 3. 入库文档
|
||
|
||
```bash
|
||
# 单文件
|
||
uv run md-vector-db ingest docs/intro.md
|
||
|
||
# 批量导入目录（递归扫描所有支持的文档格式）
|
||
uv run md-vector-db ingest-dir ./md_docs/
|
||
|
||
# 指定集合（多项目数据隔离）
|
||
uv run md-vector-db ingest docs/intro.md -C my_project
|
||
|
||
# 查看统计
|
||
uv run md-vector-db stats
|
||
```
|
||
|
||
### 4. 搜索
|
||
|
||
```bash
|
||
# 搜索（注意指定集合，默认为 default）
|
||
uv run md-vector-db search "如何配置向量数据库" -k 5 -C my_project
|
||
|
||
# JSON 输出（供 MCP 等工具消费）
|
||
uv run md-vector-db search "Docker 部署" -k 3 -C obsidian_blog --json
|
||
```
|
||
|
||
输出示例:
|
||
|
||
```
|
||
--- 结果 1 (相似度: 0.8231) ---
|
||
📄 来源: setup-guide.md
|
||
📑 章节: 配置向量数据库
|
||
## 配置向量数据库
|
||
编辑 config.yaml 中的 embed 段来选择嵌入模式...
|
||
|
||
--- 结果 2 (相似度: 0.6104) ---
|
||
📄 来源: faq.md
|
||
📑 章节: 常见问题
|
||
...
|
||
```
|
||
|
||
### 5. 启动 HTTP 服务
|
||
|
||
```bash
|
||
uv run md-vector-db serve --port 8000
|
||
```
|
||
|
||
浏览器打开 `http://localhost:8000/docs` 查看 Swagger UI，可直接在页面中调试所有 API。
|
||
|
||
---
|
||
|
||
## HTTP API
|
||
|
||
| 方法 | 路径 | 需要认证 | 说明 |
|
||
| ---------- | --------------------------------- | -------- | ----------------------------------------------------------------------------- |
|
||
| `GET` | `/` | — | 重定向到`/docs` (Swagger UI) |
|
||
| `GET` | `/api/v1/health` | — | 健康检查（验证 ChromaDB + Embedder 可用性） |
|
||
| `GET` | `/api/v1/collections` | — | 列出集合和已入库的源文件 |
|
||
| `POST` | `/api/v1/ingest` | API Key | 入库文档（`file_path` 或 `content` + `file_name`，可选 `collection`） |
|
||
| `POST` | `/api/v1/search` | API Key | 语义检索（`query` + `top_k`，可选 `collection`） |
|
||
| `DELETE` | `/api/v1/documents/{file_name}` | API Key | 按文件名删除所有关联 chunks |
|
||
|
||
### 调用示例
|
||
|
||
```bash
|
||
# 入库内容
|
||
curl -X POST http://localhost:8000/api/v1/ingest \
|
||
-H "Content-Type: application/json" \
|
||
-H "x-api-key: your-secret-key" \
|
||
-d '{"content": "# 服务配置\n端口设置为 8080。", "file_name": "config.md"}'
|
||
|
||
# 语义搜索
|
||
curl -X POST http://localhost:8000/api/v1/search \
|
||
-H "Content-Type: application/json" \
|
||
-H "x-api-key: your-secret-key" \
|
||
-d '{"query": "端口设置", "top_k": 5}'
|
||
|
||
# 删除文档
|
||
curl -X DELETE http://localhost:8000/api/v1/documents/config.md \
|
||
-H "x-api-key: your-secret-key"
|
||
```
|
||
|
||
### 其他项目集成示例
|
||
|
||
```python
|
||
# Python
|
||
import httpx
|
||
|
||
resp = httpx.post(
|
||
"http://localhost:8000/api/v1/search",
|
||
json={"query": "向量数据库", "top_k": 5},
|
||
headers={"x-api-key": "your-secret-key"},
|
||
)
|
||
for r in resp.json()["results"]:
|
||
print(f"[{r['score']:.2f}] {r['source_file']} — {r['section_title']}")
|
||
```
|
||
|
||
---
|
||
|
||
---
|
||
|
||
## 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 <文件路径> --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） |
|
||
|
||
---
|
||
|
||
## 配置说明
|
||
|
||
### config.yaml
|
||
|
||
```yaml
|
||
chroma:
|
||
persist_dir: ./data # ChromaDB 数据存储目录
|
||
collection_name: markdown_docs # 集合名称
|
||
|
||
embed:
|
||
mode: local # local（本地模型） | api（云端 API）
|
||
local_model: BAAI/bge-small-zh-v1.5 # 本地模型名
|
||
provider: openai # api 模式下的服务商
|
||
api_base: "" # 覆盖默认 API 地址
|
||
model: "" # 覆盖默认模型名
|
||
|
||
chunk:
|
||
max_size: 1000 # 分块最大字符数
|
||
overlap: 100 # 相邻 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 环境变量
|
||
|
||
```bash
|
||
EMBED_API_KEY=sk-xxx # 嵌入 API 密钥（api 模式必需）
|
||
MD_VECTOR_API_KEY=secret # HTTP API 访问密钥（不设置则跳过认证）
|
||
```
|
||
|
||
### Provider 速查
|
||
|
||
| provider | 默认 API 地址 | 默认模型 | 维度 |
|
||
| ------------- | ----------------------------- | -------------------------- | ---- |
|
||
| `local` | 本地（sentence-transformers） | `BAAI/bge-small-zh-v1.5` | 512 |
|
||
| `openai` | `https://api.openai.com/v1` | `text-embedding-3-small` | 1536 |
|
||
| `dashscope` | 阿里云 DashScope | `text-embedding-v4` | 1536 |
|
||
|
||
使用 OpenAI 兼容的第三方服务（硅基流动、智谱、DeepSeek 等）只需设置 `provider: openai` 并覆盖 `api_base` 和 `model`:
|
||
|
||
```yaml
|
||
embed:
|
||
mode: api
|
||
provider: openai
|
||
api_base: https://api.siliconflow.cn/v1
|
||
model: BAAI/bge-large-zh-v1.5
|
||
```
|
||
|
||
---
|
||
|
||
## 项目结构
|
||
|
||
```
|
||
md-vector-db/
|
||
├── config.yaml # 主配置文件
|
||
├── .env.example # 环境变量模板
|
||
├── pyproject.toml
|
||
├── src/
|
||
│ ├── core/ # 核心逻辑（不依赖 server/cli）
|
||
│ │ ├── config.py # YAML + .env → dataclass
|
||
│ │ ├── db.py # ChromaDB 封装（线程安全）
|
||
│ │ ├── embedder.py # 嵌入器（Local/OpenAI/Dashscope）
|
||
│ │ ├── ingest.py # 混合分块 + 入库
|
||
│ │ ├── security.py # 路径遍历防护
|
||
│ │ ├── search.py # 语义检索
|
||
│ │ └── splitters/ # 文档分块器包
|
||
│ │ ├── base.py # Splitter(Protocol) + BaseTextSplitter(ABC)
|
||
│ │ ├── markdown.py # MarkdownSplitter
|
||
│ │ ├── text.py # TextSplitter
|
||
│ │ ├── pdf.py # PDFSplitter
|
||
│ │ ├── html.py # HTMLSplitter
|
||
│ │ ├── epub.py # EPUBSplitter
|
||
│ │ └── registry.py # 扩展名 → Splitter 自动选择
|
||
│ ├── server/ # FastAPI HTTP 层
|
||
│ │ ├── app.py # 路由 + 中间件
|
||
│ │ ├── deps.py # 依赖注入（AppState）
|
||
│ │ └── auth.py # 认证 + 速率限制
|
||
│ └── cli/
|
||
│ └── main.py # Typer CLI
|
||
├── data/ # ChromaDB 持久化目录
|
||
├── md_docs/ # 待入库文档目录
|
||
├── scripts/
|
||
│ ├── serve.py # 快速启动脚本
|
||
│ └── ingest_obsidian.py # 批量入库 Obsidian 知识库
|
||
└── tests/ # 测试（120 个）
|
||
```
|
||
|
||
## 测试
|
||
|
||
```bash
|
||
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 # 覆盖率
|
||
```
|
||
|
||
## 架构
|
||
|
||
**数据流**: 文件 → `get_splitter(path)` 自动选择 → `Splitter.split()` → `batch_embed()` (分批嵌入) → `ChromaDB collection.add()` → `Searcher.search()` 查询
|
||
|
||
**依赖方向**: `config` ← `db` ← `embedder` ← `ingest` / `search` ← `server` / `cli`
|
||
|
||
**分块策略（混合切分）**: 先按 `#`/`##`/`###` 标题拆分 → 章节超过 1000 字符则按段落边界继续拆 → 单个段落仍超限则按标点符号硬切，所有拆分保留 100 字符重叠。
|
||
|
||
**线程安全**: 所有 ChromaDB 写操作（add/delete）通过 `threading.Lock` 保护，支持 FastAPI 多请求并发。
|
||
|
||
---
|
||
|
||
## GPU 加速
|
||
|
||
本地嵌入模式自动检测 CUDA 设备，优先使用 GPU。实测性能对比（RTX 4060 Laptop，bge-small-zh-v1.5）：
|
||
|
||
| 文件大小 | chunks | CPU 耗时 | GPU 耗时 |
|
||
| -------- | ------ | ------------ | -------- |
|
||
| 60KB | 320 | 数分钟至卡死 | 0.7s |
|
||
| 96KB | 729 | 卡死 | 1.8s |
|
||
|
||
**GPU 环境前提**：安装 CUDA 版 torch。本项目 `pyproject.toml` 已配置从本地 wheel 目录获取 CUDA 版 torch，`uv sync` 即可。
|
||
|
||
```toml
|
||
[tool.uv]
|
||
find-links = ["D:/settings/Language/Python/库"] # 本地 CUDA wheel
|
||
index-strategy = "unsafe-best-match" # 允许跨源查找
|
||
```
|
||
|
||
---
|
||
|
||
## License
|
||
|
||
MIT
|