"""批量入库 Obsidian 知识库. 用法: uv run python scripts/ingest_obsidian.py [目录1] [目录2] ... 若不传参数, 默认读取 OBSIDIAN_DIRS 环境变量 (逗号分隔的目录列表), 例如: OBSIDIAN_DIRS="D:/Code/Obsidian/博客,D:/Code/Obsidian/其他" uv run python scripts/ingest_obsidian.py """ import os import sys import time from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent / "src")) 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.security import is_safe_path MAX_SIZE = int(os.environ.get("INGEST_MAX_SIZE", "200000")) # 跳过超过此大小的文件 (字节), 默认 200KB progress_file = Path(__file__).parent.parent / "ingest_progress.txt" def log(msg): print(msg, flush=True) with open(progress_file, "a", encoding="utf-8") as f: f.write(msg + "\n") t0 = time.time() log("初始化...") try: cfg = load_config() db = VectorDB(persist_dir=cfg.chroma.persist_dir) embedder = create_embedder(cfg.embed) ingestor = DocumentIngestor(db, embedder, "obsidian_blog") except Exception as e: log(f"初始化失败: {e}") import traceback log(traceback.format_exc()) sys.exit(1) # 从 CLI 参数或环境变量获取目录列表 if len(sys.argv) > 1: targets = [(Path(d).name, d) for d in sys.argv[1:]] else: env_dirs = os.environ.get("OBSIDIAN_DIRS", "") if env_dirs: dirs = [d.strip() for d in env_dirs.split(",") if d.strip()] targets = [(Path(d).name, d) for d in dirs] else: log("错误: 未指定目标目录。请通过命令行参数或 OBSIDIAN_DIRS 环境变量提供。") log("用法: uv run python scripts/ingest_obsidian.py <目录1> [目录2] ...") sys.exit(1) files = [] skipped = [] for label, d in targets: if not Path(d).exists(): log(f"跳过不存在的目录: {d}") continue for f in Path(d).rglob("*.md"): if any(p.startswith(".") for p in f.parts): continue if "node_modules" in f.parts: continue size = f.stat().st_size if size > MAX_SIZE: skipped.append((f.name, size)) continue files.append((label, str(f))) log(f"待处理: {len(files)} 个文件") if skipped: log(f"跳过超大: {len(skipped)} 个") for name, sz in sorted(skipped, key=lambda x: -x[1]): log(f" SKIP {name} ({sz//1024}KB)") total = 0 for i, (label, fp) in enumerate(files, 1): name = Path(fp).name if not is_safe_path(fp): log(f"[{i}/{len(files)}] SKIP {label}/{name}: 不安全的路径 (路径遍历)") continue t1 = time.time() try: n = ingestor.ingest_file(fp) total += n dt = time.time() - t1 log(f"[{i}/{len(files)}] {n:>4d} chunks | {label}/{name} ({dt:.1f}s)") except Exception as e: log(f"[{i}/{len(files)}] ERROR {label}/{name}: {e}") elapsed = time.time() - t0 log(f"[DONE] 完成: {len(files)} 文件, {total} chunks ({elapsed:.0f}s)")