vetcopilot-backend/scripts/init_knowledge_base.py

"""
知识库初始化脚本 - 递归扫描 veterinary_articles 目录下所有 PDF

使用方法：
    py -3.12 scripts/init_knowledge_base.py
"""
import os
import sys
import shutil
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.database import SessionLocal, init_db
from app.models.knowledge import KnowledgeDocument, KnowledgeChunk
from app.services import knowledge_service as ks
from app.config import settings
from app.utils.logger import logger

# 扫描目录：本地 veterinary_articles（含 site1/site2/site3）
ARTICLES_DIR = Path(r"C:\Users\HP\Documents\666666666666666666\veterinary_articles")
PDF_DEST_DIR = Path(settings.UPLOAD_DIR) / "knowledge_pdfs"

EXCLUDE_PATTERNS = [
    "cover", "index", "contents", "front", "preface", "acknowledg",
    "abbreviat", "glossary", "appendi", "reference", "bibliograph",
]


def should_include(filename: str) -> bool:
    name_lower = filename.lower()
    for pattern in EXCLUDE_PATTERNS:
        if pattern in name_lower:
            return False
    return True


def main():
    if not ARTICLES_DIR.exists():
        print(f"错误：目录不存在: {ARTICLES_DIR}")
        return

    # 确保目标目录存在
    PDF_DEST_DIR.mkdir(parents=True, exist_ok=True)

    # 初始化数据库
    print("初始化数据库...")
    init_db()

    # 初始化 FAISS 索引
    print("初始化 FAISS 索引...")
    ks.init_faiss_index()

    # 递归扫描所有 PDF
    print(f"\n扫描目录: {ARTICLES_DIR}")
    pdf_files = sorted(ARTICLES_DIR.rglob("*.pdf"))
    pdf_files = [f for f in pdf_files if not f.name.startswith("_")]
    print(f"发现 {len(pdf_files)} 个 PDF 文件")

    # 过滤
    pdf_files = [f for f in pdf_files if should_include(f.name)]
    print(f"过滤后剩余 {len(pdf_files)} 个文件待导入\n")

    # 按站点分组统计
    from collections import Counter
    site_counts = Counter()
    for f in pdf_files:
        parts = f.parts
        for p in parts:
            if p.startswith("site"):
                site_counts[p] += 1
                break
        else:
            site_counts["other"] += 1
    print("按站点分布:")
    for site, cnt in site_counts.most_common():
        print(f"  {site}: {cnt} 个")

    # 检查已存在
    db = SessionLocal()
    try:
        existing_hashes = set()
        all_docs = db.query(KnowledgeDocument).all()
        for d in all_docs:
            existing_hashes.add(d.file_hash)

        imported = 0
        skipped = 0
        failed = 0
        total = len(pdf_files)

        for idx, pdf_path in enumerate(pdf_files):
            file_hash = ks._compute_sha256(pdf_path)

            if file_hash in existing_hashes:
                if (idx + 1) % 100 == 0:
                    print(f"  [{idx+1}/{total}] 已导入={imported} 跳过={skipped} 失败={failed}")
                skipped += 1
                continue

            # 复制到知识库 PDF 目录
            dest_name = pdf_path.name
            dest_path = PDF_DEST_DIR / dest_name
            # 处理重名
            counter = 1
            while dest_path.exists():
                stem = pdf_path.stem
                dest_name = f"{stem}_{counter}.pdf"
                dest_path = PDF_DEST_DIR / dest_name
                counter += 1

            shutil.copy2(pdf_path, dest_path)

            try:
                result = ks.upload_document(dest_path, dest_name)

                if result.get("duplicate"):
                    skipped += 1
                    continue

                doc_id = result["id"]

                # 同步处理（大文件可能需要较长时间）
                ks.process_document(doc_id)

                doc = db.query(KnowledgeDocument).filter(
                    KnowledgeDocument.id == doc_id
                ).first()
                if doc and doc.status == "ready":
                    imported += 1
                else:
                    status = doc.status if doc else "unknown"
                    error = doc.error_message if doc else ""
                    failed += 1
                    if idx < 5:
                        print(f"  [失败] {pdf_path.name}: status={status} error={error[:80]}")

            except Exception as e:
                failed += 1
                if idx < 5:
                    print(f"  [异常] {pdf_path.name}: {e}")

            if (idx + 1) % 50 == 0 or idx < 5:
                print(f"  [{idx+1}/{total}] 已导入={imported} 跳过={skipped} 失败={failed}")

        print(f"\n{'='*50}")
        print(f"导入完成！")
        print(f"  成功导入: {imported}")
        print(f"  跳过（已存在）: {skipped}")
        print(f"  失败: {failed}")
        print(f"{'='*50}")

    finally:
        db.close()


if __name__ == "__main__":
    main()