AIyaodianzhushou/data-pipeline/docx_ingest.py

"""
DOCX 药典批量导入脚本
遍历指定目录下所有 DOCX 文件 → 提取药名和 sections → 向量化 → 写入 PostgreSQL
药品名以文件名（去 .docx）为准
"""
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": "四部"}

# 四部子目录 → 具体分类
CATEGORY_OVERRIDE = {
    "药用辅料": "药用辅料",
    "通用技术要求": "通则",
    "指导原则": "通则",
    "通用技术要求/指导原则": "通则",
}


def load_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 parse_docx(filepath: str) -> dict | None:
    """
    解析单个 DOCX 文件为药典条目。
    药品名 = 文件名（去掉 .docx）
    凡例/纲要类：名称加上部别前缀
    """
    filename = Path(filepath).stem
    volume = get_volume_group(filepath)
    parent_dir = Path(filepath).parent.name

    # 凡例/纲要/通则类：名称格式为"一部-凡例"以区分
    is_meta = parent_dir in ("凡例", "纲要", "通用技术要求", "指导原则")
    display_name = f"{volume}-{filename}" if is_meta else filename

    try:
        doc = Document(filepath)
    except Exception:
        return None

    pinyin = ""
    sections = OrderedDict()
    current_section = "正文"
    current_text = []

    # 收集所有段落
    for p in doc.paragraphs:
        text = p.text.strip()
        if not text or text.isspace():
            continue

        style = p.style.name if p.style else ""

        # Heading 1 通常是药品名（与文件名一致），直接跳过，不加入 sections
        if 'Heading 1' in style:
            continue

        # 拼音行（紧跟药名后）
        if not pinyin and is_pinyin_line(text):
            pinyin = text.strip()
            continue

        # 独立的 section 标题（2-6 字短行）
        section_matched = detect_section_header(text)
        if section_matched and len(text) <= 10:
            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 and detect_section_header(m.group(1)):
            if current_text:
                sections[current_section] = '\n'.join(current_text).strip()
                current_text = []
            current_section = detect_section_header(m.group(1))
            if m.group(2):
                current_text.append(m.group(2))
            continue

        # 普通段落
        current_text.append(text)

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

    # ============================================
    # 凡例/纲要特殊处理：按编号拆分（一、二、三... 或 1. 2. 3.）
    # ============================================
    if parent_dir in ("凡例", "纲要") and "正文" in sections:
        sections = split_numbered_content(sections["正文"])

    if len(sections) == 0:
        return None

    category = infer_category(filepath, parent_dir)

    return {
        "drug_id": generate_drug_id(display_name, filepath),
        "name": display_name,
        "name_en": "",
        "pinyin": pinyin,
        "category": category,
        "subcategory": infer_subcategory(filepath),
        "sections": dict(sections),
        "source": {
            "version": "2025年版",
            "volume": volume,
            "page": "",
        },
    }


