feat: multi-file, stdin, and glob support for ingest; fix stderr UTF-8
This commit is contained in:
+45
-4
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user