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 docx.oxml.ns import qn
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 _extract_table_text(table) -> str:
    """将 python-docx Table 对象转换为可读文本。"""
    rows = []
    for row in table.rows:
        cells = [cell.text.strip() for cell in row.cells]
        # 跳过全空行
        if any(c for c in cells):
            rows.append(" | ".join(cells))
    return "\n".join(rows)


def parse_docx(filepath: str) -> dict | None:
    """
    解析单个 DOCX 文件为药典条目。
    药品名 = 文件名（去掉 .docx）
    凡例/纲要类：名称加上部别前缀
    按文档顺序遍历段落和表格，表格内容归入当前 section。
    """
    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 = []

    # 遍历文档 body 中的所有元素（按文档顺序，同时处理段落和表格）
    body = doc.element.body
    for child in body:
        # --- 表格 ---
        if child.tag == qn("w:tbl"):
            # 找到对应的 python-docx Table 对象
            tbl = None
            for t in doc.tables:
                if t._element is child:
                    tbl = t
                    break
            if tbl is not None:
                tbl_text = _extract_table_text(tbl)
                if tbl_text:
                    current_text.append(f"[表格]\n{tbl_text}")
            continue

        # --- 段落 ---
        if child.tag != qn("w:p"):
            continue

        # 从段落 XML 中提取文本和样式
        p_text = ""
        p_style = ""
        for p_child in child:
            if p_child.tag == qn("w:pPr"):
                for style_child in p_child:
                    if style_child.tag == qn("w:pStyle"):
                        p_style = style_child.get(qn("w:val"), "")
            elif p_child.tag == qn("w:r"):
                for r_child in p_child:
                    if r_child.tag == qn("w:t"):
                        t = r_child.text
                        if t:
                            p_text += t

        text = p_text.strip()
        if not text or text.isspace():
            continue

        # Heading 1 通常是药品名（与文件名一致），直接跳过，不加入 sections
        if p_style and (p_style.startswith("Heading") or "Heading" in p_style):
            continue

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

        # 独立的 section 标题（短行，阈值放宽以适应较长标题）
        section_matched = detect_section_header(text)
        if section_matched and len(text) <= 15:
            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一、 开头（支持"二十一"等复合数字，最大3字）
        re.compile(r'(?:^|\n)([' + cn_nums + r']{1,3})[、，,]'),
        # （一）开头
        re.compile(r'(?:^|\n)（([' + cn_nums + r']{1,3})）'),
        # 数字编号
        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:
    """检测文本是否为拼音行。

    判断逻辑：
    1. 字母占比 > 30% 且长度在合理范围
    2. 不含常见英文单词（如 drug, tablet, injection 等）
    3. 不含化学元素符号（如 Na, K, HCl 等大写字母组合）
    4. 拼音通常全小写或带声调符号
    """
    if len(text) <= 3 or len(text) >= 200:
        return False

    # 检查是否包含英文单词（排除化学药英文名）
    english_indicators = [
        "tablet", "injection", "capsule", "solution", "sodium",
        "chloride", "hydrochloride", "acid", "oxide", "sulfate",
        "phosphate", "acetate", "extract", "powder", "granule",
        "oral", "intravenous", "topical", "ophthalmic",
        "for", "and", "the", "with", "mg", "ml", "μg",
        "Tablet", "Injection", "Capsule", "Solution",
    ]
    text_lower = text.lower()
    if any(ind in text_lower for ind in english_indicators):
        return False

    # 检查是否含大写字母组合（化学式如 HCl, NaOH 等）
    if re.search(r'[A-Z]{2,}', text):
        return False

    # 计算字母占比
    alpha_count = sum(1 for c in text if c.isascii() and c.isalpha())
    total = len(text.replace(" ", ""))
    if total == 0:
        return False
    alpha_ratio = alpha_count / total

    # 拼音行通常字母占比 > 50% 且以空格分词
    if alpha_ratio < 0.5:
        return False

    # 检查是否有拼音声调标记（āáǎàēéěèīíǐìōóǒòūúǔùǖǘǚǜ）
    tone_marks = set("āáǎàēéěèīíǐìōóǒòūúǔùǖǘǚǜ")
    has_tones = any(c in tone_marks for c in text)

    # 拼音行特征：全小写字母 + 空格分词 + 可能带声调
    # 单词长度短（拼音音节通常 1-6 个字母）
    words = text.split()
    if len(words) >= 2 and all(len(w) <= 8 for w in words):
        if has_tones:
            return True
        # 无明确声调时提高阈值，减少误判
        return alpha_ratio > 0.7

    return has_tones


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())