AIyaodianzhushou/data-pipeline/import_from_openclaw.py

"""
Open Claw 药典数据 → PostgreSQL pgvector 知识库导入脚本

用法：
  # 仅转换格式（不连接数据库）
  python3 import_from_openclaw.py --convert-only

  # 转换 + 入库（需要PG + DashScope API Key）
  python3 import_from_openclaw.py

  # 跳过向量化（仅入库药品元数据）
  python3 import_from_openclaw.py --no-embedding
"""
import json
import asyncio
import os
import sys
import re
import hashlib
from pathlib import Path
from datetime import datetime, timezone

import httpx

# 项目根目录
ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = ROOT / "data"
OUTPUT_DIR = ROOT / "data" / "processed"

OPENCLAW_JSON = DATA_DIR / "all_drugs_final.json"
CONVERTED_JSON = OUTPUT_DIR / "converted_drugs.json"
CHUNKS_JSON = OUTPUT_DIR / "drug_chunks.json"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ============ 配置 ============
API_KEY = os.environ.get("QWEN_API_KEY", "")
EMBEDDING_URL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
EMBEDDING_MODEL = "text-embedding-v3"
DB_URL = os.environ.get(
    "DATABASE_URL",
    f"postgresql+asyncpg://{os.environ.get('POSTGRES_USER', 'postgres')}:{os.environ.get('POSTGRES_PASSWORD', 'change-me')}@{os.environ.get('POSTGRES_HOST', 'localhost')}:{os.environ.get('POSTGRES_PORT', '5432')}/{os.environ.get('POSTGRES_DB', 'pharmacopoeia')}"
)

# 尝试从 .env 加载
env_file = ROOT / ".env"
if env_file.exists():
    for line in open(env_file):
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            key, _, val = line.partition("=")
            os.environ.setdefault(key.strip(), val.strip())
    API_KEY = os.environ.get("QWEN_API_KEY", API_KEY)


def generate_drug_id(record: dict, idx: int) -> str:
    """生成 drug_id"""
    docid = record.get("docid", "")
    bid = record.get("蒲标网ID", "")
    eid = record.get("id", "")
    unique = docid or bid or eid or f"{idx:05d}"

    category = record.get("药品类别", "")
    if category == "化学药":
        prefix = "H"
    elif category == "生物制品":
        prefix = "S"
    elif category == "药用辅料":
        prefix = "F"
    elif "中药" in category:
        prefix = "Z"
    else:
        prefix = "X"

    # 使用 MD5 前8位确保唯一性
    hash_id = hashlib.md5(str(unique).encode()).hexdigest()[:8]
    return f"{prefix}2025{hash_id}"


def parse_source(record: dict) -> dict:
    """解析标准来源字段"""
    source_text = record.get("标准来源", "")
    version = "2025年版"
    volume = ""

    if "二部" in source_text:
        volume = "二部"
    elif "三部" in source_text:
        volume = "三部"
    elif "四部" in source_text:
        volume = "四部"
    elif "一部" in source_text:
        volume = "一部"
    elif "生物" in record.get("药品类别", ""):
        volume = "三部"
    elif "辅料" in record.get("药品类别", ""):
        volume = "四部"

    return {"version": version, "volume": volume, "page": ""}


def extract_sections(record: dict) -> dict:
    """提取药品 sections"""
    sections = {}

    # 优先使用章节字段
    zhangjie = record.get("章节")
    if zhangjie and isinstance(zhangjie, dict) and len(zhangjie) > 0:
        for key, value in zhangjie.items():
            if value and isinstance(value, str) and len(value.strip()) >= 10:
                sections[key] = value.strip()
        if len(sections) > 0:
            return sections

    # 从展开字段构建
    field_map = {
        "性状": ["性状"],
        "功效主治": ["功效主治", "定义", "定义/描述"],
        "用法用量": ["用法用量"],
        "禁忌": ["禁忌"],
        "注意事项": ["注意事项"],
        "成分": ["成分"],
    }

    for target_key, source_keys in field_map.items():
        for sk in source_keys:
            val = record.get(sk, "")
            if val and isinstance(val, str) and len(val.strip()) >= 5:
                sections[target_key] = val.strip()
                break

    return sections


