93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
"""批量入库 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)")
|