AIyaodianzhushou/data-pipeline/docx_ingest.py

"""
DOCX 药典批量导入脚本
遍历指定目录下所有 DOCX 文件 → 提取药名和 sections → 向量化 → 写入 PostgreSQL
"""
import json
import asyncio
import os
import sys
import re
import hashlib
from pathlib import Path
from collections import OrderedDict

import httpx
from docx import Document
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text

# ============================================
# 配置
# ============================================
SOURCE_DIR = os.environ.get("DOCX_SOURCE_DIR", "/opt/2025")
SKIP_DIRS = {"凡例", "通则", "凡例与通则", "附录", "索引", "目录"}

# 从环境变量读取
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  # 每批向量化的文本数

# Section 头部关键词 → 标准名称映射
SECTION_HEADERS = OrderedDict([
    ("处方", "处方"),
    ("制法", "制法"),
    ("性状", "性状"),
    ("鉴别", "鉴别"),
    ("检查", "检查"),
    ("浸出物", "浸出物"),
    ("含量测定", "含量测定"),
    ("含量", "含量测定"),
    ("功能与主治", "功能主治"),
    ("功能", "功能主治"),
    ("主治", "功能主治"),
    ("用法与用量", "用法用量"),
    ("用法", "用法用量"),
    ("用量", "用法用量"),
    ("注意", "注意事项"),
    ("注意事项", "注意事项"),
    ("规格", "规格"),
    ("贮藏", "贮藏"),
    ("类别", "类别"),
    ("制剂", "制剂"),
    ("附注", "附注"),
])

# 药典部别推断
VOLUME_MAP = {"output": "一部", "output2": "二部", "output3": "三部", "output4": "四部"}


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


def normalize_name(name: str) -> str:
    """清理药名，去掉拼音行、空格等"""
    # 移除拼音行（包含英文字母的行）
    lines = name.split('\n')
    result = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        # 跳过纯拼音行（大量英文字母）
        if sum(1 for c in line if c.isascii() and c.isalpha()) > len(line) * 0.4:
            continue
        # 跳过纯数字行
        if line.replace('.', '').replace(' ', '').isdigit():
            continue
        result.append(line)
    return result[0] if result else name.strip()


def parse_docx(filepath: str) -> dict | None:
    """解析单个 DOCX 文件为药典条目"""
    try:
        doc = Document(filepath)
    except Exception:
        return None

    drug_name = None
    pinyin = ""
    sections = OrderedDict()
    current_section = "正文"  # 第一个 heading 之前的内容
    current_text = []

    # 检测并提取 section 头部的正则
    section_pattern = re.compile(
        r'^【(.+?)】'  # 【性状】
        r'|^(\S{2,4})$'  # 性状, 鉴别 等单独的 section header
    )

    # 第一次 pass：收集所有段落文本
    all_paras = []
    for p in doc.paragraphs:
        text = p.text.strip()
        if not text or text.isspace():
            continue
        all_paras.append((p.style.name, text))

    for style, text in all_paras:
        # Heading 1 = 药品名
        if 'Heading 1' in style and drug_name is None:
            drug_name = normalize_name(text)
            continue

        # 检测拼音行（紧跟药名之后，单独一段包含字母）
        if drug_name and not pinyin and is_pinyin_line(text):
            pinyin = text.strip()
            continue

        # 检测 section 头部
        section_matched = detect_section_header(text)
        if section_matched and len(text) <= 10:
            # 保存上一个 section
            if current_text:
                sections[current_section] = '\n'.join(current_text).strip()
                current_text = []
            current_section = section_matched
            continue

        # 检测行内 section 头部：【性状】xxx 格式
        m = re.match(r'^【(.+?)】\s*(.*)', text)
        if m:
            header = m.group(1)
            rest = m.group(2)
            # 看看 header 是否匹配已知 section
            matched = detect_section_header(f"【{header}】")
            if matched:
                if current_text:
                    sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                current_section = matched
                if rest:
                    current_text.append(rest)
                continue

        # 普通段落，追加到当前 section
        current_text.append(text)

    # 最后一个 section
    if current_text:
        sections[current_section] = '\n'.join(current_text).strip()

    if not drug_name or len(sections) == 0:
        return None

    # 推断分类和部别
    parent_dir = Path(filepath).parent.name
    volume = get_volume_group(filepath)

    return {
        "drug_id": generate_drug_id(drug_name, filepath),
        "name": drug_name,
        "name_en": "",
        "pinyin": pinyin,
        "category": infer_category(filepath),
        "subcategory": parent_dir if parent_dir not in ("output", "output2", "output3", "output4") else "",
        "sections": dict(sections),
        "source": {
            "version": "2025年版",
            "volume": volume,
            "page": "",
        },
    }


