AIyaodianzhushou/data-pipeline/pdf2020_ingest.py

"""
2020 年版药典 PDF 导入脚本
解析 pharmacopoeia_2020_volume1_toc.pdf → 提取药品条目 → 向量化 → 写入 PostgreSQL
"""
import json, os, sys, re, hashlib, asyncio
from pathlib import Path
from collections import OrderedDict

import httpx
import fitz  # pymupdf
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text

# ============================================
# 配置
# ============================================
PDF_PATH = os.environ.get("PDF_2020_PATH", os.path.join(
    os.path.dirname(__file__), "..", "data", "pharmacopoeia_2020_volume1_toc.pdf"))

PG_HOST = os.environ.get("POSTGRES_HOST", "localhost")
PG_PORT = os.environ.get("POSTGRES_PORT", "5432")
PG_DB = os.environ.get("POSTGRES_DB", "pharmacopoeia")
PG_USER = os.environ.get("POSTGRES_USER", "postgres")
PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "postgres")
DB_URL = f"postgresql+asyncpg://{PG_USER}:{PG_PASSWORD}@{PG_HOST}:{PG_PORT}/{PG_DB}"

QWEN_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"
BATCH_SIZE = 10

# 2020 药典 section 标题识别
SECTION_HEADERS = OrderedDict([
    ("性状", "性状"), ("鉴别", "鉴别"), ("检查", "检查"),
    ("含量测定", "含量测定"), ("浸出物", "浸出物"),
    ("性味与归经", "性味与归经"), ("性味", "性味与归经"),
    ("功能与主治", "功能主治"), ("功能", "功能主治"), ("主治", "功能主治"),
    ("用法与用量", "用法用量"), ("用法", "用法用量"), ("用量", "用法用量"),
    ("注意", "注意事项"), ("注意事项", "注意事项"),
    ("规格", "规格"), ("贮藏", "贮藏"), ("类别", "类别"),
    ("制剂", "制剂"), ("附注", "附注"), ("禁忌", "禁忌"),
    ("不良反应", "不良反应"),
    ("处方", "处方"), ("制法", "制法"),
    ("包装", "包装"), ("有效期", "有效期"),
    ("执行标准", "执行标准"), ("批准文号", "批准文号"),
])


def load_env():
    env_file = Path(__file__).resolve().parent.parent / ".env"
    if env_file.exists():
        for line in open(env_file, encoding="utf-8"):
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, _, val = line.partition("=")
                os.environ.setdefault(key.strip(), val.strip())


def extract_drug_entries(pdf_path: str) -> list[dict]:
    """从 PDF 中提取药品条目"""
    doc = fitz.open(pdf_path)
    print(f"📖 PDF: {doc.page_count} 页")

    # 收集所有页面的文本
    all_text = ""
    for i in range(doc.page_count):
        text = doc[i].get_text()
        all_text += text + "\n"

    # 清理：去掉页眉页码、多余空白
    all_text = re.sub(r'\n{3,}', '\n\n', all_text)
    all_text = re.sub(r'^\s*\d+\s*$', '', all_text, flags=re.MULTILINE)

    # 找到正文开始位置（跳过前言/TOC）
    # 正文以药名开头，通常格式为 "药名" 后紧跟 "【性状】" 等 section
    # 扫描找到第一个药品条目
    lines = all_text.split('\n')

    # 策略：扫描【性状】标记，向前找药名
    entries = []
    current_name = None
    current_sections = OrderedDict()
    current_section = "正文"
    current_text = []
    in_entry = False

    for line in lines:
        line = line.strip()
        if not line:
            continue

        # 跳过纯页码和标题行
        if re.match(r'^\d{1,4}$', line):
            continue
        if line.startswith("中国药典") or line.startswith("ISBN"):
            continue
        if "图书在版" in line:
            continue

        # 检测 section 标题
        section_found = None
        section_content = ""

        # 匹配 【xxx】 格式
        m = re.match(r'^【(.+?)】\s*(.*)', line)
        if m:
            sec_name = m.group(1)
            for key, std_name in SECTION_HEADERS.items():
                if key in sec_name:
                    section_found = std_name
                    section_content = m.group(2)
                    break

        if section_found:
            if in_entry and current_name:
                # 保存上一个 section
                if current_text:
                    current_sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                current_section = section_found
                if section_content:
                    current_text.append(section_content)
            elif not in_entry and current_name:
                # 第一个 section，标志着药品条目开始
                in_entry = True
                current_sections = OrderedDict()
                current_section = section_found
                current_text = [section_content] if section_content else []
            continue

        # 可能是药名行（短行，没有 section 标记，以中文开头）
        # 药名通常在 section 之前的一两行
        if not in_entry and re.match(r'^[一-鿿]{2,20}$', line):
            # 可能是新药名
            if current_name and current_sections:
                # 保存上一个药品
                if current_text:
                    current_sections[current_section] = '\n'.join(current_text).strip()
                entries.append({
                    "name": current_name,
                    "sections": dict(current_sections),
                })
                current_text = []
                current_sections = OrderedDict()
            current_name = line
            in_entry = False
            current_section = "正文"
            current_text = []
            continue

        # 正文内容行
        if in_entry and current_name:
            current_text.append(line)

    # 最后一个药品
    if current_name and current_sections:
        if current_text:
            current_sections[current_section] = '\n'.join(current_text).strip()
        entries.append({
            "name": current_name,
            "sections": dict(current_sections),
        })

    doc.close()

    # 过滤无效条目（至少要有 2 个 section 或内容 > 100 字）
    valid = []
    for e in entries:
        content_len = sum(len(v) for v in e["sections"].values())
        if len(e["sections"]) >= 2 or content_len > 100:
            # 生成 drug_id
            hash_suffix = hashlib.md5(e["name"].encode()).hexdigest()[:6].upper()
            e["drug_id"] = f"Z2020-{hash_suffix}"
            e["category"] = "中药"
            e["subcategory"] = ""
            e["name_en"] = ""
            e["pinyin"] = ""
            e["source"] = {
                "version": "2020年版",
                "volume": "一部",
                "page": "",
            }
            valid.append(e)

    return valid


