tps-dashboard/backend/scripts/import_obsidian_kb.py

"""
批量导入脚本：将 Obsidian 知识库目录下的 .md 文件
导入 TPS 知识库（knowledge_items + knowledge_embeddings）。

运行方式：
    cd backend && python -m scripts.import_obsidian_kb --limit 10
    cd backend && python -m scripts.import_obsidian_kb              # 全量
    cd backend && python -m scripts.import_obsidian_kb --dry-run    # 只预览不写库
"""

import argparse
import sys
from pathlib import Path

from sqlmodel import Session, select

from app.db.session import engine
from app.models.knowledge import KnowledgeItem
from app.services.knowledge_service import KnowledgeService

# ── 默认 Obsidian 知识库根目录 ────────────────────────────────
DEFAULT_OBSIDIAN_DIR = r"E:\AIworks\obsidian\MSTknowledge_base\知识库"

# ── 跳过的目录名 ────────────────────────────────────────────────
SKIP_DIRS = {".obsidian", "99-原始素材"}


def collect_md_files(root: Path) -> list[Path]:
    """递归收集 root 下所有 .md 文件，跳过 SKIP_DIRS 中的目录。"""
    md_files = []
    for p in root.rglob("*.md"):
        # 检查路径中是否包含跳过目录
        parts = p.relative_to(root).parts
        if any(skip in parts for skip in SKIP_DIRS):
            continue
        md_files.append(p)
    return sorted(md_files, key=lambda p: str(p.relative_to(root)))


def check_exists(source_file_name: str) -> bool:
    """查 knowledge_items 表，判断 source_file_name 是否已存在。"""
    with Session(engine) as session:
        stmt = select(KnowledgeItem).where(
            KnowledgeItem.source_file_name == source_file_name
        )
        existing = session.exec(stmt).first()
        return existing is not None


def has_frontmatter(content: str) -> bool:
    """检查文件内容是否包含 YAML frontmatter（--- 包裹）。"""
    stripped = content.strip()
    if not stripped.startswith("---"):
        return False
    # 找到第二个 ---
    second = stripped.find("---", 3)
    return second > 0


def main():
    parser = argparse.ArgumentParser(description="批量导入 Obsidian 知识库 .md 文件到 TPS 知识库")
    parser.add_argument("--dir", type=str, default=DEFAULT_OBSIDIAN_DIR,
                        help="Obsidian 知识库根目录")
    parser.add_argument("--limit", type=int, default=0,
                        help="只处理前 N 个文件（测试用，0=全量）")
    parser.add_argument("--dry-run", action="store_true",
                        help="只打印解析结果，不真正调 API 写库")
    args = parser.parse_args()

    root = Path(args.dir)
    if not root.is_dir():
        print(f"[ERROR] 目录不存在: {root}")
        sys.exit(1)

    # 收集文件
    md_files = collect_md_files(root)
    if not md_files:
        print("[ERROR] 未找到任何 .md 文件")
        sys.exit(1)

    total = len(md_files)
    if args.limit > 0:
        md_files = md_files[:args.limit]
        print(f"共发现 {total} 个 .md 文件，本次处理前 {len(md_files)} 个\n")
    else:
        print(f"共发现 {total} 个 .md 文件，开始全量导入...\n")

    if args.dry_run:
        print("【DRY-RUN 模式】只预览解析结果，不写入数据库\n")

    service = KnowledgeService()
    success_count = 0
    skip_exists_count = 0
    skip_empty_count = 0
    skip_no_fm_count = 0
    fail_count = 0
    processed = len(md_files)

    for idx, file_path in enumerate(md_files, 1):
        rel_path = str(file_path.relative_to(root)).replace("\\", "/")

        # 跳过空文件
        if file_path.stat().st_size == 0:
            print(f"[{idx}/{processed}] [SKIP] {rel_path} -- 空文件，跳过")
            skip_empty_count += 1
            continue

        # 读取内容
        try:
            content_bytes = file_path.read_bytes()
        except Exception as e:
            print(f"[{idx}/{processed}] [FAIL] {rel_path} -- 读取失败：{e}")
            fail_count += 1
            continue

        # 检查 frontmatter
        content_str = content_bytes.decode("utf-8", errors="replace")
        if not has_frontmatter(content_str):
            print(f"[{idx}/{processed}] [SKIP] {rel_path} -- 无 frontmatter，跳过")
            skip_no_fm_count += 1
            continue

        # 查重（用 Obsidian 相对路径作为 source_file_name）
        source_file_name = rel_path
        if check_exists(source_file_name):
            print(f"[{idx}/{processed}] [SKIP] {rel_path} -- 已存在，跳过")
            skip_exists_count += 1
            continue

        # dry-run 模式：只解析不入库
        if args.dry_run:
            try:
                parsed = service.parse_md_file(content_bytes, source_file_name)
                char_count = len(parsed["content_md"])
                print(f"[{idx}/{processed}] [PREVIEW] {rel_path}")
                print(f"         title={parsed['title']}")
                print(f"         source_type={parsed['source_type']}, author={parsed['author']}")
                print(f"         metadata={parsed['metadata']}")
                print(f"         entities={len(parsed['related_entities'] or [])}条, 正文{char_count}字")
                success_count += 1
            except Exception as e:
                print(f"[{idx}/{processed}] [FAIL] {rel_path} -- 解析失败：{e}")
                fail_count += 1
            continue

        # 正式入库
        try:
            parsed = service.parse_md_file(content_bytes, source_file_name)
            char_count = len(parsed["content_md"])
            service.store_md_file(file_content=content_bytes, file_name=source_file_name)
            print(f"[{idx}/{processed}] [OK] {rel_path} -- 入库成功（{char_count}字，类型={parsed['source_type']}）")
            success_count += 1
        except Exception as e:
            print(f"[{idx}/{processed}] [FAIL] {rel_path} -- 入库失败：{e}")
            fail_count += 1

    # 汇总
    print(f"\n{'='*60}")
    mode_label = "预览" if args.dry_run else "导入"
    print(f"{mode_label}完成：成功 {success_count} 篇"
          f" / 跳过(已存在) {skip_exists_count} 篇"
          f" / 跳过(空文件) {skip_empty_count} 篇"
          f" / 跳过(无frontmatter) {skip_no_fm_count} 篇"
          f" / 失败 {fail_count} 篇")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()