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 is_cjk_name(s):
    """严格药名检测：2-6 个纯 CJK 字符"""
    if not s or len(s) < 2 or len(s) > 6:
        return False
    cjk = sum(1 for c in s if ord(c) >= 0x4E00 and ord(c) <= 0x9FFF)
    return cjk >= 2 and cjk == len(s)


def is_pinyin_line(s):
    """拼音行：纯拉丁字母空格，首字母大写"""
    if not s or len(s) < 3 or len(s) > 80:
        return False
    alpha_spaces = sum(1 for c in s if c.isalpha() or c.isspace())
    return alpha_spaces == len(s) and any(c.isalpha() for c in s) and s[0].isupper()


def extract_drug_entries(pdf_path: str) -> list[dict]:
    """解析策略：扫描【性状】锚点，向前找药名。
    药名特征：2-6 纯 CJK + 下一行是拼音（首字母大写）或全大写拉丁名。
    只有通过拼音/拉丁名确认的才作为药名候选。
    """
    doc = fitz.open(pdf_path)
    print(f"PDF: {doc.page_count} pages")

    # 跳过前言（前 50 页）
    all_text = ""
    for i in range(32, doc.page_count):
        all_text += doc[i].get_text() + "\n"
    doc.close()

    lines = all_text.split('\n')
    clean_lines = []
    for line in lines:
        s = line.strip()
        if not s: continue
        if s.startswith("中国药典") and len(s) < 30: continue
        if "ISBN" in s or "CIP" in s: continue
        if re.match(r'^[・•·]\s*\d+\s*[・•·]$', s): continue
        clean_lines.append(s)

    # 第一遍：标记所有"已确认的药名"——CJK + 下一行是拼音或拉丁名
    confirmed_names = set()
    for i in range(len(clean_lines) - 2):
        s = clean_lines[i]
        nxt = clean_lines[i + 1]
        if is_cjk_name(s):
            if is_pinyin_line(nxt) or (nxt.isupper() and len(nxt) > 5):
                confirmed_names.add(s)

    print(f"   确认药名: {len(confirmed_names)} 个")

    # 第二遍：以已确认药名为锚点，提取药品条目
    entries = []
    current_drug = None
    current_sections = OrderedDict()
    current_section = "正文"
    current_text = []
    in_drug = False

    for s in clean_lines:
        # 检测 section
        m = re.match(r'^【(.+?)】\s*(.*)', s)
        if m:
            sec_raw = m.group(1)
            sec_content = m.group(2)
            section_found = None
            for key, std_name in SECTION_HEADERS.items():
                if key in sec_raw:
                    section_found = std_name
                    break

            if section_found and in_drug:
                if current_text:
                    current_sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                current_section = section_found
                if sec_content:
                    current_text.append(sec_content)
                continue

        # 检测已确认的药名
        if s in confirmed_names:
            # 保存上一个药品
            if in_drug and current_drug and current_sections:
                if current_text:
                    current_sections[current_section] = '\n'.join(current_text).strip()
                entries.append({"name": current_drug, "sections": dict(current_sections)})
            # 开始新药品
            current_drug = s
            in_drug = True
            current_sections = OrderedDict()
            current_section = "正文"
            current_text = []
            # 描述段落（药名和【性状】之间的文本）归入"正文"
            continue

        # 跳过拼音行和拉丁名行
        if in_drug and not current_sections and (is_pinyin_line(s) or (s.isupper() and len(s) > 5)):
            continue

        # 正文内容
        if in_drug:
            current_text.append(s)

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

    print(f"   初步提取: {len(entries)} 个条目")

    # 去重：同一药名保留内容最多的条目
    best = {}
    for e in entries:
        content_len = sum(len(v) for v in e["sections"].values())
        if len(e["sections"]) >= 1 and content_len > 50:
            name = e["name"]
            if name not in best or content_len > sum(len(v) for v in best[name]["sections"].values()):
                hash_suffix = hashlib.md5(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": ""}
                best[name] = e

    return list(best.values())


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"  [OK] 向量化完成, 维度={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[DONE] 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())