def split_numbered_content(text: str) -> OrderedDict:
    """
    凡例/纲要类内容按编号拆分为独立 section。
    匹配模式：
      - 一、二、三、... 十四、
      - （一）（二）...
      - 1. 2. 3. 或 1）2）3）
    """
    # 中文数字
    cn_nums = "一二三四五六七八九十"
    patterns = [
        # ^一、或 \n一、 开头
        re.compile(r'(?:^|\n)([' + cn_nums + r']{1,2})[、，,]'),
        # （一）开头
        re.compile(r'(?:^|\n)（([' + cn_nums + r']{1,2})）'),
        # 数字编号
        re.compile(r'(?:^|\n)(\d{1,2})[\.\)）]'),
    ]

    # 用第一个模式尝试
    result = OrderedDict()
    pattern = patterns[0]
    matches = list(pattern.finditer(text))

    if len(matches) < 2:
        # 尝试第二个模式
        pattern = patterns[1]
        matches = list(pattern.finditer(text))

    if len(matches) < 2:
        # 尝试第三个模式
        pattern = patterns[2]
        matches = list(pattern.finditer(text))

    if len(matches) < 2:
        # 确实没有编号结构，保留原样但限制长度
        if len(text) > 3000:
            # 按段落粗略切分
            paras = [p.strip() for p in text.split('\n') if p.strip()]
            buf = ""
            idx = 0
            for p in paras:
                if len(buf) + len(p) > 2000 and buf:
                    result[f"第{idx + 1}段"] = buf
                    idx += 1
                    buf = p
                else:
                    buf += ("\n" if buf else "") + p
            if buf:
                result[f"第{idx + 1}段"] = buf
            return result if result else OrderedDict([("正文", text)])
        return OrderedDict([("正文", text)])

    # 按编号切分
    for i, match in enumerate(matches):
        start = match.start()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        content = text[start:end].strip()
        # 编号字符串（去除换行和括号）
        raw_num = match.group(0).lstrip('\n').strip()
        section_name = f"凡例-{raw_num}"
        if len(content) > 50:
            result[section_name] = content

    return result if result else OrderedDict([("正文", text)])


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:
    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, parent_dir: str = "") -> str:
    path_lower = filepath.lower()
    # 凡例/纲要/通用技术要求
    if parent_dir in ("凡例", "纲要"):
        return "凡例与纲要"
    if parent_dir in ("通用技术要求", "指导原则"):
        return "通则"
    if "output2" in path_lower: return "化学药"
    if "output3" in path_lower: return "生物制品"
    if "output4" in path_lower: return _infer_category_vol4(filepath)
    return "中药"


def _infer_category_vol4(filepath: str) -> str:
    """四部特殊处理"""
    for key, cat in CATEGORY_OVERRIDE.items():
        if key in filepath:
            return cat
    return "通则"


def infer_subcategory(filepath: str) -> str:
    """从目录层级推断子分类"""
    parts = Path(filepath).parts
    # 找到 output/output2/output3/output4 之后的子目录名
    for i, p in enumerate(parts):
        if p in ("output", "output2", "output3", "output4"):
            if i + 1 < len(parts):
                sub = parts[i + 1]
                # 排除数字编号前缀
                return sub
    return ""


def generate_drug_id(name: str, filepath: str) -> str:
    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]]:
    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):
    """批量向量化 + 入库"""
    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,
                        subcategory = EXCLUDED.subcategory,
                        source_version = EXCLUDED.source_version,
                        source_volume = EXCLUDED.source_volume,
                        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 dedup_old_versions(engine):
    """覆盖性新增：同名药品如果存在非 2025 版的旧记录，标记为 inactive"""
    async with engine.begin() as conn:
        result = await conn.execute(
            text("""
                UPDATE drugs
                SET is_active = FALSE, updated_at = NOW()
                WHERE is_active = TRUE
                  AND source_version != '2025年版'
                  AND name IN (
                      SELECT name FROM drugs
                      WHERE source_version = '2025年版' AND is_active = TRUE
                  )
            """)
        )
        # asyncpg 不支持 rowcount，用另一种方式
        count = getattr(result, 'rowcount', 0)
        if count > 0:
            print(f"  🧹 去重: 标记 {count} 个旧版药品为 inactive")


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 = []
    skipped = 0
    by_category = {}
    for i, fp in enumerate(files):
        entry = parse_docx(fp)
        if entry:
            entries.append(entry)
            cat = entry["category"]
            by_category[cat] = by_category.get(cat, 0) + 1
        else:
            skipped += 1
        if (i + 1) % 1000 == 0:
            print(f"   解析进度: {i + 1}/{len(files)} (有效: {len(entries)})")

    print(f"\n📦 解析完成: {len(entries)} 个有效条目 (跳过 {skipped} 个)")
    print("   分类分布:")
    for cat, cnt in sorted(by_category.items(), key=lambda x: -x[1]):
        print(f"     {cat}: {cnt}")

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

    # 连接数据库
    engine = create_async_engine(DB_URL)

    # 分批入库
    total_drugs = 0
    total_chunks = 0
    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']}")

    # 去重：标记同名旧版药品为 inactive
    print("\n🧹 清理旧版重复数据...")
    await dedup_old_versions(engine)

    await engine.dispose()

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


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