diff --git a/src/cli/main.py b/src/cli/main.py new file mode 100644 index 0000000..8aff262 --- /dev/null +++ b/src/cli/main.py @@ -0,0 +1,90 @@ +"""命令行工具入口.""" +from pathlib import Path +import sys + +import typer +import uvicorn + +# 确保 src 在路径中 +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from src.core.config import load_config +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 文档向量数据库管理工具") + + +def _get_components(config_path: str = "config.yaml"): + """初始化所有组件.""" + 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 + + +@app.command() +def ingest(file_path: str): + """入库单个 Markdown 文件.""" + _, _, ingestor = _get_components() + count = ingestor.ingest_file(file_path) + typer.echo(f"✅ 已入库: {file_path} ({count} 个 chunks)") + + +@app.command() +def ingest_dir(dir_path: str): + """入库目录下所有 Markdown 文件.""" + _, _, ingestor = _get_components() + 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") + + +@app.command() +def search(query: str, top_k: int = 10): + """语义检索.""" + _, searcher, _ = _get_components() + results = searcher.search(query, top_k=top_k) + 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() +def serve(port: int = 8000): + """启动 HTTP 服务.""" + typer.echo(f"🚀 启动服务: http://0.0.0.0:{port}") + uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False) + + +@app.command() +def stats(): + """查看统计信息.""" + _, searcher, _ = _get_components() + info = searcher.get_collection_info() + sources = searcher.list_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: + typer.echo(f" - {s}") + + +if __name__ == "__main__": + app()