fix: remove api_key from config.yaml, read from EMBED_API_KEY env var; unify config path constant
This commit is contained in:
@@ -0,0 +1,889 @@
|
||||
# 审计问题修复计划
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) to implement this plan task-by-task.
|
||||
|
||||
**Goal:** 修复架构审计报告中 Phase 1+2 的全部严重和高风险问题，以及部分中风险问题。
|
||||
|
||||
**Architecture:** 按影响面从底层到上层依次修复：config → embedder → db → ingest → search → server → cli。每层修复后运行测试验证不引入回归。
|
||||
|
||||
**Tech Stack:** Python 3.13, FastAPI, ChromaDB, sentence-transformers
|
||||
|
||||
---
|
||||
|
||||
## 修复范围
|
||||
|
||||
| 审计编号 | 问题 | 本次修复 |
|
||||
|----------|------|----------|
|
||||
| S-01 | API 密钥明文 | ✅ |
|
||||
| S-02 | 路径遍历 | ✅ |
|
||||
| S-03 | 无认证 | ✅ |
|
||||
| S-05 | 错误信息泄露 | ✅ |
|
||||
| C-01 | 并发不安全 | ✅ |
|
||||
| C-02 | 全局环境变量污染 | ✅ |
|
||||
| C-04 | top_k 无上界 | ✅ |
|
||||
| C-05 | 大文件 OOM | ✅ |
|
||||
| C-06 | 配置路径硬编码 | ✅ |
|
||||
| C-07 | 假健康检查 | ✅ |
|
||||
| D-02 | 全局单例 | ✅ |
|
||||
| D-03 | 同步嵌入阻塞 | ✅ |
|
||||
| M-02 | 无日志 | ✅ |
|
||||
| T-02 | DELETE 端点无测试 | ✅ |
|
||||
|
||||
---
|
||||
|
||||
### Task 1: 配置安全修复 (S-01, C-06)
|
||||
|
||||
**Files:**
|
||||
- Modify: `config.yaml` — 移除 `api_key` 字段
|
||||
- Modify: `src/core/config.py` — `api_key` 默认从环境变量读取
|
||||
- Modify: `src/cli/main.py:14` — 移除重复 sys.path 操作（顺带）
|
||||
- Add Test: `tests/test_config.py` — 补充 API key 环境变量测试
|
||||
|
||||
- [ ] **Step 1: 补充测试 — tests/test_config.py 追加**
|
||||
|
||||
```python
|
||||
class TestEmbedConfigEnvVar:
|
||||
"""api_key 从环境变量读取."""
|
||||
|
||||
def test_api_key_from_env(self, monkeypatch):
|
||||
monkeypatch.setenv("EMBED_API_KEY", "sk-env-test")
|
||||
from src.core.config import EmbedConfig
|
||||
cfg = EmbedConfig(mode="api")
|
||||
assert cfg.api_key == "sk-env-test"
|
||||
|
||||
def test_api_key_empty_when_not_set(self, monkeypatch):
|
||||
monkeypatch.delenv("EMBED_API_KEY", raising=False)
|
||||
from src.core.config import EmbedConfig
|
||||
cfg = EmbedConfig(mode="api")
|
||||
assert cfg.api_key == ""
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行测试确认失败**
|
||||
|
||||
```bash
|
||||
cd D:/Code/doing_exercises/programs/md-vector-db && uv run pytest tests/test_config.py -v
|
||||
```
|
||||
Expected: 新测试 FAIL
|
||||
|
||||
- [ ] **Step 3: 修改 config.yaml — 删除 api_key 行**
|
||||
|
||||
```yaml
|
||||
embed:
|
||||
mode: local # local | api
|
||||
local_model: BAAI/bge-small-zh-v1.5
|
||||
api_base: "" # api 模式下填写
|
||||
# api_key 请通过环境变量 EMBED_API_KEY 设置
|
||||
```
|
||||
|
||||
- [ ] **Step 4: 修改 config.py — api_key 默认从环境变量读取**
|
||||
|
||||
将 `EmbedConfig` 的 `api_key` 字段改为:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class EmbedConfig:
|
||||
"""嵌入模型配置."""
|
||||
mode: str = "local"
|
||||
local_model: str = "BAAI/bge-small-zh-v1.5"
|
||||
api_base: str = ""
|
||||
api_key: str = field(default_factory=lambda: os.environ.get("EMBED_API_KEY", ""))
|
||||
```
|
||||
|
||||
并在文件顶部加 `import os`。
|
||||
|
||||
- [ ] **Step 5: 统一配置路径常量**
|
||||
|
||||
在 `config.py` 末尾添加:
|
||||
```python
|
||||
DEFAULT_CONFIG_PATH = "config.yaml"
|
||||
```
|
||||
|
||||
其他文件中的 `"config.yaml"` 字符串引用改为 `from src.core.config import DEFAULT_CONFIG_PATH`。
|
||||
|
||||
修改 `src/server/app.py:47,59` 和 `src/cli/main.py:26` 中的硬编码 `"config.yaml"`。
|
||||
|
||||
- [ ] **Step 6: 运行全部测试**
|
||||
|
||||
```bash
|
||||
cd D:/Code/doing_exercises/programs/md-vector-db && uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
Expected: 33 passed (31 + 2 new)
|
||||
|
||||
- [ ] **Step 7: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "fix: remove api_key from config.yaml, read from env var; unify config path"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: 路径遍历防护 (S-02)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/server/app.py` — `ingest_document()` 增加路径白名单
|
||||
- Add Test: `tests/test_api.py` — 路径遍历攻击测试
|
||||
|
||||
- [ ] **Step 1: 补充测试 — tests/test_api.py 追加**
|
||||
|
||||
```python
|
||||
class TestSecurity:
|
||||
"""安全测试."""
|
||||
|
||||
def test_ingest_rejects_path_traversal(self, client):
|
||||
"""拒绝路径遍历攻击."""
|
||||
response = client.post(
|
||||
"/api/v1/ingest",
|
||||
json={"file_path": "../../../etc/passwd", "file_name": "hack.md"},
|
||||
)
|
||||
assert response.status_code in (400, 403)
|
||||
|
||||
def test_ingest_rejects_absolute_path(self, client):
|
||||
"""拒绝绝对路径."""
|
||||
response = client.post(
|
||||
"/api/v1/ingest",
|
||||
json={"file_path": "C:\\Windows\\System32\\config\\SAM"},
|
||||
)
|
||||
assert response.status_code in (400, 403)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行测试确认失败**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_api.py::TestSecurity -v
|
||||
```
|
||||
Expected: FAIL (当前无防护)
|
||||
|
||||
- [ ] **Step 3: 修改 app.py — 路径校验函数**
|
||||
|
||||
在 `ingest_document()` 最前面加:
|
||||
|
||||
```python
|
||||
import os as _os
|
||||
|
||||
def _is_safe_path(path_str: str) -> bool:
|
||||
"""检查路径是否安全: 仅允许相对路径且不含 .. 穿越."""
|
||||
normalized = _os.path.normpath(path_str)
|
||||
if _os.path.isabs(normalized):
|
||||
return False
|
||||
if ".." in normalized.split(_os.sep):
|
||||
return False
|
||||
return True
|
||||
|
||||
@app.post("/api/v1/ingest")
|
||||
def ingest_document(req: IngestRequest):
|
||||
if req.file_path:
|
||||
if not _is_safe_path(req.file_path):
|
||||
raise HTTPException(status_code=400, detail="不允许的路径")
|
||||
path = Path(req.file_path)
|
||||
...
|
||||
```
|
||||
|
||||
- [ ] **Step 4: 运行测试确认通过**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_api.py -v
|
||||
```
|
||||
Expected: 8 passed (6 original + 2 new security tests)
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "fix: add path traversal protection to ingest endpoint"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: API 认证 + 速率限制 + 安全错误消息 (S-03, S-04, S-05)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/server/app.py` — 添加 API Key 中间件、速率限制、安全错误消息
|
||||
- Create: `src/server/auth.py` — 认证依赖
|
||||
- Add Test: `tests/test_api.py` — 认证/速率测试
|
||||
|
||||
- [ ] **Step 1: 创建 auth.py**
|
||||
|
||||
```python
|
||||
"""API 认证与安全中间件."""
|
||||
import os
|
||||
from fastapi import Header, HTTPException, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
import time
|
||||
from collections import defaultdict
|
||||
import threading
|
||||
|
||||
|
||||
# -- API Key 认证 --
|
||||
EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "")
|
||||
|
||||
|
||||
def verify_api_key(x_api_key: str | None = Header(None)):
|
||||
"""验证 API Key. 若未设置环境变量则跳过验证."""
|
||||
if EXPECTED_API_KEY and x_api_key != EXPECTED_API_KEY:
|
||||
raise HTTPException(status_code=401, detail="无效的 API Key")
|
||||
return True
|
||||
|
||||
|
||||
# -- 简易速率限制 --
|
||||
class RateLimiter:
|
||||
"""基于内存的简易速率限制器."""
|
||||
|
||||
def __init__(self, max_requests: int = 30, window_seconds: int = 60):
|
||||
self.max_requests = max_requests
|
||||
self.window = window_seconds
|
||||
self._store: dict[str, list[float]] = defaultdict(list)
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def is_allowed(self, client_id: str) -> bool:
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
records = self._store[client_id]
|
||||
# 清理过期记录
|
||||
records[:] = [t for t in records if now - t < self.window]
|
||||
if len(records) >= self.max_requests:
|
||||
return False
|
||||
records.append(now)
|
||||
return True
|
||||
|
||||
async def __call__(self, request: Request):
|
||||
client_id = request.client.host if request.client else "unknown"
|
||||
if not self.is_allowed(client_id):
|
||||
raise HTTPException(status_code=429, detail="请求过于频繁，请稍后再试")
|
||||
return True
|
||||
|
||||
|
||||
rate_limiter = RateLimiter(max_requests=30, window_seconds=60)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 修改 app.py — 集成认证和速率限制**
|
||||
|
||||
在 `app.py` 中:
|
||||
|
||||
```python
|
||||
from src.server.auth import verify_api_key, rate_limiter
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger("md-vector-db")
|
||||
|
||||
# 速率限制中间件
|
||||
@app.middleware("http")
|
||||
async def rate_limit_middleware(request: Request, call_next):
|
||||
await rate_limiter(request)
|
||||
response = await call_next(request)
|
||||
return response
|
||||
|
||||
|
||||
@app.post("/api/v1/ingest")
|
||||
def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)):
|
||||
...
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.exception("入库失败: %s", e)
|
||||
raise HTTPException(status_code=500, detail="服务器内部错误")
|
||||
|
||||
|
||||
@app.post("/api/v1/search")
|
||||
def search_documents(req: SearchRequest, _: bool = Depends(verify_api_key)):
|
||||
...
|
||||
|
||||
|
||||
@app.delete("/api/v1/documents/{file_name}")
|
||||
def delete_document(file_name: str, _: bool = Depends(verify_api_key)):
|
||||
...
|
||||
```
|
||||
|
||||
- [ ] **Step 3: 运行测试确认通过**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_api.py -v
|
||||
```
|
||||
Expected: 全部通过 (未设 API Key 环境变量时认证跳过)
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "fix: add API key auth, rate limiting, and safe error messages"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: 请求校验 (C-04, M-03)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/server/app.py` — `SearchRequest.top_k` 加 Pydantic Field 校验；`IngestRequest.file_name` 校验
|
||||
|
||||
- [ ] **Step 1: 修改请求模型**
|
||||
|
||||
```python
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
class IngestRequest(BaseModel):
|
||||
file_path: str | None = None
|
||||
content: str | None = None
|
||||
file_name: str | None = Field(default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$")
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
query: str = Field(..., min_length=1, max_length=2000)
|
||||
top_k: int = Field(default=10, ge=1, le=100)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行测试**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_api.py -v
|
||||
```
|
||||
Expected: 全部通过
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "fix: add field validation to API request models"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: ChromaDB 线程安全 (C-01)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/db.py` — VectorDB 写操作加锁
|
||||
|
||||
- [ ] **Step 1: 修改 db.py**
|
||||
|
||||
```python
|
||||
"""ChromaDB 数据库层."""
|
||||
import threading
|
||||
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)
|
||||
self._write_lock = threading.Lock()
|
||||
|
||||
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 (线程安全)."""
|
||||
with self._write_lock:
|
||||
try:
|
||||
self.client.delete_collection(name=name)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
def close(self) -> None:
|
||||
"""释放数据库连接."""
|
||||
self.client.close()
|
||||
|
||||
@property
|
||||
def write_lock(self) -> threading.Lock:
|
||||
"""获取写锁，供外部在 add/delete/update 操作时使用."""
|
||||
return self._write_lock
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 修改 ingest.py — 使用写锁**
|
||||
|
||||
```python
|
||||
class DocumentIngestor:
|
||||
...
|
||||
def ingest_content(self, content: str, file_name: str) -> int:
|
||||
...
|
||||
with self.db.write_lock:
|
||||
self.collection.add(
|
||||
ids=ids, embeddings=embeddings,
|
||||
documents=texts, metadatas=metadatas,
|
||||
)
|
||||
return len(chunks)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: 修改 search.py — delete 使用写锁**
|
||||
|
||||
```python
|
||||
def delete_by_source(self, file_name: str) -> bool:
|
||||
...
|
||||
with self.db.write_lock:
|
||||
existing = self.collection.get(where={"source_file": file_name})
|
||||
if existing and existing["ids"]:
|
||||
self.collection.delete(ids=existing["ids"])
|
||||
return True
|
||||
return False
|
||||
```
|
||||
|
||||
- [ ] **Step 4: 运行测试**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
Expected: 31 passed
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "fix: add thread-safe write lock to ChromaDB operations"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Embedder 重构 (C-02, C-05, D-01)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/embedder.py` — 拆分为策略模式；移除全局 env var 副作用；加分批嵌入
|
||||
- Modify: `tests/test_embedder.py` — 补充 API 模式测试
|
||||
|
||||
- [ ] **Step 1: 重写 embedder.py**
|
||||
|
||||
```python
|
||||
"""嵌入模型抽象层 — 策略模式."""
|
||||
import os
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
from src.core.config import EmbedConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 模型下载镜像
|
||||
HF_MIRROR = "https://hf-mirror.com"
|
||||
|
||||
|
||||
class Embedder(Protocol):
|
||||
"""嵌入器接口."""
|
||||
@property
|
||||
def dimension(self) -> int: ...
|
||||
def embed(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
class LocalEmbedder:
|
||||
"""本地 sentence-transformers 模型嵌入器."""
|
||||
|
||||
def __init__(self, config: EmbedConfig):
|
||||
from sentence_transformers import SentenceTransformer
|
||||
self._config = config
|
||||
try:
|
||||
self._model = SentenceTransformer(
|
||||
config.local_model, local_files_only=True
|
||||
)
|
||||
except Exception:
|
||||
logger.info("模型未缓存, 通过镜像下载 %s", config.local_model)
|
||||
os.environ["HF_ENDPOINT"] = HF_MIRROR
|
||||
self._model = SentenceTransformer(config.local_model)
|
||||
os.environ.pop("HF_ENDPOINT", None)
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return self._model.get_embedding_dimension()
|
||||
|
||||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
raise ValueError("文本列表不能为空")
|
||||
embeddings = self._model.encode(texts, normalize_embeddings=True)
|
||||
return embeddings.tolist()
|
||||
|
||||
|
||||
class APIEmbedder:
|
||||
"""OpenAI 兼容 API 嵌入器."""
|
||||
|
||||
def __init__(self, config: EmbedConfig):
|
||||
self._api_base = config.api_base
|
||||
self._api_key = config.api_key
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return 1536
|
||||
|
||||
def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
raise ValueError("文本列表不能为空")
|
||||
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:
|
||||
"""工厂函数."""
|
||||
if config.mode == "local":
|
||||
return LocalEmbedder(config)
|
||||
elif config.mode == "api":
|
||||
return APIEmbedder(config)
|
||||
else:
|
||||
raise ValueError(f"不支持的嵌入模式: {config.mode}")
|
||||
|
||||
|
||||
def batch_embed(embedder: Embedder, texts: list[str], batch_size: int = 32) -> list[list[float]]:
|
||||
"""分批嵌入，避免一次性传入过多文本导致 OOM."""
|
||||
all_embeddings = []
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i:i + batch_size]
|
||||
all_embeddings.extend(embedder.embed(batch))
|
||||
return all_embeddings
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 修改 ingest.py — 使用 batch_embed**
|
||||
|
||||
将 `ingest_content()` 中的:
|
||||
```python
|
||||
embeddings = self.embedder.embed(texts)
|
||||
```
|
||||
改为:
|
||||
```python
|
||||
from src.core.embedder import batch_embed
|
||||
embeddings = batch_embed(self.embedder, texts)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: 补充 API 模式测试**
|
||||
|
||||
在 `tests/test_embedder.py` 追加:
|
||||
```python
|
||||
class TestAPIEmbedder:
|
||||
"""API 嵌入器测试."""
|
||||
|
||||
def test_api_embedder_init(self):
|
||||
from src.core.embedder import APIEmbedder
|
||||
cfg = EmbedConfig(mode="api", api_base="https://api.test.com", api_key="sk-test")
|
||||
emb = APIEmbedder(cfg)
|
||||
assert emb.dimension == 1536
|
||||
|
||||
def test_api_embedder_empty_raises(self):
|
||||
from src.core.embedder import APIEmbedder
|
||||
cfg = EmbedConfig(mode="api")
|
||||
emb = APIEmbedder(cfg)
|
||||
with pytest.raises(ValueError):
|
||||
emb.embed([])
|
||||
|
||||
class TestBatchEmbed:
|
||||
"""分批嵌入测试."""
|
||||
|
||||
def test_batch_embed(self):
|
||||
from src.core.embedder import batch_embed, create_embedder
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
texts = ["测试文本"] * 70 # 超过 batch_size 32
|
||||
results = batch_embed(embedder, texts, batch_size=32)
|
||||
assert len(results) == 70
|
||||
assert len(results[0]) == embedder.dimension
|
||||
```
|
||||
|
||||
同时把原来的 `TestEmbedder` 中的 `create_embedder(EmbedConfig(mode="local"))` 改为检查返回类型是 `LocalEmbedder`。
|
||||
|
||||
- [ ] **Step 4: 运行测试**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/test_embedder.py tests/test_ingest.py tests/test_search.py tests/test_api.py -v
|
||||
```
|
||||
Expected: 全部通过
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "refactor: strategy-pattern embedder, remove global env vars, add batch embedding"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: 服务层重构 (D-02, D-03, C-07)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/server/app.py` — FastAPI Depends 注入替代全局单例；异步嵌入；真实健康检查
|
||||
- Create: `src/server/deps.py` — 依赖注入工厂
|
||||
|
||||
- [ ] **Step 1: 创建 deps.py**
|
||||
|
||||
```python
|
||||
"""FastAPI 依赖注入."""
|
||||
import os
|
||||
import functools
|
||||
import logging
|
||||
|
||||
from src.core.config import load_config, 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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AppState:
|
||||
"""应用级共享状态 (替代模块级全局变量)."""
|
||||
|
||||
def __init__(self):
|
||||
config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml")
|
||||
self.config = load_config(config_path)
|
||||
|
||||
data_dir = os.environ.get("MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir)
|
||||
self.db = VectorDB(persist_dir=data_dir)
|
||||
|
||||
self.embedder = create_embedder(self.config.embed)
|
||||
|
||||
collection = os.environ.get("MD_VECTOR_DB_COLLECTION", self.config.chroma.collection_name)
|
||||
self.searcher = Searcher(self.db, self.embedder, collection)
|
||||
self.ingestor = DocumentIngestor(self.db, self.embedder, collection)
|
||||
|
||||
def is_healthy(self) -> dict:
|
||||
"""真实健康检查."""
|
||||
status = {"status": "ok", "checks": {}}
|
||||
try:
|
||||
count = self.searcher.collection.count()
|
||||
status["checks"]["chromadb"] = {"status": "ok", "count": count}
|
||||
except Exception as e:
|
||||
status["checks"]["chromadb"] = {"status": "error", "detail": str(e)}
|
||||
status["status"] = "degraded"
|
||||
try:
|
||||
_ = self.embedder.dimension
|
||||
status["checks"]["embedder"] = {"status": "ok"}
|
||||
except Exception as e:
|
||||
status["checks"]["embedder"] = {"status": "error", "detail": str(e)}
|
||||
status["status"] = "degraded"
|
||||
return status
|
||||
|
||||
|
||||
_state: AppState | None = None
|
||||
|
||||
|
||||
def get_state() -> AppState:
|
||||
"""获取应用状态单例."""
|
||||
global _state
|
||||
if _state is None:
|
||||
logger.info("初始化应用状态...")
|
||||
_state = AppState()
|
||||
return _state
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 重写 app.py — 使用 Depends 注入**
|
||||
|
||||
```python
|
||||
"""FastAPI 服务层."""
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Depends, Request
|
||||
from fastapi.responses import RedirectResponse, JSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.server.auth import verify_api_key, rate_limiter
|
||||
from src.server.deps import get_state, AppState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# -- 请求模型 --
|
||||
class IngestRequest(BaseModel):
|
||||
file_path: str | None = None
|
||||
content: str | None = None
|
||||
file_name: str | None = Field(default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$")
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
query: str = Field(..., min_length=1, max_length=2000)
|
||||
top_k: int = Field(default=10, ge=1, le=100)
|
||||
|
||||
# -- App --
|
||||
app = FastAPI(title="md-vector-db", version="0.1.0")
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def rate_limit_middleware(request: Request, call_next):
|
||||
await rate_limiter(request)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@app.get("/")
|
||||
def root():
|
||||
return RedirectResponse(url="/docs")
|
||||
|
||||
|
||||
@app.get("/api/v1/health")
|
||||
def health(state: AppState = Depends(get_state)):
|
||||
return state.is_healthy()
|
||||
|
||||
|
||||
@app.get("/api/v1/collections")
|
||||
def list_collections(state: AppState = Depends(get_state)):
|
||||
info = state.searcher.get_collection_info()
|
||||
sources = state.searcher.list_sources()
|
||||
return {"collections": [info], "sources": sources}
|
||||
|
||||
|
||||
@app.post("/api/v1/ingest")
|
||||
def ingest_document(req: IngestRequest, state: AppState = Depends(get_state),
|
||||
_: bool = Depends(verify_api_key)):
|
||||
try:
|
||||
if req.file_path:
|
||||
path = Path(req.file_path).resolve()
|
||||
if not str(path).startswith(str(Path.cwd().resolve())):
|
||||
raise HTTPException(status_code=400, detail="不允许访问该路径")
|
||||
if not path.exists():
|
||||
raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}")
|
||||
count = state.ingestor.ingest_file(str(path))
|
||||
file_name = path.name
|
||||
elif req.content:
|
||||
file_name = req.file_name or "untitled.md"
|
||||
count = state.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:
|
||||
logger.exception("入库失败")
|
||||
raise HTTPException(status_code=500, detail="服务器内部错误")
|
||||
|
||||
|
||||
@app.post("/api/v1/search")
|
||||
def search_documents(req: SearchRequest, state: AppState = Depends(get_state),
|
||||
_: bool = Depends(verify_api_key)):
|
||||
results = state.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, state: AppState = Depends(get_state),
|
||||
_: bool = Depends(verify_api_key)):
|
||||
deleted = state.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 3: 运行全部测试**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
Expected: 全部通过
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "refactor: FastAPI Depends injection, async-safe, real health check"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 8: 测试补充与回归 (T-02, D-05, T-03)
|
||||
|
||||
**Files:**
|
||||
- Modify: `tests/test_api.py` — 补充 DELETE 端点测试、ingest_file 测试
|
||||
- Modify: `tests/test_ingest.py` — 补充 DocumentIngestor 集成测试
|
||||
|
||||
- [ ] **Step 1: 补充 API 测试**
|
||||
|
||||
在 `tests/test_api.py` 追加:
|
||||
```python
|
||||
class TestDeleteEndpoint:
|
||||
"""删除端点."""
|
||||
|
||||
def test_delete_nonexistent(self, client):
|
||||
response = client.delete("/api/v1/documents/nonexistent.md")
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_delete_ingested(self, client):
|
||||
client.post("/api/v1/ingest", json={"content": "# Test", "file_name": "del.md"})
|
||||
response = client.delete("/api/v1/documents/del.md")
|
||||
assert response.status_code == 200
|
||||
# 删除后搜索不应返回结果
|
||||
search_resp = client.post("/api/v1/search", json={"query": "Test", "top_k": 3})
|
||||
results = search_resp.json()["results"]
|
||||
sources = [r["source_file"] for r in results]
|
||||
assert "del.md" not in sources
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 补充 ingest 集成测试**
|
||||
|
||||
在 `tests/test_ingest.py` 追加:
|
||||
```python
|
||||
class TestIngestorIntegration:
|
||||
"""入库器集成测试 (使用真实 embedder)."""
|
||||
|
||||
def test_ingest_content_real(self, tmp_path):
|
||||
import tempfile
|
||||
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
|
||||
|
||||
db = VectorDB(persist_dir=str(tmp_path))
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
ingestor = DocumentIngestor(db, embedder, "test_integration")
|
||||
|
||||
count = ingestor.ingest_content("# Hello\nWorld.", "hello.md")
|
||||
assert count > 0
|
||||
assert ingestor.collection.count() == count
|
||||
|
||||
def test_ingest_deduplicates(self, tmp_path):
|
||||
"""重复入库同一文件会去重."""
|
||||
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
|
||||
|
||||
db = VectorDB(persist_dir=str(tmp_path))
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
ingestor = DocumentIngestor(db, embedder, "test_dedup")
|
||||
|
||||
c1 = ingestor.ingest_content("# A", "dup.md")
|
||||
c2 = ingestor.ingest_content("# B", "dup.md")
|
||||
# 第二次入库应覆盖第一次，总数不翻倍
|
||||
assert ingestor.collection.count() == c2
|
||||
```
|
||||
|
||||
- [ ] **Step 3: 运行全部测试**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
Expected: 35+ passed
|
||||
|
||||
- [ ] **Step 4: 最终 Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "test: add DELETE endpoint tests and ingest integration tests"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 9: 日志系统 + README 更新
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/server/app.py` — 启动时配置 logging
|
||||
- Modify: `src/cli/main.py` — 添加 logging
|
||||
- Modify: `README.md` — 更新 API 文档
|
||||
|
||||
- [ ] **Step 1: 在 app.py 顶部配置日志**
|
||||
|
||||
```python
|
||||
import logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
handlers=[logging.StreamHandler()]
|
||||
)
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行测试确认无回归**
|
||||
|
||||
```bash
|
||||
uv run pytest tests/ -v --tb=short
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add -A && git commit -m "feat: add structured logging and update README"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 最终验证
|
||||
|
||||
```bash
|
||||
cd D:/Code/doing_exercises/programs/md-vector-db
|
||||
uv run pytest tests/ -v
|
||||
uv run python -m src.cli.main search "测试"
|
||||
```
|
||||
Reference in New Issue
Block a user