def extract_subcategory(record: dict) -> str:
    """提取子分类"""
    cat = record.get("药品类别", "")
    if cat == "生物制品":
        detail = record.get("详细分类", "")
        sub = record.get("子分类", "")
        if sub and detail:
            return f"{detail} - {sub}"
        return detail or sub or "生物制品"
    if cat == "药用辅料":
        return "药用辅料"
    return record.get("剂型", record.get("类型", ""))


def convert_record(record: dict, idx: int) -> dict:
    """将 Open Claw 格式转换为项目标准格式"""
    source = parse_source(record)
    sections = extract_sections(record)

    drug = {
        "drug_id": generate_drug_id(record, idx),
        "name": record.get("药品名称", "").strip(),
        "name_en": record.get("英文名", record.get("拼音", "")).strip(),
        "pinyin": record.get("拼音", "").strip(),
        "category": record.get("药品类别", "").strip(),
        "subcategory": extract_subcategory(record).strip(),
        "sections": sections,
        "source": source,
    }

    return drug


async def get_embeddings(texts: list[str]) -> list[list[float]]:
    """批量获取 DashScope 嵌入向量"""
    async with httpx.AsyncClient(timeout=120) as client:
        resp = await client.post(
            EMBEDDING_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {API_KEY}",
            },
            json={
                "model": EMBEDDING_MODEL,
                "input": {"texts": texts},
                "parameters": {"text_type": "document"},
            },
        )
        data = resp.json()

    if data.get("code") and data["code"] != "":
        raise RuntimeError(f"Embedding API error: {data.get('message', data)}")

    return [item["embedding"] for item in data["output"]["embeddings"]]


def chunk_drug(drug: dict) -> list[dict]:
    """将一个药品切分为多个 chunk"""
    source_str = f"{drug['source']['version']} {drug['source']['volume']} {drug['source']['page']}".strip()
    chunks = []
    base_content = (
        f"药品名称：{drug['name']}\n"
        f"英文名：{drug['name_en']}\n"
        f"拼音：{drug['pinyin']}\n"
        f"类别：{drug['category']} / {drug['subcategory']}"
    )
    if source_str:
        base_content += f"\n来源：{source_str}"

    # 概览 chunk
    chunks.append({
        "drug_id": drug["drug_id"],
        "section": "概览",
        "content": base_content.strip(),
        "source": source_str,
        "chunk_index": 0,
    })

    # 各章节 chunk
    idx = 1
    for section_key, section_text in drug["sections"].items():
        if not section_text or len(section_text.strip()) < 20:
            continue
        content = f"【{drug['name']} - {section_key}】\n{section_text}\n\n来源：{source_str}"
        chunks.append({
            "drug_id": drug["drug_id"],
            "section": section_key,
            "content": content.strip(),
            "source": source_str,
            "chunk_index": idx,
        })
        idx += 1

    return chunks


