# md-vector-db 实现计划 > **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. **Goal:** 构建一个 Markdown 文档向量数据库,支持文档入库、嵌入、语义检索,并通过 FastAPI HTTP 服务供其他项目调用。 **Architecture:** 分层设计 — `core` 层封装 ChromaDB + 嵌入模型逻辑,`server` 层提供 FastAPI HTTP 接口,`cli` 层提供命令行工具。core 不依赖 server/cli,两者各自调用 core。 **Tech Stack:** Python 3.13, ChromaDB, sentence-transformers (BAAI/bge-small-zh-v1.5), FastAPI + uvicorn, Typer, PyYAML, markdown-it-py --- ## 文件结构 | 文件 | 职责 | | -------------------------- | -------------------------------------------------------- | | `pyproject.toml` | 项目元数据和依赖声明 | | `config.yaml` | 用户配置文件(嵌入模式、分块参数、服务端口等) | | `src/core/config.py` | 配置加载,将 YAML 解析为 dataclass 对象 | | `src/core/db.py` | ChromaDB PersistentClient 初始化,collection 管理 | | `src/core/embedder.py` | 嵌入模型抽象(本地/API 双模),EmbedConfig → Embedder | | `src/core/ingest.py` | Markdown 解析、标题提取、混合分块、入库、去重 | | `src/core/search.py` | 语义检索,query 嵌入→ChromaDB 相似度搜索→格式化结果 | | `src/server/app.py` | FastAPI 路由:ingest/search/collections/documents/health | | `src/cli/main.py` | Typer CLI:ingest/ingest-dir/search/serve/stats | | `scripts/serve.py` | 一行启动脚本:`uvicorn src.server.app:app` | | `tests/test_config.py` | 配置加载测试 | | `tests/test_embedder.py` | 嵌入模型测试(本地模式核心验证) | | `tests/test_ingest.py` | 分块逻辑与入库测试 | | `tests/test_search.py` | 检索接口测试 | | `tests/test_api.py` | FastAPI 端点集成测试 | --- ### Task 1: 项目脚手架 **Files:** - Create: `pyproject.toml` - Create: `config.yaml` - Create: `src/__init__.py` - Create: `src/core/__init__.py` - Create: `src/server/__init__.py` - Create: `src/cli/__init__.py` - Create: `tests/__init__.py` - Create: `scripts/__init__.py` (空,不创建也行) - Create: `.gitignore` - [ ] **Step 1: 创建 pyproject.toml** ```toml [project] name = "md-vector-db" version = "0.1.0" description = "Markdown 文档向量数据库,支持语义检索" requires-python = ">=3.13" dependencies = [ "chromadb>=0.5.0", "sentence-transformers>=3.0.0", "fastapi>=0.115.0", "uvicorn[standard]>=0.30.0", "pyyaml>=6.0", "markdown-it-py>=3.0.0", "typer>=0.12.0", ] [project.optional-dependencies] dev = [ "pytest>=8.0", "httpx>=0.27.0", ] [build-system] requires = ["hatchling"] build-backend = "hatchling.build" [tool.pytest.ini_options] testpaths = ["tests"] pythonpath = ["src"] ``` - [ ] **Step 2: 创建 config.yaml** ```yaml chroma: persist_dir: ./data collection_name: markdown_docs embed: mode: local # local | api local_model: BAAI/bge-small-zh-v1.5 api_base: "" # api 模式下填写 api_key: "" # api 模式下填写 chunk: max_size: 1000 # 分块最大字符数 overlap: 100 # 相邻块重叠字符数 server: host: 0.0.0.0 port: 8000 ``` - [ ] **Step 3: 创建所有 __init__.py(空文件)** 创建以下空文件: - `src/__init__.py` - `src/core/__init__.py` - `src/server/__init__.py` - `src/cli/__init__.py` - `tests/__init__.py` - [ ] **Step 4: 创建 .gitignore** ```gitignore __pycache__/ *.pyc data/ .env *.egg-info/ .pytest_cache/ ``` - [ ] **Step 5: 安装依赖** ```bash cd D:/Code/doing_exercises/programs/md-vector-db uv sync ``` - [ ] **Step 6: Commit** ```bash git init git add -A git commit -m "chore: scaffold md-vector-db project structure" ``` --- ### Task 2: 配置加载模块 (config.py) **Files:** - Create: `src/core/config.py` - Create: `tests/test_config.py` - [ ] **Step 1: 写测试 — tests/test_config.py** ```python """配置加载模块测试.""" import tempfile from pathlib import Path import yaml import pytest from src.core.config import AppConfig, EmbedConfig, ChunkConfig, load_config class TestEmbedConfig: """嵌入配置解析.""" def test_local_mode_defaults(self): """默认 local 模式,带默认模型名.""" data = {"embed": {"mode": "local"}} cfg = AppConfig(**data) assert cfg.embed.mode == "local" assert cfg.embed.local_model == "BAAI/bge-small-zh-v1.5" assert cfg.embed.api_base == "" def test_api_mode_fields(self): """api 模式下 api_base 和 api_key 可设置.""" data = { "embed": { "mode": "api", "api_base": "https://api.openai.com/v1", "api_key": "sk-test", } } cfg = AppConfig(**data) assert cfg.embed.mode == "api" assert cfg.embed.api_base == "https://api.openai.com/v1" assert cfg.embed.api_key == "sk-test" class TestChunkConfig: """分块配置解析.""" def test_default_values(self): """分块默认值正确.""" data = {"chunk": {}} cfg = AppConfig(**data) assert cfg.chunk.max_size == 1000 assert cfg.chunk.overlap == 100 class TestLoadConfig: """load_config 函数测试.""" def test_load_from_yaml_file(self): """从 YAML 文件加载配置.""" yaml_content = """ chroma: persist_dir: /tmp/test_data collection_name: test_collection embed: mode: local chunk: max_size: 500 server: port: 9000 """ with tempfile.NamedTemporaryFile( mode="w", suffix=".yaml", delete=False ) as f: f.write(yaml_content) tmp_path = f.name try: cfg = load_config(tmp_path) assert cfg.chroma.persist_dir == "/tmp/test_data" assert cfg.chroma.collection_name == "test_collection" assert cfg.chunk.max_size == 500 assert cfg.server.port == 9000 finally: Path(tmp_path).unlink() def test_load_missing_file_uses_defaults(self): """配置文件不存在时使用默认值.""" cfg = load_config("/nonexistent/config.yaml") assert cfg.chroma.persist_dir == "./data" assert cfg.embed.mode == "local" ``` - [ ] **Step 2: 运行测试验证失败** ```bash cd D:/Code/doing_exercises/programs/md-vector-db uv run pytest tests/test_config.py -v ``` Expected: FAIL (module not found) - [ ] **Step 3: 实现 config.py** ```python """应用配置加载模块.""" from dataclasses import dataclass, field from pathlib import Path import yaml @dataclass class ChromaConfig: """ChromaDB 配置.""" persist_dir: str = "./data" collection_name: str = "markdown_docs" @dataclass class EmbedConfig: """嵌入模型配置.""" mode: str = "local" # "local" | "api" local_model: str = "BAAI/bge-small-zh-v1.5" api_base: str = "" api_key: str = "" @dataclass class ChunkConfig: """文档分块配置.""" max_size: int = 1000 overlap: int = 100 @dataclass class ServerConfig: """HTTP 服务配置.""" host: str = "0.0.0.0" port: int = 8000 @dataclass class AppConfig: """应用总配置.""" chroma: ChromaConfig = field(default_factory=ChromaConfig) embed: EmbedConfig = field(default_factory=EmbedConfig) chunk: ChunkConfig = field(default_factory=ChunkConfig) server: ServerConfig = field(default_factory=ServerConfig) def __init__(self, **kwargs): self.chroma = ChromaConfig(**kwargs.get("chroma", {})) self.embed = EmbedConfig(**kwargs.get("embed", {})) self.chunk = ChunkConfig(**kwargs.get("chunk", {})) self.server = ServerConfig(**kwargs.get("server", {})) def load_config(path: str | None = None) -> AppConfig: """从 YAML 文件加载配置, 若文件不存在则返回默认配置.""" config_path = path or "config.yaml" if not Path(config_path).exists(): return AppConfig() with open(config_path, "r", encoding="utf-8") as f: data = yaml.safe_load(f) or {} return AppConfig(**data) ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_config.py -v ``` Expected: PASS - [ ] **Step 5: Commit** ```bash git add src/core/config.py tests/test_config.py git commit -m "feat: add config loading module with YAML support" ``` --- ### Task 3: 数据库层 (db.py) **Files:** - Create: `src/core/db.py` - Create: `tests/test_db.py` - [ ] **Step 1: 写测试 — tests/test_db.py** ```python """数据库层测试.""" import tempfile from pathlib import Path import pytest from src.core.db import VectorDB class TestVectorDB: """VectorDB 测试.""" @pytest.fixture def temp_dir(self): with tempfile.TemporaryDirectory() as d: yield d def test_init_creates_persist_directory(self, temp_dir): """初始化时自动创建持久化目录.""" db = VectorDB(persist_dir=temp_dir + "/subdir") assert Path(temp_dir + "/subdir").exists() assert db.client is not None def test_get_or_create_collection(self, temp_dir): """创建或获取 collection.""" db = VectorDB(persist_dir=temp_dir) col = db.get_or_create_collection("test_col") assert col.name == "test_col" # 再次获取应返回同一个 col2 = db.get_or_create_collection("test_col") assert col2.name == "test_col" def test_count_returns_zero_for_empty_collection(self, temp_dir): """空 collection 文档数为 0.""" db = VectorDB(persist_dir=temp_dir) col = db.get_or_create_collection("test_col") assert col.count() == 0 def test_delete_collection(self, temp_dir): """删除 collection.""" db = VectorDB(persist_dir=temp_dir) db.get_or_create_collection("tmp_col") db.delete_collection("tmp_col") # 再次获取会创建新的 col = db.get_or_create_collection("tmp_col") assert col.count() == 0 ``` - [ ] **Step 2: 运行测试验证失败** ```bash uv run pytest tests/test_db.py -v ``` Expected: FAIL - [ ] **Step 3: 实现 db.py** ```python """ChromaDB 数据库层.""" import chromadb from chromadb.api.models.Collection import Collection class VectorDB: """向量数据库封装.""" def __init__(self, persist_dir: str = "./data"): self.client = chromadb.PersistentClient(path=persist_dir) def get_or_create_collection(self, name: str) -> Collection: """获取或创建 collection.""" return self.client.get_or_create_collection(name=name) def delete_collection(self, name: str) -> None: """删除 collection.""" try: self.client.delete_collection(name=name) except ValueError: pass # collection 不存在则忽略 ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_db.py -v ``` Expected: PASS - [ ] **Step 5: Commit** ```bash git add src/core/db.py tests/test_db.py git commit -m "feat: add ChromaDB database layer" ``` --- ### Task 4: 嵌入层 (embedder.py) **Files:** - Create: `src/core/embedder.py` - Create: `tests/test_embedder.py` - [ ] **Step 1: 写测试 — tests/test_embedder.py** ```python """嵌入模型测试.""" import pytest from src.core.config import EmbedConfig from src.core.embedder import Embedder, create_embedder class TestEmbedder: """Embedder 单元测试. 注意:本地模型测试需要下载 sentence-transformers 模型(约 100MB), 首次运行耗时较长。API 模式测试使用 mock 避免网络依赖。 """ @pytest.fixture def local_config(self): return EmbedConfig(mode="local") def test_create_local_embedder(self, local_config): """创建本地嵌入器,验证维度正确.""" embedder = create_embedder(local_config) assert embedder.dimension > 0 assert isinstance(embedder.dimension, int) def test_embed_single_text(self, local_config): """嵌入单条文本返回正确维度向量.""" embedder = create_embedder(local_config) result = embedder.embed(["你好世界"]) assert len(result) == 1 assert len(result[0]) == embedder.dimension assert all(isinstance(v, float) for v in result[0]) def test_embed_multiple_texts(self, local_config): """嵌入多条文本返回对应数量的向量.""" embedder = create_embedder(local_config) texts = ["第一段文本", "第二段文本", "第三段文本"] result = embedder.embed(texts) assert len(result) == 3 for vec in result: assert len(vec) == embedder.dimension def test_embed_empty_list_raises(self, local_config): """空列表应抛出异常.""" embedder = create_embedder(local_config) with pytest.raises(ValueError): embedder.embed([]) def test_create_embedder_from_factory_function(self, local_config): """工厂函数正确创建 Embedder 实例.""" embedder = create_embedder(local_config) assert isinstance(embedder, Embedder) ``` - [ ] **Step 2: 运行测试验证失败** ```bash uv run pytest tests/test_embedder.py -v ``` Expected: FAIL - [ ] **Step 3: 实现 embedder.py** ```python """嵌入模型抽象层.""" from sentence_transformers import SentenceTransformer from src.core.config import EmbedConfig class Embedder: """文本嵌入器, 支持本地模型和 API 两种模式.""" def __init__(self, config: EmbedConfig): self._config = config if config.mode == "local": self._model = SentenceTransformer(config.local_model) elif config.mode == "api": self._model = None # 延迟初始化, 需要 openai 包 self._api_base = config.api_base self._api_key = config.api_key else: raise ValueError(f"不支持的嵌入模式: {config.mode}") @property def mode(self) -> str: """当前嵌入模式.""" return self._config.mode @property def dimension(self) -> int: """嵌入向量维度.""" if self.mode == "local": return self._model.get_sentence_embedding_dimension() else: # 默认 OpenAI text-embedding-ada-002 / text-embedding-3-small 维度 return 1536 def embed(self, texts: list[str]) -> list[list[float]]: """对文本列表进行嵌入, 返回向量列表.""" if not texts: raise ValueError("文本列表不能为空") if self.mode == "local": embeddings = self._model.encode(texts, normalize_embeddings=True) return embeddings.tolist() else: return self._embed_via_api(texts) def _embed_via_api(self, texts: list[str]) -> list[list[float]]: """通过 OpenAI 兼容 API 嵌入(延迟导入 openai).""" from openai import OpenAI client = OpenAI(base_url=self._api_base, api_key=self._api_key) response = client.embeddings.create( model="text-embedding-3-small", input=texts, ) return [d.embedding for d in response.data] def create_embedder(config: EmbedConfig) -> Embedder: """工厂函数: 根据配置创建嵌入器.""" return Embedder(config) ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_embedder.py -v ``` Expected: PASS (首次运行会下载模型) - [ ] **Step 5: Commit** ```bash git add src/core/embedder.py tests/test_embedder.py git commit -m "feat: add embedder with local/API dual mode" ``` --- ### Task 5: 文档入库模块 (ingest.py) **Files:** - Create: `src/core/ingest.py` - Create: `tests/test_ingest.py` - [ ] **Step 1: 写测试 — tests/test_ingest.py** ```python """文档入库测试.""" import tempfile from pathlib import Path import pytest from src.core.ingest import MarkdownSplitter, DocumentIngestor class TestMarkdownSplitter: """Markdown 分块器测试.""" @pytest.fixture def splitter(self): return MarkdownSplitter(max_size=1000, overlap=100) def test_split_simple_document(self, splitter): """简单文档按标题拆分.""" md = """# 标题一 这是第一段内容。 ## 标题二 这是第二段内容。 # 标题三 这是第三段内容。""" chunks = splitter.split(md, source_file="test.md") assert len(chunks) >= 3 # 每个 chunk 有元数据 for chunk in chunks: assert "content" in chunk assert chunk["source_file"] == "test.md" def test_chunk_has_heading_metadata(self, splitter): """chunk 附带标题元数据.""" md = "# 配置指南\n这里是配置说明。" chunks = splitter.split(md, source_file="config.md") assert len(chunks) >= 1 title = chunks[0]["section_title"] assert "配置指南" in title or title == "" def test_long_section_is_split(self, splitter): """超长章节被进一步拆分.""" # 创建一个超过 max_size 的段落 long_text = "这是很长的文本。" * 300 # ~3000 字符 md = f"# 长章节\n{long_text}" small_splitter = MarkdownSplitter(max_size=500, overlap=50) chunks = small_splitter.split(md, source_file="long.md") assert len(chunks) > 1 def test_empty_document(self, splitter): """空文档返回空列表.""" chunks = splitter.split("", source_file="empty.md") assert chunks == [] def test_code_blocks_preserved(self, splitter): """代码块不被拆分.""" md = """# 代码示例 ```python def hello(): print("world") ``` """ chunks = splitter.split(md, source_file="code.md") assert len(chunks) >= 1 # 代码块内容应在某个 chunk 中 all_content = " ".join(c["content"] for c in chunks) assert "def hello()" in all_content class TestDocumentIngestor: """文档入库器测试.""" @pytest.fixture def temp_md_dir(self): with tempfile.TemporaryDirectory() as d: # 创建测试 Markdown 文件 md_path = Path(d) / "test.md" md_path.write_text("# 测试\n这是测试内容。", encoding="utf-8") yield d def test_read_markdown_file(self, temp_md_dir): """读取 Markdown 文件.""" ingestor = DocumentIngestor.__new__(DocumentIngestor) content = Path(temp_md_dir + "/test.md").read_text(encoding="utf-8") assert "测试" in content assert "这是测试内容" in content ``` - [ ] **Step 2: 运行测试验证失败** ```bash uv run pytest tests/test_ingest.py -v ``` Expected: FAIL - [ ] **Step 3: 实现 ingest.py** ```python """Markdown 文档解析与入库模块.""" import re from dataclasses import dataclass, field from pathlib import Path from src.core.db import VectorDB from src.core.embedder import Embedder @dataclass class Chunk: """文档分块.""" content: str source_file: str section_title: str = "" heading_level: int = 0 chunk_index: int = 0 def to_metadata(self) -> dict: """转为 ChromaDB 元数据.""" return { "source_file": self.source_file, "section_title": self.section_title, "heading_level": self.heading_level, "chunk_index": self.chunk_index, } class MarkdownSplitter: """Markdown 混合分块器:先按标题拆,超长再按段落拆.""" def __init__(self, max_size: int = 1000, overlap: int = 100): self.max_size = max_size self.overlap = overlap def split(self, text: str, source_file: str = "") -> list[dict]: """将 Markdown 文本拆分为带元数据的 chunk 列表.""" if not text.strip(): return [] sections = self._split_by_headings(text) chunks = [] for section in sections: if len(section["content"]) <= self.max_size: chunks.append(section) else: sub_chunks = self._split_by_paragraphs( section["content"], section["section_title"], section["heading_level"], ) chunks.extend(sub_chunks) # 为所有 chunk 补充 source_file 和 chunk_index for i, chunk in enumerate(chunks): chunk["source_file"] = source_file or chunk.get("source_file", "") chunk["chunk_index"] = i return chunks def _split_by_headings(self, text: str) -> list[dict]: """按 Markdown 标题拆分.""" # 匹配行首的 # 标题 heading_pattern = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE) matches = list(heading_pattern.finditer(text)) if not matches: return [{ "content": text.strip(), "section_title": "", "heading_level": 0, }] sections = [] for i, match in enumerate(matches): level = len(match.group(1)) title = match.group(2).strip() start = match.end() end = matches[i + 1].start() if i + 1 < len(matches) else len(text) content = text[start:end].strip() if content: sections.append({ "content": f"{match.group(0)}\n{content}", "section_title": title, "heading_level": level, }) # 处理第一个标题之前的内容 if matches and matches[0].start() > 0: preamble = text[:matches[0].start()].strip() if preamble: sections.insert(0, { "content": preamble, "section_title": "", "heading_level": 0, }) return sections def _split_by_paragraphs( self, text: str, section_title: str, heading_level: int ) -> list[dict]: """按段落边界拆分超长章节.""" paragraphs = re.split(r"\n\n+", text) chunks = [] current = "" count = 0 for para in paragraphs: if len(current) + len(para) > self.max_size and current: chunks.append({ "content": current.strip(), "section_title": section_title, "heading_level": heading_level, }) # overlap: 保留上一块的末尾部分 if self.overlap > 0 and len(current) > self.overlap: current = current[-self.overlap:] + "\n\n" + para else: current = para else: if current: current += "\n\n" + para else: current = para if current.strip(): chunks.append({ "content": current.strip(), "section_title": section_title, "heading_level": heading_level, }) return chunks class DocumentIngestor: """文档入库器: 读取 MD 文件 → 分块 → 嵌入 → 入库.""" def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str): self.db = db self.embedder = embedder self.collection_name = collection_name self.splitter = MarkdownSplitter() @property def collection(self): return self.db.get_or_create_collection(self.collection_name) def ingest_file(self, file_path: str) -> int: """入库单个 Markdown 文件, 返回 chunk 数量.""" path = Path(file_path) content = path.read_text(encoding="utf-8") file_name = path.name return self.ingest_content(content, file_name) def ingest_content(self, content: str, file_name: str) -> int: """入库 Markdown 内容(无需实际文件).""" # 去重:先删旧 chunks self._remove_by_source(file_name) # 分块 chunks = self.splitter.split(content, source_file=file_name) if not chunks: return 0 # 嵌入 texts = [c["content"] for c in chunks] embeddings = self.embedder.embed(texts) # 入库 ids = [f"{file_name}_{i}" for i in range(len(chunks))] metadatas = [ { "source_file": c.get("source_file", file_name), "section_title": c.get("section_title", ""), "heading_level": c.get("heading_level", 0), "chunk_index": i, } for i, c in enumerate(chunks) ] self.collection.add( ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas, ) return len(chunks) def ingest_directory(self, dir_path: str) -> dict[str, int]: """入库目录下所有 Markdown 文件.""" results = {} for md_file in Path(dir_path).rglob("*.md"): count = self.ingest_file(str(md_file)) results[md_file.name] = count return results def _remove_by_source(self, file_name: str) -> None: """按 source_file 删除已有 chunks.""" try: existing = self.collection.get( where={"source_file": file_name} ) if existing and existing["ids"]: self.collection.delete(ids=existing["ids"]) except Exception: pass # collection 为空时 get 可能抛异常 ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_ingest.py -v ``` Expected: PASS - [ ] **Step 5: Commit** ```bash git add src/core/ingest.py tests/test_ingest.py git commit -m "feat: add Markdown splitter and document ingestor" ``` --- ### Task 6: 检索模块 (search.py) **Files:** - Create: `src/core/search.py` - Create: `tests/test_search.py` - [ ] **Step 1: 写测试 — tests/test_search.py** ```python """检索模块测试.""" import tempfile from pathlib import Path import pytest from src.core.config import EmbedConfig from src.core.db import VectorDB from src.core.embedder import create_embedder from src.core.ingest import DocumentIngestor from src.core.search import Searcher @pytest.fixture def searcher(): """创建带测试数据的 Searcher.""" tmpdir = tempfile.mkdtemp() db = VectorDB(persist_dir=tmpdir) embedder = create_embedder(EmbedConfig(mode="local")) ingestor = DocumentIngestor(db, embedder, "test_search") # 入库一些测试文档 content = """# Python 入门 Python 是一种解释型编程语言。 ## 安装 Python 从 python.org 下载安装包。 # 向量数据库 ChromaDB 是一个轻量级向量数据库。 ## ChromaDB 安装 使用 pip install chromadb 安装。""" ingestor.ingest_content(content, "guide.md") return Searcher(db, embedder, "test_search") class TestSearcher: """检索器测试.""" def test_search_returns_results(self, searcher): """搜索返回至少一条结果.""" results = searcher.search("Python 编程", top_k=3) assert len(results) > 0 for r in results: assert r["content"] assert r["source_file"] assert "score" in r def test_search_scores_are_descending(self, searcher): """搜索结果按相似度降序排列.""" results = searcher.search("向量数据库", top_k=5) scores = [r["score"] for r in results] assert scores == sorted(scores) def test_search_respects_top_k(self, searcher): """top_k 参数限制返回数量.""" results = searcher.search("安装", top_k=2) assert len(results) <= 2 def test_search_returns_all_fields(self, searcher): """搜索结果包含完整字段.""" results = searcher.search("ChromaDB", top_k=1) if results: r = results[0] assert "content" in r assert "source_file" in r assert "section_title" in r assert "heading_level" in r assert "chunk_index" in r assert "score" in r def test_search_no_results(self, searcher): """无语义匹配时不崩溃.""" results = searcher.search("xyzxyz不存在的内容abcabc", top_k=3) assert isinstance(results, list) # ChromaDB 总是返回最近的向量, 所以即使不匹配也会返回结果, 只是 score 低 # 这里只验证不抛异常 ``` - [ ] **Step 2: 运行测试验证失败** ```bash uv run pytest tests/test_search.py -v ``` Expected: FAIL - [ ] **Step 3: 实现 search.py** ```python """语义检索模块.""" from src.core.db import VectorDB from src.core.embedder import Embedder class Searcher: """向量检索器.""" def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str): self.db = db self.embedder = embedder self.collection_name = collection_name @property def collection(self): return self.db.get_or_create_collection(self.collection_name) def search( self, query: str, top_k: int = 10, source_file: str | None = None, ) -> list[dict]: """语义检索, 返回格式化结果列表.""" query_embedding = self.embedder.embed([query])[0] where_filter = None if source_file: where_filter = {"source_file": source_file} results = self.collection.query( query_embeddings=[query_embedding], n_results=top_k, where=where_filter, include=["documents", "metadatas", "distances"], ) formatted = [] if results["ids"] and results["ids"][0]: for i, doc_id in enumerate(results["ids"][0]): metadata = results["metadatas"][0][i] if results["metadatas"] else {} distance = results["distances"][0][i] if results["distances"] else 0.0 # ChromaDB 默认用余弦距离, 转为相似度分数 (0~1) score = round(1.0 - distance, 4) formatted.append({ "id": doc_id, "content": results["documents"][0][i] if results["documents"] else "", "source_file": metadata.get("source_file", ""), "section_title": metadata.get("section_title", ""), "heading_level": metadata.get("heading_level", 0), "chunk_index": metadata.get("chunk_index", 0), "score": max(0.0, score), }) return formatted def get_collection_info(self) -> dict: """获取 collection 信息.""" return { "name": self.collection_name, "count": self.collection.count(), } def list_sources(self) -> list[str]: """列出所有已入库的源文件.""" if self.collection.count() == 0: return [] result = self.collection.get(include=["metadatas"]) sources = set() if result and result["metadatas"]: for m in result["metadatas"]: if m and "source_file" in m: sources.add(m["source_file"]) return sorted(sources) def delete_by_source(self, file_name: str) -> bool: """按文件名删除文档.""" try: existing = self.collection.get( where={"source_file": file_name} ) if existing and existing["ids"]: self.collection.delete(ids=existing["ids"]) return True except Exception: pass return False ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_search.py -v ``` Expected: PASS - [ ] **Step 5: Commit** ```bash git add src/core/search.py tests/test_search.py git commit -m "feat: add semantic search module" ``` --- ### Task 7: FastAPI 服务层 (server/app.py) **Files:** - Create: `src/server/app.py` - Create: `tests/test_api.py` - [ ] **Step 1: 写测试 — tests/test_api.py** ```python """API 端点集成测试.""" import tempfile from pathlib import Path import pytest from fastapi.testclient import TestClient @pytest.fixture def app_client(monkeypatch): """创建测试客户端(使用临时目录).""" tmpdir = tempfile.mkdtemp() # 在导入 app 前设置环境 monkeypatch.setattr( "src.server.app._get_config", lambda: type( "C", (), { "chroma": type("Ch", (), {"persist_dir": tmpdir, "collection_name": "test_api"})(), "embed": type("Em", (), {"mode": "local"})(), "chunk": type("Ck", (), {"max_size": 1000, "overlap": 100})(), "server": type("Sv", (), {"host": "0.0.0.0", "port": 8000})(), }, )(), ) from src.server.app import app return TestClient(app) class TestHealthEndpoint: """健康检查.""" def test_health_returns_ok(self, app_client): response = app_client.get("/api/v1/health") assert response.status_code == 200 data = response.json() assert data["status"] == "ok" class TestCollectionEndpoint: """Collection 端点.""" def test_list_collections(self, app_client): response = app_client.get("/api/v1/collections") assert response.status_code == 200 data = response.json() assert "collections" in data class TestSearchEndpoint: """搜索端点.""" def test_search_requires_query(self, app_client): response = app_client.post("/api/v1/search", json={}) assert response.status_code == 422 # 缺少必填字段 def test_search_empty_collection(self, app_client): response = app_client.post( "/api/v1/search", json={"query": "test", "top_k": 5}, ) assert response.status_code == 200 data = response.json() assert "results" in data assert data["results"] == [] class TestIngestEndpoint: """入库端点.""" def test_ingest_content(self, app_client): response = app_client.post( "/api/v1/ingest", json={ "content": "# Test\nHello world.", "file_name": "test.md", }, ) assert response.status_code == 200 data = response.json() assert data["status"] == "ok" assert data["chunks"] > 0 def test_ingest_then_search(self, app_client): """入库后能检索到.""" # 先入库 app_client.post( "/api/v1/ingest", json={"content": "# 配置说明\nChromaDB 配置很简单。", "file_name": "config.md"}, ) # 再搜索 response = app_client.post( "/api/v1/search", json={"query": "配置", "top_k": 3}, ) assert response.status_code == 200 data = response.json() assert len(data["results"]) > 0 ``` - [ ] **Step 2: 运行测试验证失败** ```bash uv run pytest tests/test_api.py -v ``` Expected: FAIL - [ ] **Step 3: 实现 server/app.py** ```python """FastAPI 服务层.""" import os from pathlib import Path from fastapi import FastAPI, HTTPException from pydantic import BaseModel from src.core.config import load_config, AppConfig from src.core.db import VectorDB from src.core.embedder import create_embedder from src.core.ingest import DocumentIngestor from src.core.search import Searcher # -- 模型 -- class IngestRequest(BaseModel): file_path: str | None = None content: str | None = None file_name: str | None = None class SearchRequest(BaseModel): query: str top_k: int = 10 # -- 懒加载单例 -- _config: AppConfig | None = None _db: VectorDB | None = None _embedder = None _searcher: Searcher | None = None _ingestor: DocumentIngestor | None = None def _get_config(): global _config if _config is None: _config = load_config(os.getenv("MD_VECTOR_CONFIG", "config.yaml")) return _config def _get_db(): global _db if _db is None: cfg = _get_config() _db = VectorDB(persist_dir=cfg.chroma.persist_dir) return _db def _get_embedder(): global _embedder if _embedder is None: cfg = _get_config() _embedder = create_embedder(cfg.embed) return _embedder def _init_services(): global _searcher, _ingestor cfg = _get_config() db = _get_db() embedder = _get_embedder() _searcher = Searcher(db, embedder, cfg.chroma.collection_name) _ingestor = DocumentIngestor(db, embedder, cfg.chroma.collection_name) def _get_searcher(): if _searcher is None: _init_services() return _searcher def _get_ingestor(): if _ingestor is None: _init_services() return _ingestor # -- App -- app = FastAPI( title="md-vector-db", description="Markdown 文档向量数据库 API", version="0.1.0", ) @app.on_event("startup") def startup(): _init_services() @app.get("/api/v1/health") def health(): return {"status": "ok"} @app.get("/api/v1/collections") def list_collections(): searcher = _get_searcher() info = searcher.get_collection_info() sources = searcher.list_sources() return {"collections": [info], "sources": sources} @app.post("/api/v1/ingest") def ingest_document(req: IngestRequest): ingestor = _get_ingestor() try: if req.file_path: path = Path(req.file_path) if not path.exists(): raise HTTPException(status_code=404, detail=f"文件不存在: {req.file_path}") count = ingestor.ingest_file(str(path)) file_name = path.name elif req.content: file_name = req.file_name or "untitled.md" count = ingestor.ingest_content(req.content, file_name) else: raise HTTPException( status_code=400, detail="需要提供 file_path 或 content" ) return {"status": "ok", "chunks": count, "file": file_name} except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/api/v1/search") def search_documents(req: SearchRequest): searcher = _get_searcher() results = searcher.search(req.query, top_k=req.top_k) return {"results": results} @app.delete("/api/v1/documents/{file_name}") def delete_document(file_name: str): searcher = _get_searcher() deleted = searcher.delete_by_source(file_name) if not deleted: raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}") return {"status": "ok", "file": file_name} ``` - [ ] **Step 4: 运行测试验证通过** ```bash uv run pytest tests/test_api.py -v ``` Expected: PASS - [ ] **Step 5: Commit** ```bash git add src/server/app.py tests/test_api.py git commit -m "feat: add FastAPI HTTP API layer" ``` --- ### Task 8: CLI 工具 (cli/main.py) **Files:** - Create: `src/cli/main.py` - [ ] **Step 1: 实现 CLI** ```python """命令行工具入口.""" import sys from pathlib import Path import typer import uvicorn # 添加 src 到路径 sys.path.insert(0, str(Path(__file__).parent.parent)) from src.core.config import load_config from src.core.db import VectorDB from src.core.embedder import create_embedder from src.core.ingest import DocumentIngestor from src.core.search import Searcher app = typer.Typer(name="md-vector-db", help="Markdown 文档向量数据库管理工具") def _get_components(config_path: str = "config.yaml"): """初始化所有组件.""" cfg = load_config(config_path) db = VectorDB(persist_dir=cfg.chroma.persist_dir) embedder = create_embedder(cfg.embed) searcher = Searcher(db, embedder, cfg.chroma.collection_name) ingestor = DocumentIngestor(db, embedder, cfg.chroma.collection_name) return cfg, searcher, ingestor @app.command() def ingest(file_path: str): """入库单个 Markdown 文件.""" _, _, ingestor = _get_components() count = ingestor.ingest_file(file_path) typer.echo(f"✅ 已入库: {file_path} ({count} 个 chunks)") @app.command() def ingest_dir(dir_path: str): """入库目录下所有 Markdown 文件.""" _, _, ingestor = _get_components() results = ingestor.ingest_directory(dir_path) total = sum(results.values()) for name, count in results.items(): typer.echo(f" 📄 {name}: {count} chunks") typer.echo(f"✅ 共入库 {len(results)} 个文件, {total} 个 chunks") @app.command() def search(query: str, top_k: int = 10): """语义检索.""" _, searcher, _ = _get_components() results = searcher.search(query, top_k=top_k) if not results: typer.echo("未找到匹配结果。") return for i, r in enumerate(results, 1): typer.echo(f"\n--- 结果 {i} (相似度: {r['score']:.4f}) ---") typer.echo(f"📄 来源: {r['source_file']}") if r["section_title"]: typer.echo(f"📑 章节: {r['section_title']}") preview = r["content"][:200] + "..." if len(r["content"]) > 200 else r["content"] typer.echo(preview) @app.command() def serve(port: int = 8000): """启动 HTTP 服务.""" typer.echo(f"🚀 启动服务: http://0.0.0.0:{port}") uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False) @app.command() def stats(): """查看统计信息.""" _, searcher, _ = _get_components() info = searcher.get_collection_info() sources = searcher.list_sources() typer.echo(f"📊 Collection: {info['name']}") typer.echo(f"📦 总 chunks: {info['count']}") typer.echo(f"📄 源文件数: {len(sources)}") if sources: typer.echo("\n源文件列表:") for s in sources: typer.echo(f" - {s}") if __name__ == "__main__": app() ``` - [ ] **Step 2: 验证 CLI 可用** ```bash cd D:/Code/doing_exercises/programs/md-vector-db uv run python -m src.cli.main --help ``` Expected: 显示帮助信息,列出 ingest/ingest-dir/search/serve/stats 命令 - [ ] **Step 3: Commit** ```bash git add src/cli/main.py git commit -m "feat: add CLI management tool" ``` --- ### Task 9: 启动脚本与入口完善 (scripts/serve.py + pyproject.toml 更新) **Files:** - Create: `scripts/serve.py` - Modify: `pyproject.toml` - [ ] **Step 1: 创建 scripts/serve.py** ```python """便捷启动脚本.""" import uvicorn if __name__ == "__main__": uvicorn.run("src.server.app:app", host="0.0.0.0", port=8000, reload=True) ``` - [ ] **Step 2: 更新 pyproject.toml 添加 scripts 入口** 在 `pyproject.toml` 的 `[project]` 段末尾添加: ```toml [project.scripts] md-vector-db = "src.cli.main:app" ``` - [ ] **Step 3: 验证 pip install 后可命令行调用** ```bash uv pip install -e . md-vector-db --help ``` - [ ] **Step 4: Commit** ```bash git add scripts/serve.py pyproject.toml git commit -m "feat: add serve script and console entry point" ``` --- ### Task 10: 端到端验证与文档 **Files:** - Create: `README.md` - [ ] **Step 1: 创建测试用的 Markdown 文件** ```bash mkdir -p md_docs cat > md_docs/test-guide.md << 'EOF' # 向量数据库入门指南 ## 什么是向量数据库 向量数据库是一种专门用于存储和检索向量嵌入的数据库。 ## 为什么需要向量数据库 传统的数据库只能做精确匹配,而向量数据库可以做语义相似度搜索。 ## 常用的向量数据库 - ChromaDB:轻量级,适合小项目 - Qdrant:高性能,适合生产环境 - Milvus:分布式,适合大规模数据 ## ChromaDB 快速上手 安装 ChromaDB: ```bash pip install chromadb ``` 创建 collection 并添加数据: ```python import chromadb client = chromadb.PersistentClient(path="./data") collection = client.get_or_create_collection("my_docs") collection.add( documents=["这是第一篇文档", "这是第二篇文档"], ids=["doc1", "doc2"], ) ``` EOF ``` - [ ] **Step 2: 端到端测试** ```bash # 入库测试文档 uv run python -m src.cli.main ingest md_docs/test-guide.md # 命令行搜索 uv run python -m src.cli.main search "如何安装ChromaDB" --top-k 3 # 查看统计 uv run python -m src.cli.main stats ``` Expected: 搜索返回相关结果,统计显示 1 个源文件 - [ ] **Step 3: 运行全部测试验证覆盖率** ```bash uv run pytest tests/ -v --tb=short ``` Expected: 所有测试 PASS - [ ] **Step 4: 创建 README.md** ```markdown # md-vector-db Markdown 文档向量数据库,支持文档入库、语义检索,可通过 HTTP API 供其他项目调用。 ## 快速开始 ### 安装 \`\`\`bash cd md-vector-db uv sync \`\`\` ### CLI 使用 \`\`\`bash # 入库文档 uv run python -m src.cli.main ingest path/to/file.md # 批量入库 uv run python -m src.cli.main ingest-dir ./md_docs/ # 搜索 uv run python -m src.cli.main search "关键词" --top-k 10 # 启动 HTTP 服务 uv run python -m src.cli.main serve --port 8000 \`\`\` ### HTTP API \`\`\`bash # 启动服务 uv run python -m src.cli.main serve # 入库 curl -X POST http://localhost:8000/api/v1/ingest \\ -H "Content-Type: application/json" \\ -d '{"content": "# 标题\\n内容...", "file_name": "doc.md"}' # 搜索 curl -X POST http://localhost:8000/api/v1/search \\ -H "Content-Type: application/json" \\ -d '{"query": "关键词", "top_k": 5}' \`\`\` ## 配置 编辑 `config.yaml` 切换嵌入模式、分块参数等。 ## 项目结构 - `src/core/` — 核心逻辑(数据库、嵌入、分块、检索) - `src/server/` — FastAPI HTTP 服务 - `src/cli/` — 命令行工具 - `tests/` — 测试 \`\`\` - [ ] **Step 5: 最终 Commit** ```bash git add README.md md_docs/ git commit -m "docs: add README and test data" ```