Serendipity
fc010cda06
- 用 load_config(config) 替代 get_state().config - 避免 CLI 主进程加载嵌入模型后 uvicorn 子进程再次加载 - 移除对 get_state 的隐式依赖（已有 DEFAULT_CONFIG_PATH 导入） Co-Authored-By: Claude <noreply@anthropic.com>
218 lines
7.6 KiB
Python
218 lines
7.6 KiB
Python
"""命令行工具入口 — 可作为 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
|
|
# 使用 load_config 而非 get_state()，避免在 CLI 主进程加载嵌入模型
|
|
from src.core.config import load_config
|
|
|
|
cfg = load_config(config)
|
|
typer.echo(f"启动服务: http://localhost:{port}")
|
|
typer.echo(f"API 文档: http://localhost:{port}/docs")
|
|
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()
|