diff --git a/CLAUDE.md b/CLAUDE.md index 7b47a55..3462a95 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,14 +4,14 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## 项目概述 -Markdown 文档向量数据库 — 将 .md 文件分块 → 嵌入 → 存入 ChromaDB，通过 FastAPI HTTP 或 CLI 提供语义检索。 +文档向量数据库 — 将 .md/.txt/.pdf/.html/.epub 文件分块 → 嵌入 → 存入 ChromaDB，通过 FastAPI HTTP 或 CLI 提供语义检索。 ## 常用命令 -```bash +```Shell # --- 安装与测试 --- uv sync # 安装依赖（lockfile 已锁定 CUDA torch） -uv run pytest tests/ -v # 全部测试 (115+ 个) +uv run pytest tests/ -v # 全部测试 (120 个) uv run pytest tests/test_api.py -v # 单个测试模块 uv run pytest tests/ -v -k "test_search" # 按名称过滤 @@ -36,6 +36,7 @@ uv run python scripts/ingest_obsidian.py │ ├── db.py # VectorDB: 线程安全的 ChromaDB 封装 │ ├── embedder.py # 策略模式: LocalEmbedder / OpenAIEmbedder / DashscopeEmbedder │ ├── ingest.py # DocumentIngestor: 按扩展名自动选择 Splitter +│ ├── security.py # 路径遍历防护 (is_safe_path + is_path_within_workspace) │ ├── search.py # Searcher: 语义检索 + 源文件管理 │ └── splitters/ # 文档分块器包 │ ├── base.py # Splitter(Protocol) + BaseTextSplitter(ABC) @@ -43,6 +44,7 @@ uv run python scripts/ingest_obsidian.py │ ├── text.py # TextSplitter: 纯文本段落切分 │ ├── pdf.py # PDFSplitter: pymupdf 提取文字 │ ├── html.py # HTMLSplitter: bs4 去标签 +│ ├── epub.py # EPUBSplitter: ebooklib 提取章节 │ └── registry.py # 扩展名 → Splitter 自动选择 server/ # FastAPI HTTP 层 @@ -55,7 +57,7 @@ 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` +**支持的格式**: `.md` / `.txt` / `.pdf` / `.html` / `.epub` — 安装可选依赖: `uv sync --extra all` **依赖方向**: `config` ← `db` ← `embedder` ← `ingest`/`search` ← `server`/`cli` @@ -90,7 +92,7 @@ cli/main.py # Typer CLI，5 个命令 + --config 选项 - **API Key**: 环境变量 `MD_VECTOR_API_KEY` → `verify_api_key` 依赖注入到 ingest/search/delete 端点；未设置则跳过 - **速率限制**: `RateLimiter` 中间件，默认 60s 窗口内最多 30 请求 -- **路径遍历防护**: `_is_safe_path()` 拒绝绝对路径和 `..` 穿越 +- **路径遍历防护**: `is_path_within_workspace()` 拒绝绝对路径、`..` 穿越和目录外访问 - **错误信息**: 500 返回通用消息，详细错误记入 `logger.exception` ### 配置系统 (config.py) @@ -118,15 +120,16 @@ cli/main.py # Typer CLI，5 个命令 + --config 选项 ## 已入库知识库 -| 集合名 | 来源 | 文件数 | 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 | +| 集合名 | 来源 | 文件数 | 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 中的测试数据。 **小说搜索示例**： + ```bash uv run md-vector-db search "张莽和孙太后的关系" -k 3 -C novel_taohou uv run md-vector-db search "抄家事件" -k 5 -C novel_taohou @@ -149,6 +152,7 @@ index-strategy = "unsafe-best-match" # 允许跨源查找 **不要删除 `[tool.uv]` 配置**，否则 `uv sync` 会重新解析为 CPU 版 torch。 ### MarkdownSplitter 边界情况 + `_split_single_paragraph` 中，当段落分隔符（。！？等）距 chunk 起点 < overlap(100) 时， `start` 会回退为负数，Python `str.rfind` 的负索引会绕回文本末尾，造成死循环。 此 bug 已被修复（`start = max(start + 1, next_start)`），但给超长段落测试时需留意类似问题。 diff --git a/README.md b/README.md index f2bc35b..371b95a 100644 --- a/README.md +++ b/README.md @@ -4,14 +4,14 @@ Markdown 文档向量数据库 — 将 Markdown 文件自动分块、嵌入、 ## 功能特性 -- **文档入库**: 支持单文件、目录批量导入 Markdown 文档，自动按标题+段落智能分块 +- **文档入库**: 支持单文件、目录批量导入多种格式文档，自动按标题+段落智能分块 - **语义检索**: 自然语言查询，返回最相关的文档片段及来源定位（文件名、章节标题） - **多 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`，按扩展名自动选择分块器，可通过 `Splitter` Protocol 扩展 +- **多格式文档**: 支持 `.md` / `.txt` / `.pdf` / `.html` / `.epub`，按扩展名自动选择分块器，可通过 `Splitter` Protocol 扩展 - **去重**: 同一文件重复入库自动覆盖旧版本（基于路径 SHA256 哈希） ## 快速开始 @@ -52,7 +52,7 @@ embed: # 单文件 uv run md-vector-db ingest docs/intro.md -# 批量导入目录（递归扫描所有 .md 文件） +# 批量导入目录（递归扫描所有支持的文档格式） uv run md-vector-db ingest-dir ./md_docs/ # 指定集合（多项目数据隔离） @@ -151,8 +151,8 @@ for r in resp.json()["results"]: | 命令 | 说明 | | --------------------------- | -------------------------------------------------------- | -| `ingest <文件路径>` | 入库单个 .md 文件，支持`-C` 指定集合 | -| `ingest-dir <目录路径>` | 递归入库目录下所有 .md 文件 | +| `ingest <文件路径>` | 入库单个文件（自动识别格式），支持`-C` 指定集合 | +| `ingest-dir <目录路径>` | 递归入库目录下所有支持的文档格式 | | `search <查询> -k <数量>` | 语义检索，`-k` 默认 10、最大 100，`--json` JSON 输出 | | `stats` | 显示 chunks 总数、源文件列表 | | `serve -p <端口>` | 启动 HTTP 服务（默认 8000） | @@ -224,6 +224,7 @@ md-vector-db/ │ │ ├── db.py # ChromaDB 封装（线程安全） │ │ ├── embedder.py # 嵌入器（Local/OpenAI/Dashscope） │ │ ├── ingest.py # 混合分块 + 入库 +│ │ ├── security.py # 路径遍历防护 │ │ ├── search.py # 语义检索 │ │ └── splitters/ # 文档分块器包 │ │ ├── base.py # Splitter(Protocol) + BaseTextSplitter(ABC) @@ -242,15 +243,15 @@ md-vector-db/ ├── data/ # ChromaDB 持久化目录 ├── md_docs/ # 待入库文档目录 ├── scripts/ -│ ├── serve.py # 快速启动脚本 -│ └── ingest_obsidian.py # 批量入库 Obsidian 知识库 -└── tests/ # 测试（115+ 个） +│ ├── serve.py # 快速启动脚本 +│ └── ingest_obsidian.py # 批量入库 Obsidian 知识库 +└── tests/ # 测试（120 个） ``` ## 测试 ```bash -uv run pytest tests/ -v # 全部测试 (115+ 个) +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 # 覆盖率 diff --git a/src/cli/main.py b/src/cli/main.py index 4ab0fa3..2562d1e 100644 --- a/src/cli/main.py +++ b/src/cli/main.py @@ -17,7 +17,7 @@ if sys.stdout.encoding != "utf-8": sys.path.insert(0, str(Path(__file__).parent.parent)) from src.core.config import DEFAULT_CONFIG_PATH -from src.core.security import is_path_within_workspace +from src.core.security import is_safe_cli_path from src.server.deps import get_state, get_default_collection app = typer.Typer( @@ -82,7 +82,7 @@ def ingest( if file_paths: total = 0 for fp in file_paths: - if not is_path_within_workspace(fp): + if not is_safe_cli_path(fp): typer.echo(f"[SKIP] 不安全的路径: {fp}", err=True) continue # 支持通配符 (shell 展开或 Python glob) @@ -90,7 +90,7 @@ def ingest( if "*" in fp or "?" in fp: matches = _glob.glob(fp, recursive=True) for m in matches: - if not is_path_within_workspace(m): + if not is_safe_cli_path(m): typer.echo(f"[SKIP] 不安全的路径: {m}", err=True) continue c = ingestor.ingest_file(m) @@ -117,7 +117,7 @@ def ingest_dir( collection: CollectionOpt = None, ): _init_config(config) - if not is_path_within_workspace(dir_path): + if not is_safe_cli_path(dir_path): typer.echo(f"错误: 不安全的路径 — {dir_path}", err=True) raise typer.Exit(code=1) state = get_state() diff --git a/src/core/security.py b/src/core/security.py index 22aaec0..5d01168 100644 --- a/src/core/security.py +++ b/src/core/security.py @@ -46,3 +46,15 @@ def is_path_within_workspace(path_str: str) -> bool: except ValueError: return False return common == cwd + + +def is_safe_cli_path(path_str: str) -> bool: + """CLI 路径安全检查 — 仅拒绝 .. 穿越组件，允许绝对路径和任意目录。 + + CLI 是本地工具，用户有权限访问系统中任意路径。 + 与 is_path_within_workspace（API 用，绑定当前目录）相比更宽松。 + """ + parts = path_str.replace("\\", "/").split("/") + if ".." in parts: + return False + return True