# 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 全量数据
- **Docker 部署**: 提供 Dockerfile 和 docker-compose.yml，一键部署（含 GPU profile）

## 快速开始

### 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