AIyaodianzhushou/data-pipeline/import_all.py

"""
批量导入脚本 — 扫描 data/ 目录下所有兼容格式的 JSON 文件，逐个导入数据库。
兼容的文件格式：列表中每个元素包含 drug_id, name, sections, source 字段。
自动跳过 drug_index.json（只有索引无内容）等非兼容文件。
"""
import json
import asyncio
import os
import sys
from pathlib import Path

# 确保能导入 ingest 模块
sys.path.insert(0, str(Path(__file__).resolve().parent))

from ingest import ingest_drugs, load_env

# ============================================
# 配置：哪些文件跳过（非兼容格式）
# ============================================
SKIP_FILES = {
    "drug_index.json",        # 只有药品名和页码，无 sections 内容
    "catalog_volume1.json",   # 目录结构数据，非药品内容
    "catalog_volume2.json",
    "catalog_volume3.json",
    "catalog_volume4.json",
}


def is_compatible_drug_data(filepath: str) -> bool:
    """检查 JSON 文件是否为兼容的药品数据格式"""
    filename = os.path.basename(filepath)
    if filename in SKIP_FILES:
        return False

    try:
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)

        if not isinstance(data, list) or len(data) == 0:
            return False

        # 检查第一个元素是否有必要字段
        first = data[0]
        required = ["drug_id", "name", "sections"]
        return all(k in first for k in required)

    except (json.JSONDecodeError, Exception):
        return False


async def main():
    load_env()

    data_dir = Path(__file__).resolve().parent / "data"
    json_files = sorted(data_dir.glob("*.json"))

    print("=" * 60)
    print("📂 扫描数据目录:", data_dir)
    print(f"   发现 {len(json_files)} 个 JSON 文件")
    print("=" * 60)

    compatible = []
    skipped = []

    for f in json_files:
        if is_compatible_drug_data(str(f)):
            compatible.append(f)
            print(f"   ✅ {f.name} — 兼容，将导入")
        else:
            skipped.append(f)
            print(f"   ⏭️  {f.name} — 跳过（非药品数据格式）")

    if not compatible:
        print("\n❌ 没有找到可导入的数据文件！")
        print("   需要包含 drug_id, name, sections, source 字段的 JSON 数组")
        return

    print(f"\n📦 共 {len(compatible)} 个文件待导入")
    print("-" * 60)

    total_drugs = 0
    total_chunks = 0

    for f in compatible:
        try:
            print(f"\n🚀 正在导入: {f.name} ...")
            await ingest_drugs(str(f))
            # ingest_drugs 已经打印了统计信息

            # 数一下这个文件有多少条
            with open(f, "r", encoding="utf-8") as fh:
                data = json.load(fh)
            drug_count = len(data) if isinstance(data, list) else 0
            total_drugs += drug_count

        except Exception as e:
            print(f"   ❌ 导入失败: {e}")
            continue

    print("\n" + "=" * 60)
    print(f"🎉 全部导入完成！共处理 {len(compatible)} 个文件")
    print("=" * 60)

    # 提示：如果有旧数据没有 vec，执行修复 SQL
    print("""
💡 提示：如果之前导入的数据缺少 vec 列（pgvector），
   请在 psql 中执行以下 SQL 修复：

   UPDATE drug_chunks
   SET vec = (embedding::text)::vector
   WHERE vec IS NULL AND embedding IS NOT NULL;
""")


if __name__ == "__main__":
    asyncio.run(main())