diff --git a/src/cli/main.py b/src/cli/main.py index 3f39f07..6755133 100644 --- a/src/cli/main.py +++ b/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 search \"如何配置\" --top-k 5\n" - " md-vector-db serve --port 8000", + " md-vector-db ingest-dir ./md_docs/ -C my_project\n" + " md-vector-db search \"如何配置\" -k 5 --json\n" + " md-vector-db serve -p 8000", ) -def _get_components(config_path: str = DEFAULT_CONFIG_PATH): - """初始化所有组件.""" - 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 +# -- 共享组件 (懒加载) -- +_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 _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( - "--config", "-c", - help="配置文件路径 (默认: config.yaml)", - show_default="config.yaml", - ), + typer.Option("--config", "-c", help="配置文件路径", 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 文件的目录路径")], - config: ConfigOption = DEFAULT_CONFIG_PATH, + dir_path: Annotated[str, typer.Argument(help="Markdown 文件目录")], + 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)}")