fix: 修复 44 个代码审查问题 (CRITICAL/HIGH/MEDIUM/LOW)
Batch 1 — CRITICAL (1): - 提取 is_safe_path() 到 src/core/security.py 公共模块 - CLI 和 ingest_obsidian.py 统一添加路径遍历防护 Batch 2 — HIGH (13) + 架构重构: - CLI 复用 deps.py AppState, 消除 30 行重复代码 - AppState/get_state 添加线程安全锁 - serve 命令传递 --config 到 uvicorn (H1) - OpenAIEmbedder 懒创建+复用 HTTP 客户端 (H2) - DashscopeEmbedder import 移到模块顶部 (H3) - 路径检查改用 os.path.commonpath (H4) - embedder.embed() 返回值长度检查 (H5) - 健康检查不泄露内部错误详情 (H7) - /api/v1/collections 添加 API Key 认证 (H8) - API Key 使用 hmac.compare_digest 恒定时间比较 (H9) - 添加 CORS 中间件 (H10) - ServerConfig 支持 SSL 配置 (H11) - HF_ENDPOINT 修改添加详细注释 (H12) Batch 3 — MEDIUM (20) + Splitter Protocol: - 定义 Splitter(Protocol) 接口, DocumentIngestor 接受可选 splitter - DashScope 响应添加结构验证 (M2) - ingest_obsidian.py 支持 CLI 参数和 OBSIDIAN_DIRS 环境变量 (M6) - scripts/serve.py 添加废弃警告 (M7) - content 限制 500KB, collection 正则限制字符集 (M12-M14) - 默认监听地址 127.0.0.1 (M16) - 添加安全响应头中间件 (M17) - verify_api_key 认证失败记录日志 (M19) Batch 4 — LOW (10): - CLI emoji 清理为纯文本标记 (L5) - logging.basicConfig 移到 FastAPI lifespan (L1) - VectorDB 添加 write_guard() 上下文管理器 (L3) - IngestRequest file_path/content 互斥校验 (L10) - ingest_obsidian.py 注释修正 (L6) 测试: 46 → 70 (+24) - tests/test_security.py: 11 个路径安全测试 - tests/test_deps.py: 11 个依赖注入测试 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+60
-53
@@ -16,9 +16,9 @@ if sys.stdout.encoding != "utf-8":
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from src.core.config import load_config, DEFAULT_CONFIG_PATH
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import create_embedder
|
||||
from src.core.config import DEFAULT_CONFIG_PATH
|
||||
from src.core.security import is_safe_path
|
||||
from src.server.deps import get_state, get_default_collection
|
||||
|
||||
app = typer.Typer(
|
||||
name="md-vector-db",
|
||||
@@ -31,30 +31,14 @@ app = typer.Typer(
|
||||
)
|
||||
|
||||
|
||||
# -- 共享组件 (懒加载) --
|
||||
_db: VectorDB | None = None
|
||||
_embedder = None
|
||||
_cfg = None
|
||||
|
||||
|
||||
def _init_shared(config_path: str = DEFAULT_CONFIG_PATH):
|
||||
"""初始化 db + embedder（全局共享）."""
|
||||
global _db, _embedder, _cfg
|
||||
if _db is None:
|
||||
_cfg = load_config(config_path)
|
||||
_db = VectorDB(persist_dir=_cfg.chroma.persist_dir)
|
||||
_embedder = create_embedder(_cfg.embed)
|
||||
|
||||
|
||||
def _get_default_collection() -> str:
|
||||
return os.environ.get(
|
||||
"MD_VECTOR_DB_COLLECTION",
|
||||
_cfg.chroma.collection_name if _cfg else "markdown_docs",
|
||||
)
|
||||
# -- 共享初始化 --
|
||||
def _init_config(config_path: str = DEFAULT_CONFIG_PATH):
|
||||
"""确保配置已加载并设置到环境变量 (供 deps.get_state 复用)."""
|
||||
os.environ["MD_VECTOR_CONFIG"] = config_path
|
||||
|
||||
|
||||
def _resolve_collection(collection: str | None) -> str:
|
||||
return collection or _get_default_collection()
|
||||
return collection or get_default_collection()
|
||||
|
||||
|
||||
# -- 共享选项 --
|
||||
@@ -82,40 +66,46 @@ def ingest(
|
||||
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||
collection: CollectionOpt = None,
|
||||
):
|
||||
_init_shared(config)
|
||||
from src.core.ingest import DocumentIngestor
|
||||
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
|
||||
_init_config(config)
|
||||
state = get_state()
|
||||
ingestor = state.get_ingestor(_resolve_collection(collection))
|
||||
|
||||
# 标准输入模式
|
||||
if file_paths and file_paths[0] == "-":
|
||||
content = sys.stdin.read()
|
||||
file_name = name or "stdin.md"
|
||||
count = ingestor.ingest_content(content, file_name)
|
||||
typer.echo(f"✅ 已入库: stdin → {count} chunks [{ingestor.collection_name}]")
|
||||
typer.echo(f"[OK] 已入库: stdin -> {count} chunks [{ingestor.collection_name}]")
|
||||
return
|
||||
|
||||
# 多文件模式
|
||||
if file_paths:
|
||||
total = 0
|
||||
for fp in file_paths:
|
||||
if not is_safe_path(fp):
|
||||
typer.echo(f"[SKIP] 不安全的路径: {fp}", err=True)
|
||||
continue
|
||||
# 支持通配符 (shell 展开或 Python glob)
|
||||
p = Path(fp)
|
||||
if "*" in fp or "?" in fp:
|
||||
matches = _glob.glob(fp, recursive=True)
|
||||
for m in matches:
|
||||
if not is_safe_path(m):
|
||||
typer.echo(f"[SKIP] 不安全的路径: {m}", err=True)
|
||||
continue
|
||||
c = ingestor.ingest_file(m)
|
||||
typer.echo(f" 📄 {m}: {c} chunks")
|
||||
typer.echo(f" {m}: {c} chunks")
|
||||
total += c
|
||||
elif p.is_file():
|
||||
c = ingestor.ingest_file(fp)
|
||||
typer.echo(f" 📄 {fp}: {c} chunks")
|
||||
typer.echo(f" {fp}: {c} chunks")
|
||||
total += c
|
||||
else:
|
||||
typer.echo(f"⚠️ 跳过 (非文件): {fp}", err=True)
|
||||
typer.echo(f"✅ 共入库 {total} chunks [{ingestor.collection_name}]")
|
||||
typer.echo(f"[SKIP] 非文件: {fp}", err=True)
|
||||
typer.echo(f"[OK] 共入库 {total} chunks [{ingestor.collection_name}]")
|
||||
return
|
||||
|
||||
# 无参数 → 显示帮助
|
||||
# 无参数 -> 显示帮助
|
||||
typer.echo("用法: md-vector-db ingest <文件1> [文件2 ...] 或 echo '内容' | md-vector-db ingest - --name doc.md", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
@@ -126,17 +116,20 @@ def ingest_dir(
|
||||
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||
collection: CollectionOpt = None,
|
||||
):
|
||||
_init_shared(config)
|
||||
from src.core.ingest import DocumentIngestor
|
||||
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
|
||||
_init_config(config)
|
||||
if not is_safe_path(dir_path):
|
||||
typer.echo(f"错误: 不安全的路径 — {dir_path}", err=True)
|
||||
raise typer.Exit(code=1)
|
||||
state = get_state()
|
||||
ingestor = state.get_ingestor(_resolve_collection(collection))
|
||||
results = ingestor.ingest_directory(dir_path)
|
||||
if not results:
|
||||
typer.echo(f"⚠️ 目录中未找到 .md 文件: {dir_path}")
|
||||
typer.echo(f"[SKIP] 目录中未找到 .md 文件: {dir_path}")
|
||||
return
|
||||
total = sum(results.values())
|
||||
for name, count in results.items():
|
||||
typer.echo(f" 📄 {name}: {count} chunks")
|
||||
typer.echo(f"✅ 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]")
|
||||
typer.echo(f" {name}: {count} chunks")
|
||||
typer.echo(f"[OK] 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]")
|
||||
|
||||
|
||||
@app.command(help="语义检索已入库的文档. 加 --json 输出机器可读 JSON.")
|
||||
@@ -147,9 +140,9 @@ def search(
|
||||
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||
collection: CollectionOpt = None,
|
||||
):
|
||||
_init_shared(config)
|
||||
from src.core.search import Searcher
|
||||
searcher = Searcher(_db, _embedder, _resolve_collection(collection))
|
||||
_init_config(config)
|
||||
state = get_state()
|
||||
searcher = state.get_searcher(_resolve_collection(collection))
|
||||
results = searcher.search(query, top_k=top_k)
|
||||
|
||||
if json_output:
|
||||
@@ -161,9 +154,9 @@ def search(
|
||||
return
|
||||
for i, r in enumerate(results, 1):
|
||||
typer.echo(f"\n--- 结果 {i} (相似度: {r['score']:.4f}) ---")
|
||||
typer.echo(f"📄 来源: {r['source_file']}")
|
||||
typer.echo(f"来源: {r['source_file']}")
|
||||
if r["section_title"]:
|
||||
typer.echo(f"📑 章节: {r['section_title']}")
|
||||
typer.echo(f"章节: {r['section_title']}")
|
||||
preview = r["content"][:200] + "..." if len(r["content"]) > 200 else r["content"]
|
||||
typer.echo(preview)
|
||||
|
||||
@@ -173,9 +166,23 @@ def serve(
|
||||
port: Annotated[int, typer.Option("--port", "-p", help="监听端口")] = 8000,
|
||||
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||
):
|
||||
typer.echo(f"🚀 启动服务: http://localhost:{port}")
|
||||
typer.echo(f"📖 API 文档: http://localhost:{port}/docs")
|
||||
uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False)
|
||||
# 传递 config 给 uvicorn 子进程 (通过环境变量)
|
||||
os.environ["MD_VECTOR_CONFIG"] = config
|
||||
typer.echo(f"启动服务: http://localhost:{port}")
|
||||
typer.echo(f"API 文档: http://localhost:{port}/docs")
|
||||
# 加载配置以获取 SSL 设置
|
||||
cfg = get_state().config
|
||||
ssl_kwargs = {}
|
||||
if cfg.server.ssl_keyfile and cfg.server.ssl_certfile:
|
||||
ssl_kwargs["ssl_keyfile"] = cfg.server.ssl_keyfile
|
||||
ssl_kwargs["ssl_certfile"] = cfg.server.ssl_certfile
|
||||
uvicorn.run(
|
||||
"src.server.app:app",
|
||||
host=cfg.server.host,
|
||||
port=port or cfg.server.port,
|
||||
reload=False,
|
||||
**ssl_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@app.command(help="查看向量库统计信息. 加 --json 输出机器可读 JSON.")
|
||||
@@ -184,9 +191,9 @@ def stats(
|
||||
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||
collection: CollectionOpt = None,
|
||||
):
|
||||
_init_shared(config)
|
||||
from src.core.search import Searcher
|
||||
searcher = Searcher(_db, _embedder, _resolve_collection(collection))
|
||||
_init_config(config)
|
||||
state = get_state()
|
||||
searcher = state.get_searcher(_resolve_collection(collection))
|
||||
info = searcher.get_collection_info()
|
||||
sources = searcher.list_sources()
|
||||
data = {"collection": info["name"], "total_chunks": info["count"], "sources": sources}
|
||||
@@ -195,9 +202,9 @@ def stats(
|
||||
typer.echo(json.dumps(data, ensure_ascii=False, indent=2))
|
||||
return
|
||||
|
||||
typer.echo(f"📊 Collection: {info['name']}")
|
||||
typer.echo(f"📦 总 chunks: {info['count']}")
|
||||
typer.echo(f"📄 源文件数: {len(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:
|
||||
|
||||
Reference in New Issue
Block a user