def is_pinyin_line(text: str) -> bool:
    """判断是否为拼音行"""
    alpha_count = sum(1 for c in text if c.isascii() and c.isalpha())
    return alpha_count > len(text) * 0.3 and len(text) > 3 and len(text) < 200


def detect_section_header(text: str) -> str | None:
    """检测是否为 section 标题行，返回标准 section 名称"""
    clean = text.replace("【", "").replace("】", "").strip()
    if clean in SECTION_HEADERS:
        return SECTION_HEADERS[clean]
    return None


def get_volume_group(filepath: str) -> str:
    """推断药典部别"""
    for key, vol in VOLUME_MAP.items():
        if key in filepath:
            return vol
    return "一部"


def infer_category(filepath: str) -> str:
    """推断药品分类"""
    path_lower = filepath.lower()
    if "output2" in path_lower:
        return "化学药"
    if "output3" in path_lower:
        return "生物制品"
    return "中药"


def generate_drug_id(name: str, filepath: str) -> str:
    """生成 drug_id"""
    vol = get_volume_group(filepath)
    prefix = {"一部": "Z", "二部": "H", "三部": "S", "四部": "T"}.get(vol, "Z")
    hash_suffix = hashlib.md5(name.encode()).hexdigest()[:6].upper()
    return f"{prefix}2025-{hash_suffix}"


def find_docx_files(root_dir: str) -> list[str]:
    """扫描所有 DOCX 文件，跳过非药品目录"""
    files = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        # 跳过非药品目录
        dir_basename = os.path.basename(dirpath)
        if dir_basename in SKIP_DIRS:
            dirnames.clear()
            continue

        # 跳过临时文件
        for f in sorted(filenames):
            if f.startswith('~') or f.startswith('.'):
                continue
            if f.endswith('.docx'):
                files.append(os.path.join(dirpath, f))
    return files


async def get_embeddings(texts: list[str], text_type: str = "document") -> list[list[float]]:
    """调用 Qwen Embedding API"""
    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_docx_entries(entries: list[dict], engine, start_idx: int = 0):
    """批量向量化 + 入库"""
    # 构建 chunk 文本
    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": start_idx + idx,
                    "embedding": json.dumps(vec),
                    "vec": vec_str,
                },
            )
            chunk_count += 1

    # 写入 drugs 表
    drug_count = 0
    async with engine.begin() as conn:
        seen = set()
        for entry in entries:
            if entry["drug_id"] in seen:
                continue
            seen.add(entry["drug_id"])
            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,
                        category = EXCLUDED.category,
                        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

    return drug_count, 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，请在 .env 文件中配置")
        sys.exit(1)

    print("=" * 60)
    print("📂 扫描 DOCX 文件...")
    files = find_docx_files(SOURCE_DIR)
    print(f"   发现 {len(files)} 个 DOCX 文件")
    print("=" * 60)

    # 解析所有文件
    entries = []
    parse_errors = 0
    for i, fp in enumerate(files):
        entry = parse_docx(fp)
        if entry:
            entries.append(entry)

        if (i + 1) % 500 == 0:
            print(f"   解析进度: {i + 1}/{len(files)} (有效: {len(entries)})")

    print(f"\n📦 解析完成: {len(entries)} 个有效药品条目 (跳过 {len(files) - len(entries)} 个)")

    if not entries:
        print("❌ 没有有效数据")
        return

    # 连接数据库，分批入库
    engine = create_async_engine(DB_URL)
    total_drugs = 0
    total_chunks = 0

    # 分批处理（每批 50 个药品，避免 Qwen API 超时）
    BATCH = 50
    for i in range(0, len(entries), BATCH):
        batch = entries[i:i + BATCH]
        batch_no = i // BATCH + 1
        total_batches = (len(entries) + BATCH - 1) // BATCH
        print(f"\n🚀 第 {batch_no}/{total_batches} 批 ({len(batch)} 个药品) ...")
        try:
            dc, cc = await ingest_docx_entries(batch, engine, start_idx=total_chunks)
            total_drugs += dc
            total_chunks += cc
            print(f"   ✅ 入库: {dc} 药品, {cc} chunks")
        except Exception as e:
            print(f"   ❌ 批次失败: {e}")
            # 打印第一个药品名用于调试
            if batch:
                print(f"      首个药品: {batch[0]['name']}")

    await engine.dispose()

    print("\n" + "=" * 60)
    print(f"🎉 全部完成！药品 {total_drugs} 个, chunks {total_chunks} 条")
    print("=" * 60)


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