Files
md-vector-db/CLAUDE.md
T

8.0 KiB
Raw Blame History

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                       # 全部测试 (115+ 个)
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 CLI5 个命令 + --config 选项

数据流: 文件 → get_splitter(path) 自动选择 → Splitter.split()batch_embed()ChromaDB collection.add()Searcher.search()

支持的格式: .md / .txt / .pdf / .html — 安装可选依赖: uv sync --extra all

依赖方向: configdbembedderingest/searchserver/cli

关键实现细节

GPU 支持 (embedder.py)

LocalEmbedder.__init__ 自动检测 CUDAtorch.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_KEYverify_api_key 依赖注入到 ingest/search/delete 端点;未设置则跳过
  • 速率限制: RateLimiter 中间件,默认 60s 窗口内最多 30 请求
  • 路径遍历防护: _is_safe_path() 拒绝绝对路径和 .. 穿越
  • 错误信息: 500 返回通用消息,详细错误记入 logger.exception

配置系统 (config.py)

config.py 顶部 load_dotenv() 自动加载 .envEMBED_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)),但给超长段落测试时需留意类似问题。