md-vector-db/CLAUDE.md

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

项目概述

Markdown 文档向量数据库 — 将 .md 文件分块 → 嵌入 → 存入 ChromaDB，通过 FastAPI HTTP 或 CLI 提供语义检索。

常用命令

# --- 安装与测试 ---
uv sync                                       # 安装依赖（lockfile 已锁定 CUDA torch）
uv run pytest tests/ -v                       # 全部测试 (46 个)
uv run pytest tests/test_api.py -v            # 单个测试模块
uv run pytest tests/ -v -k "test_search"      # 按名称过滤

# --- CLI 命令（所有命令支持 --config/-c 和 --collection/-C） ---
uv run md-vector-db ingest <file.md>          # 入库单文件（或: python -m src.cli.main ingest）
uv run md-vector-db ingest-dir ./md_docs/     # 入库目录
uv run md-vector-db search "关键词" -k 5      # 搜索 default 集合
uv run md-vector-db search "关键词" -C obsidian_blog -k 5     # 搜索博客知识库
uv run md-vector-db search "..." -C obsidian_blog --json      # JSON 输出（MCP 用）
uv run md-vector-db stats -C obsidian_blog    # 集合统计
uv run md-vector-db serve --port 8000         # 启动 HTTP API

# --- 批量入库 Obsidian ---
uv run python scripts/ingest_obsidian.py

架构

├── core/                         # 核心逻辑（不依赖 server/cli）
│   ├── config.py                 # YAML → dataclass, load_dotenv() 加载 .env
│   ├── db.py                     # VectorDB: 线程安全的 ChromaDB 封装
│   ├── embedder.py               # 策略模式: LocalEmbedder / OpenAIEmbedder / DashscopeEmbedder
│   ├── ingest.py                 # DocumentIngestor: 按扩展名自动选择 Splitter
│   ├── search.py                 # Searcher: 语义检索 + 源文件管理
│   └── splitters/                # 文档分块器包
│       ├── base.py               # Splitter(Protocol) + BaseTextSplitter(ABC)
│       ├── markdown.py           # MarkdownSplitter: 标题+段落混合分块
│       ├── text.py               # TextSplitter: 纯文本段落切分
│       ├── pdf.py                # PDFSplitter: pymupdf 提取文字
│       ├── html.py               # HTMLSplitter: bs4 去标签
│       └── registry.py           # 扩展名 → Splitter 自动选择

server/                       # FastAPI HTTP 层
├── app.py                    # Depends(get_state) 依赖注入, 速率限制中间件
├── deps.py                   # AppState: 集中管理 db/embedder/searcher/ingestor
└── auth.py                   # API Key 认证 (MD_VECTOR_API_KEY) + 速率限制器

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

依赖方向: config ← db ← embedder ← ingest/search ← server/cli

关键实现细节

GPU 支持 (embedder.py)

LocalEmbedder.__init__ 自动检测 CUDA：torch.cuda.is_available() → 优先使用 GPU（RTX 4060），否则 CPU。模型首次加载用 local_files_only=True，未缓存时自动走 hf-mirror.com 镜像下载。

GPU vs CPU 对比（bge-small-zh-v1.5, 96KB 文件 / 729 chunks）：

• CPU: 文件几秒完成，但超 15 chunks 的文件逐渐变慢，320 chunks 以上可能几分钟
• GPU: 729 chunks 嵌入仅 1.8s（含分块+嵌入+ChromaDB 写入）

嵌入 Provider 架构 (embedder.py)

策略模式，接口 Embedder(Protocol):

Provider	类	API 格式	默认模型	向量维度
local	LocalEmbedder	sentence-transformers + GPU	BAAI/bge-small-zh-v1.5	512
openai	OpenAIEmbedder	OpenAI 兼容	text-embedding-3-small	1536
dashscope	DashscopeEmbedder	阿里云自定义 HTTP	text-embedding-v4	1536

config.api_base / config.model 可覆盖默认值。batch_embed() 分批嵌入（每批 32 条）防 OOM。

线程安全 (db.py + ingest.py + search.py)

VectorDB._write_lock (threading.Lock) 保护所有写操作（collection.add / delete / get+delete）。

认证与安全 (auth.py + app.py)

• API Key: 环境变量 MD_VECTOR_API_KEY → verify_api_key 依赖注入到 ingest/search/delete 端点；未设置则跳过
• 速率限制: RateLimiter 中间件，默认 60s 窗口内最多 30 请求
• 路径遍历防护: _is_safe_path() 拒绝绝对路径和 .. 穿越
• 错误信息: 500 返回通用消息，详细错误记入 logger.exception

配置系统 (config.py)

config.py 顶部 load_dotenv() 自动加载 .env。EMBED_API_KEY 从环境变量读取。DEFAULT_CONFIG_PATH = "config.yaml" 统一引用。

服务层 (deps.py + app.py)

AppState 类集中管理 db/embedder/searcher/ingestor 单例。端点通过 FastAPI Depends(get_state) 获取，可测试、可替换。is_healthy() 真实验证 ChromaDB 和 Embedder 可用性。

配置

config.yaml + .env 联合控制。复制 .env.example 为 .env 填入密钥。

HTTP API

方法	路径	认证	说明
GET	/	-	重定向到 /docs
GET	/api/v1/health	-	真实健康检查 (ChromaDB + Embedder)
GET	/api/v1/collections	-	集合和源文件列表
POST	/api/v1/ingest	API Key	入库 (file_path 或 content+file_name)
POST	/api/v1/search	API Key	语义检索 (query + top_k 1-100)
DELETE	/api/v1/documents/{file_name}	API Key	按文件名删除

已入库知识库

集合名	来源	文件数	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

搜索时务必用 -C 指定集合，否则只会搜到 default 中的测试数据。

小说搜索示例：

uv run md-vector-db search "张莽和孙太后的关系" -k 3 -C novel_taohou
uv run md-vector-db search "抄家事件" -k 5 -C novel_taohou
uv run md-vector-db search "文谦变法" -k 3 -C novel_taohou

已知问题 / 注意事项

uv 与 CUDA torch 的兼容配置

本机全局 UV_INDEX_URL 指向阿里云镜像（只有 CPU 版 torch），pyproject.toml 通过以下配置让 uv 从本地 wheel 取 CUDA 版：

[tool.uv]
find-links = ["D:/settings/Language/Python/库"]   # 本地 CUDA wheel 目录
index-strategy = "unsafe-best-match"              # 允许跨源查找

uv.lock 已锁定：Windows 平台 → torch 2.6.0+cu124（本地），Linux/Mac → CPU torch（清华源）。 不要删除 [tool.uv] 配置，否则 uv sync 会重新解析为 CPU 版 torch。

MarkdownSplitter 边界情况

_split_single_paragraph 中，当段落分隔符（。！？等）距 chunk 起点 < overlap(100) 时， start 会回退为负数，Python str.rfind 的负索引会绕回文本末尾，造成死循环。 此 bug 已被修复（start = max(start + 1, next_start)），但给超长段落测试时需留意类似问题。