async def main(convert_only: bool = False, no_embedding: bool = False):
    print("=" * 60)
    print("Open Claw 药典数据导入工具")
    print("=" * 60)

    # 1. 读取数据
    print(f"\n📖 读取数据: {OPENCLAW_JSON}")
    with open(OPENCLAW_JSON, "r", encoding="utf-8") as f:
        raw_data = json.load(f)

    print(f"   共 {len(raw_data)} 条原始记录")

    # 2. 格式转换
    print("\n🔄 格式转换中...")
    drugs = []
    skipped = 0
    for idx, record in enumerate(raw_data):
        name = record.get("药品名称", "").strip()
        if not name:
            skipped += 1
            continue
        drug = convert_record(record, idx)
        drugs.append(drug)

    print(f"   成功转换 {len(drugs)} 条，跳过 {skipped} 条（无名称）")

    # 3. 保存转换结果
    with open(CONVERTED_JSON, "w", encoding="utf-8") as f:
        json.dump(drugs, f, ensure_ascii=False, indent=2)
    print(f"   💾 已保存到: {CONVERTED_JSON}")

    if convert_only:
        print("\n✅ 格式转换完成（--convert-only）")
        return

    # 4. 切分 chunk
    print("\n✂️  切分 chunk...")
    all_chunks = []
    for drug in drugs:
        chunks = chunk_drug(drug)
        all_chunks.extend(chunks)

    print(f"   共切出 {len(all_chunks)} 个 chunk")

    with open(CHUNKS_JSON, "w", encoding="utf-8") as f:
        json.dump(all_chunks, f, ensure_ascii=False, indent=2)
    print(f"   💾 已保存到: {CHUNKS_JSON}")

    # 5. 向量化
    if no_embedding:
        print("\n⏭️  跳过向量化（--no-embedding）")
        print("\n✅ 转换完成。数据准备就绪，可后续执行入库。")
        return

    if not API_KEY:
        print("\n⚠️  未设置 QWEN_API_KEY，跳过向量化和入库")
        print("   请设置后重新运行: export QWEN_API_KEY=sk-xxx")
        return

    chunk_texts = [c["content"] for c in all_chunks]
    print(f"\n🧬 开始向量化 ({len(chunk_texts)} 个文本)...")
    all_vectors = []

    batch_size = 10
    for i in range(0, len(chunk_texts), batch_size):
        batch = chunk_texts[i : i + batch_size]
        try:
            vectors = await get_embeddings(batch)
            all_vectors.extend(vectors)
            n = min(i + batch_size, len(chunk_texts))
            print(f"   进度: {n}/{len(chunk_texts)}")
        except Exception as e:
            print(f"   ❌ 向量化失败 (batch {i}): {e}")
            print("   请检查 QWEN_API_KEY 是否正确")
            return

    print(f"   ✅ 向量化完成，维度: {len(all_vectors[0]) if all_vectors else 'N/A'}")

    # 6. 入库 PostgreSQL
    print("\n💾 写入 PostgreSQL...")
    try:
        from sqlalchemy.ext.asyncio import create_async_engine
        from sqlalchemy import text
    except ImportError:
        print("❌ 缺少 sqlalchemy 依赖，请安装:")
        print("   pip3 install sqlalchemy asyncpg")
        return

    engine = create_async_engine(DB_URL)

    # 写入 drug_chunks
    count = 0
    async with engine.begin() as conn:
        for i, (chunk, vec) in enumerate(zip(all_chunks, all_vectors)):
            await conn.execute(
                text("""
                    INSERT INTO drug_chunks (drug_id, section, content, source, chunk_index, embedding)
                    VALUES (:drug_id, :section, :content, :source, :chunk_index, :embedding)
                """),
                {
                    "drug_id": chunk["drug_id"],
                    "section": chunk["section"],
                    "content": chunk["content"],
                    "source": chunk["source"],
                    "chunk_index": chunk["chunk_index"],
                    "embedding": json.dumps(vec),
                },
            )
            count += 1

    # 写入 drugs
    drug_count = 0
    async with engine.begin() as conn:
        for drug in drugs:
            await conn.execute(
                text("""
                    INSERT INTO drugs (drug_id, name, name_en, pinyin, category, subcategory, sections, source_version, source_volume, source_page, is_active)
                    VALUES (:drug_id, :name, :name_en, :pinyin, :category, :subcategory, :sections, :source_version, :source_volume, :source_page, TRUE)
                    ON CONFLICT (drug_id) DO UPDATE SET
                        sections = EXCLUDED.sections,
                        updated_at = NOW()
                """),
                {
                    "drug_id": drug["drug_id"],
                    "name": drug["name"],
                    "name_en": drug["name_en"] or None,
                    "pinyin": drug["pinyin"] or None,
                    "category": drug["category"] or None,
                    "subcategory": drug["subcategory"] or None,
                    "sections": json.dumps(drug["sections"], ensure_ascii=False),
                    "source_version": drug["source"]["version"],
                    "source_volume": drug["source"]["volume"] or None,
                    "source_page": drug["source"]["page"] or None,
                },
            )
            drug_count += 1

    await engine.dispose()
    print(f"\n🎉 入库完成！")
    print(f"   药品: {drug_count} 条")
    print(f"   Chunk: {count} 条")
    print(f"   转换结果: {CONVERTED_JSON}")
    print(f"   Chunk 数据: {CHUNKS_JSON}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Open Claw 药典数据导入")
    parser.add_argument("--convert-only", action="store_true", help="仅转换格式，不连接数据库")
    parser.add_argument("--no-embedding", action="store_true", help="跳过向量化和入库")
    args = parser.parse_args()
    asyncio.run(main(convert_only=args.convert_only, no_embedding=args.no_embedding))