feat: add CLI management tool
This commit is contained in:
@@ -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()
|
||||||
Reference in New Issue
Block a user