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":
|
if sys.stdout.encoding != "utf-8":
|
||||||
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
|
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))
|
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(
|
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,
|
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||||
collection: CollectionOpt = None,
|
collection: CollectionOpt = None,
|
||||||
):
|
):
|
||||||
_init_shared(config)
|
_init_shared(config)
|
||||||
from src.core.ingest import DocumentIngestor
|
from src.core.ingest import DocumentIngestor
|
||||||
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
|
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 文件 (递归).")
|
@app.command(help="批量入库目录下所有 .md 文件 (递归).")
|
||||||
@@ -89,6 +127,9 @@ def ingest_dir(
|
|||||||
from src.core.ingest import DocumentIngestor
|
from src.core.ingest import DocumentIngestor
|
||||||
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
|
ingestor = DocumentIngestor(_db, _embedder, _resolve_collection(collection))
|
||||||
results = ingestor.ingest_directory(dir_path)
|
results = ingestor.ingest_directory(dir_path)
|
||||||
|
if not results:
|
||||||
|
typer.echo(f"⚠️ 目录中未找到 .md 文件: {dir_path}")
|
||||||
|
return
|
||||||
total = sum(results.values())
|
total = sum(results.values())
|
||||||
for name, count in results.items():
|
for name, count in results.items():
|
||||||
typer.echo(f" 📄 {name}: {count} chunks")
|
typer.echo(f" 📄 {name}: {count} chunks")
|
||||||
|
|||||||
Reference in New Issue
Block a user