7e77c31e16
- README: 补充 GPU 加速章节、多集合、CLI 新选项 (-C/--json) - CLAUDE: 修正 uv run 可用说明、补充 [tool.uv] 配置文档 - pyproject.toml: 添加 [tool.uv] 配置, 通过 find-links + unsafe-best-match 让 uv 从本地 CUDA wheel 解析 torch, 不再覆盖 GPU 环境 - uv.lock: 更新 lockfile, Windows 平台锁定 torch 2.6.0+cu124
281 lines
9.1 KiB
Markdown
281 lines
9.1 KiB
Markdown
# md-vector-db
|
||
|
||
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 认证、速率限制、路径遍历防护
|
||
- **去重**: 同一文件重复入库自动覆盖旧版本(基于路径 SHA256 哈希)
|
||
|
||
## 快速开始
|
||
|
||
### 1. 安装
|
||
|
||
```bash
|
||
git clone git@lhy-git.liuhangyv.top:Serendipity/md-vector-db.git
|
||
cd md-vector-db
|
||
uv sync
|
||
```
|
||
|
||
### 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
|
||
|
||
# 批量导入目录(递归扫描所有 .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']}")
|
||
```
|
||
|
||
---
|
||
|
||
## CLI 命令参考
|
||
|
||
所有命令均支持 `--config/-c`(配置文件)、`--collection/-C`(集合名,默认 `default`)。
|
||
|
||
| 命令 | 说明 |
|
||
|------|------|
|
||
| `ingest <文件路径>` | 入库单个 .md 文件,支持 `-C` 指定集合 |
|
||
| `ingest-dir <目录路径>` | 递归入库目录下所有 .md 文件 |
|
||
| `search <查询> -k <数量>` | 语义检索,`-k` 默认 10、最大 100,`--json` JSON 输出 |
|
||
| `stats` | 显示 chunks 总数、源文件列表 |
|
||
| `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 # 服务监听端口
|
||
```
|
||
|
||
### .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 # 混合分块 + 入库
|
||
│ │ └── search.py # 语义检索
|
||
│ ├── 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/ # 测试(46 个)
|
||
```
|
||
|
||
## 测试
|
||
|
||
```bash
|
||
uv run pytest tests/ -v # 全部测试 (46 个)
|
||
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 # 覆盖率
|
||
```
|
||
|
||
## 架构
|
||
|
||
**数据流**: MD 文件 → `MarkdownSplitter.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
|