feat: multi-file, stdin, and glob support for ingest; fix stderr UTF-8

This commit is contained in:
2026-07-05 02:14:17 +08:00
parent 4b2fd3d33a
commit d3cb88525c
+45 -4
View File
@@ -10,6 +10,7 @@ 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))
@@ -66,17 +67,54 @@ CollectionOpt = Annotated[
# --- 命令 ---
@app.command(help="入库单个 Markdown 文件.")
@app.command(help="入库 Markdown 文件. 支持多文件、通配符、标准输入.")
def ingest(
file_path: Annotated[str, typer.Argument(help="Markdown 文件路径")],
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_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 [{ingestor.collection_name}]")
# 标准输入模式
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"✅ 已入库: stdin → {count} chunks [{ingestor.collection_name}]")
return
# 多文件模式
if file_paths:
total = 0
for fp in file_paths:
# 支持通配符 (shell 展开或 Python glob)
from pathlib import Path as _Path
p = _Path(fp)
if "*" in fp or "?" in fp:
import glob as _glob
matches = _glob.glob(fp, recursive=True)
for m in matches:
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"⚠️ 跳过 (非文件): {fp}", err=True)
typer.echo(f"✅ 共入库 {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 文件 (递归).")
@@ -89,6 +127,9 @@ def ingest_dir(
from src.core.ingest import DocumentIngestor
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
results = ingestor.ingest_directory(dir_path)
if not results:
typer.echo(f"⚠️ 目录中未找到 .md 文件: {dir_path}")
return
total = sum(results.values())
for name, count in results.items():
typer.echo(f" 📄 {name}: {count} chunks")