feat: add --json output and --collection to CLI for MCP scripting

src/cli/main.py

@@ -1,13 +1,13 @@
-"""命令行工具入口."""
+"""命令行工具入口 — 可作为 MCP tool 直接调用."""
 from pathlib import Path
 import sys
 import io
+import json
 from typing import Annotated
 
 import typer
 import uvicorn
 
-# Windows 终端 GBK → UTF-8
 if sys.stdout.encoding != "utf-8":
     sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
 
@@ -16,74 +16,102 @@ 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.ingest import DocumentIngestor
-from src.core.search import Searcher
 
 app = typer.Typer(
     name="md-vector-db",
     help="Markdown 文档向量数据库 — 入库、嵌入、语义检索",
     epilog="示例:\n"
     "  md-vector-db ingest docs/readme.md\n"
-    "  md-vector-db ingest-dir ./md_docs/\n"
+    "  md-vector-db ingest-dir ./md_docs/ -C my_project\n"
-    "  md-vector-db search \"如何配置\" --top-k 5\n"
+    "  md-vector-db search \"如何配置\" -k 5 --json\n"
-    "  md-vector-db serve --port 8000",
+    "  md-vector-db serve -p 8000",
 )
 
 
-def _get_components(config_path: str = DEFAULT_CONFIG_PATH):
+# -- 共享组件 (懒加载) --
-    """初始化所有组件."""
+_db: VectorDB | None = None
-    cfg = load_config(config_path)
+_embedder = None
-    db = VectorDB(persist_dir=cfg.chroma.persist_dir)
+_cfg = None
-    embedder = create_embedder(cfg.embed)
+
-    searcher = Searcher(db, embedder, cfg.chroma.collection_name)
+
-    ingestor = DocumentIngestor(db, embedder, cfg.chroma.collection_name)
+def _init_shared(config_path: str = DEFAULT_CONFIG_PATH):
-    return cfg, searcher, ingestor
+    """初始化 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 _cfg.chroma.collection_name if _cfg else "markdown_docs"
+
+
+def _resolve_collection(collection: str | None) -> str:
+    return collection or _get_default_collection()
 
 
 # -- 共享选项 --
-ConfigOption = Annotated[
+ConfigOpt = Annotated[
     str,
-    typer.Option(
+    typer.Option("--config", "-c", help="配置文件路径", show_default="config.yaml"),
-        "--config", "-c",
-        help="配置文件路径 (默认: config.yaml)",
-        show_default="config.yaml",
-    ),
 ]
 
+CollectionOpt = Annotated[
+    str | None,
+    typer.Option("--collection", "-C", help="目标 collection（默认使用配置文件中的）"),
+]
+
+
 # --- 命令 ---
 
 
 @app.command(help="入库单个 Markdown 文件.")
 def ingest(
     file_path: Annotated[str, typer.Argument(help="Markdown 文件路径")],
-    config: ConfigOption = DEFAULT_CONFIG_PATH,
+    config: ConfigOpt = DEFAULT_CONFIG_PATH,
+    collection: CollectionOpt = None,
 ):
-    _, _, ingestor = _get_components(config)
+    _init_shared(config)
+    from src.core.ingest import DocumentIngestor
+    ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
     count = ingestor.ingest_file(file_path)
-    typer.echo(f"✅ 已入库: {file_path} ({count} 个 chunks)")
+    typer.echo(f"✅ 已入库: {file_path} → {count} chunks [{ingestor.collection_name}]")
 
 
 @app.command(help="批量入库目录下所有 .md 文件 (递归).")
 def ingest_dir(
-    dir_path: Annotated[str, typer.Argument(help="包含 Markdown 文件的目录路径")],
+    dir_path: Annotated[str, typer.Argument(help="Markdown 文件目录")],
-    config: ConfigOption = DEFAULT_CONFIG_PATH,
+    config: ConfigOpt = DEFAULT_CONFIG_PATH,
+    collection: CollectionOpt = None,
 ):
-    _, _, ingestor = _get_components(config)
+    _init_shared(config)
+    from src.core.ingest import DocumentIngestor
+    ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
     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")
+    typer.echo(f"✅ 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]")
 
 
-@app.command(help="语义检索已入库的文档.")
+@app.command(help="语义检索已入库的文档. 加 --json 输出机器可读 JSON.")
 def search(
     query: Annotated[str, typer.Argument(help="搜索关键词或自然语言查询")],
     top_k: Annotated[int, typer.Option("--top-k", "-k", help="返回结果数量 (1-100)")] = 10,
-    config: ConfigOption = DEFAULT_CONFIG_PATH,
+    json_output: Annotated[bool, typer.Option("--json", help="以 JSON 格式输出")] = False,
+    config: ConfigOpt = DEFAULT_CONFIG_PATH,
+    collection: CollectionOpt = None,
 ):
-    _, searcher, _ = _get_components(config)
+    _init_shared(config)
+    from src.core.search import Searcher
+    searcher = Searcher(_db, _embedder, _resolve_collection(collection))
     results = searcher.search(query, top_k=top_k)
+
+    if json_output:
+        typer.echo(json.dumps(results, ensure_ascii=False, indent=2))
+        return
+
     if not results:
         typer.echo("未找到匹配结果。")
         return
@@ -99,21 +127,30 @@ def search(
 @app.command(help="启动 HTTP API 服务.")
 def serve(
     port: Annotated[int, typer.Option("--port", "-p", help="监听端口")] = 8000,
-    config: ConfigOption = DEFAULT_CONFIG_PATH,
+    config: ConfigOpt = DEFAULT_CONFIG_PATH,
 ):
-    # serve 模式下 config 仅用于校验路径存在，实际由 server 模块自行加载
     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)
 
 
-@app.command(help="查看向量库统计信息 (collection、chunks 数量、源文件列表).")
+@app.command(help="查看向量库统计信息. 加 --json 输出机器可读 JSON.")
 def stats(
-    config: ConfigOption = DEFAULT_CONFIG_PATH,
+    json_output: Annotated[bool, typer.Option("--json", help="以 JSON 格式输出")] = False,
+    config: ConfigOpt = DEFAULT_CONFIG_PATH,
+    collection: CollectionOpt = None,
 ):
-    _, searcher, _ = _get_components(config)
+    _init_shared(config)
+    from src.core.search import Searcher
+    searcher = Searcher(_db, _embedder, _resolve_collection(collection))
     info = searcher.get_collection_info()
     sources = searcher.list_sources()
+    data = {"collection": info["name"], "total_chunks": info["count"], "sources": sources}
+
+    if json_output:
+        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)}")