async def get_embeddings(texts: list[str], text_type: str = "document") -> list[list[float]]:
    async with httpx.AsyncClient(timeout=60) as client:
        resp = await client.post(
            EMBEDDING_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {QWEN_API_KEY}",
            },
            json={
                "model": EMBEDDING_MODEL,
                "input": {"texts": texts},
                "parameters": {"text_type": text_type},
            },
        )
        data = resp.json()
    if data.get("code") and data.get("code") != "" and data.get("code") is not None:
        raise RuntimeError(f"Embedding error: {data.get('message', data)}")
    return [item["embedding"] for item in data["output"]["embeddings"]]


async def ingest_entries(entries: list[dict]):
    """向量化 + 入库"""
    engine = create_async_engine(DB_URL)

    chunks = []
    chunk_meta = []

    for entry in entries:
        source = f"{entry['source']['version']} {entry['source']['volume']}"
        for section_key, section_text in entry["sections"].items():
            if not section_text or len(section_text.strip()) < 5:
                continue
            content = f"【{entry['name']} - {section_key}】\n{section_text}\n\n来源：{source}"
            chunks.append(content)
            chunk_meta.append({
                "drug_id": entry["drug_id"],
                "section": section_key,
                "content": content,
                "source": source,
            })

    print(f"  ✂️  {len(chunks)} chunks，向量化中...")

    all_vectors = []
    for i in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[i:i + BATCH_SIZE]
        vecs = await get_embeddings(batch, text_type="document")
        all_vectors.extend(vecs)
        n = min(i + BATCH_SIZE, len(chunks))
        print(f"    向量化: {n}/{len(chunks)}")

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

    # 写入 drug_chunks
    chunk_count = 0
    async with engine.begin() as conn:
        for idx, (meta, vec) in enumerate(zip(chunk_meta, all_vectors)):
            vec_str = f"[{','.join(str(v) for v in vec)}]"
            await conn.execute(
                text("""
                    INSERT INTO drug_chunks (drug_id, section, content, source, chunk_index, embedding, vec)
                    VALUES (:drug_id, :section, :content, :source, :chunk_index, :embedding, :vec)
                    ON CONFLICT DO NOTHING
                """),
                {
                    "drug_id": meta["drug_id"],
                    "section": meta["section"],
                    "content": meta["content"],
                    "source": meta["source"],
                    "chunk_index": idx,
                    "embedding": json.dumps(vec),
                    "vec": vec_str,
                },
            )
            chunk_count += 1

    # 写入 drugs 表
    drug_count = 0
    async with engine.begin() as conn:
        for entry in entries:
            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": entry["drug_id"],
                    "name": entry["name"],
                    "name_en": entry.get("name_en", ""),
                    "pinyin": entry.get("pinyin", ""),
                    "category": entry.get("category", ""),
                    "subcategory": entry.get("subcategory", ""),
                    "sections": json.dumps(entry.get("sections", {}), ensure_ascii=False),
                    "source_version": entry["source"]["version"],
                    "source_volume": entry["source"]["volume"],
                    "source_page": entry["source"].get("page", ""),
                },
            )
            drug_count += 1

    await engine.dispose()
    print(f"\n🎉 2020 药典入库完成！药品 {drug_count} 个，chunk {chunk_count} 条")


async def main():
    load_env()
    global QWEN_API_KEY
    QWEN_API_KEY = os.environ.get("QWEN_API_KEY", "")
    if not QWEN_API_KEY:
        print("❌ 未设置 QWEN_API_KEY")
        sys.exit(1)

    print("=" * 60)
    print("📖 解析 2020 年版药典 PDF...")
    entries = extract_drug_entries(PDF_PATH)
    print(f"   提取药品条目: {len(entries)}")
    print("=" * 60)

    if not entries:
        print("❌ 未提取到有效条目，请检查 PDF 解析逻辑")
        return

    # 打印前 5 个条目作为样品
    print("\n📋 前 5 个条目预览:")
    for e in entries[:5]:
        secs = list(e["sections"].keys())
        total_len = sum(len(v) for v in e["sections"].values())
        print(f"   {e['name']} | {len(secs)} sections ({total_len} 字) | {secs}")

    print(f"\n🚀 开始向量化入库...")
    await ingest_entries(entries)


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