"""命令行工具入口 — 可作为 MCP tool 直接调用.""" from pathlib import Path import sys import io import json import os import glob as _glob from typing import Annotated import typer import uvicorn if sys.stdout.encoding != "utf-8": sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace") sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace") 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.server.deps import get_state, get_default_collection 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/ -C my_project\n" " md-vector-db search \"如何配置\" -k 5 --json\n" " md-vector-db serve -p 8000", ) # -- 共享初始化 -- 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() # -- 共享选项 -- ConfigOpt = Annotated[ str, 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_paths: Annotated[ list[str] | None, typer.Argument(help="Markdown 文件路径 (可多个, 或 - 从标准输入读取)"), ] = None, name: Annotated[str | None, typer.Option("--name", help="标准输入模式下的虚拟文件名")] = None, config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): _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"[OK] 已入库: stdin -> {count} chunks [{ingestor.collection_name}]") return # 多文件模式 if file_paths: total = 0 for fp in file_paths: if not is_path_within_workspace(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_path_within_workspace(m): typer.echo(f"[SKIP] 不安全的路径: {m}", err=True) continue c = ingestor.ingest_file(m) typer.echo(f" {m}: {c} chunks") total += c elif p.is_file(): c = ingestor.ingest_file(fp) typer.echo(f" {fp}: {c} chunks") total += c else: 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) @app.command(help="批量入库目录下所有 .md 文件 (递归).") def ingest_dir( dir_path: Annotated[str, typer.Argument(help="Markdown 文件目录")], config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): _init_config(config) if not is_path_within_workspace(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"[SKIP] 目录中未找到 .md 文件: {dir_path}") return total = sum(results.values()) for name, count in results.items(): typer.echo(f" {name}: {count} chunks") typer.echo(f"[OK] 共入库 {len(results)} 个文件, {total} 个 chunks [{ingestor.collection_name}]") @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, json_output: Annotated[bool, typer.Option("--json", help="以 JSON 格式输出")] = False, config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): _init_config(config) state = get_state() searcher = state.get_searcher(_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 for i, r in enumerate(results, 1): typer.echo(f"\n--- 结果 {i} (相似度: {r['score']:.4f}) ---") typer.echo(f"来源: {r['source_file']}") if r["section_title"]: typer.echo(f"章节: {r['section_title']}") preview = r["content"][:200] + "..." if len(r["content"]) > 200 else r["content"] typer.echo(preview) @app.command(help="启动 HTTP API 服务.") def serve( port: Annotated[int, typer.Option("--port", "-p", help="监听端口")] = 8000, config: ConfigOpt = DEFAULT_CONFIG_PATH, ): # 传递 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.") def stats( json_output: Annotated[bool, typer.Option("--json", help="以 JSON 格式输出")] = False, config: ConfigOpt = DEFAULT_CONFIG_PATH, collection: CollectionOpt = None, ): _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} 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)}") if sources: typer.echo("\n源文件列表:") for s in sources: typer.echo(f" - {s}") if __name__ == "__main__": app()