AIyaodianzhushou/data-pipeline/crawlers/wiki_spider.py

"""
维基百科药品数据采集器
从中文维基百科采集药品条目，输出标准 JSON 格式

特性：
- 频率控制：每次请求间隔 3 秒（尊重维基百科服务器）
- 批次处理：每 20 个药品暂停 60 秒
- 断点续采：已采集的自动跳过
- 输出符合 drug_schema.json 规范
"""
import json
import re
import time
import sys
import os
from pathlib import Path
from urllib.parse import quote

import httpx
from bs4 import BeautifulSoup

# ============================================
# 配置
# ============================================
REQUEST_INTERVAL = 3.0       # 单次请求间隔（秒）
BATCH_SIZE = 20              # 每批采集数量
BATCH_PAUSE = 60              # 批次间暂停（秒）
WIKI_BASE = "https://zh.wikipedia.org/wiki/"
HEADERS = {
    "User-Agent": "PharmacopoeiaAI/1.0 (Research bot; contact@example.com) Mozilla/5.0",
    "Accept-Language": "zh-CN,zh;q=0.9",
}

# 匹配药典相关 section 标题
SECTION_PATTERNS = {
    "性状": re.compile(r"(性状|外观|物理性质|Description|Appearance)"),
    "适应症": re.compile(r"(适应症|适应征|主治|功能主治|Indications)"),
    "用法与用量": re.compile(r"(用法|用量|剂量|给药|Dosage|Administration)"),
    "禁忌": re.compile(r"(禁忌|禁用|Contraindications)"),
    "不良反应": re.compile(r"(不良反|副作用|Side.effect|Adverse)"),
    "注意事项": re.compile(r"(注意|警告|谨慎|Precautions|Warnings)"),
    "贮藏": re.compile(r"(贮藏|储存|保存|Storage)"),
    "药理": re.compile(r"(药理|药效|作用机制|Pharmacology|Mechanism)"),
}

# 内容长度限制（避免单 section 过长超出 LLM 上下文窗口）
MAX_OVERVIEW_LENGTH = 5000
MAX_SECTION_LENGTH = 8000


def fetch_wiki_page(drug_name: str, client: httpx.Client) -> tuple[int, str]:
    """获取维基百科页面"""
    url = WIKI_BASE + quote(drug_name)
    resp = client.get(url, headers=HEADERS, follow_redirects=True)
    # 检查是否重定向到非药品页面（如消歧义页）
    if resp.status_code == 200:
        text = resp.text
        if "维基百科目前还没有" in text or "Wikipedia does not have" in text:
            return 404, ""
    return resp.status_code, resp.text


def parse_drug_sections(html: str, drug_name: str) -> dict:
    """解析维基百科页面，提取药品相关信息"""
    soup = BeautifulSoup(html, "lxml")
    
    # 提取基本信息
    infobox = {}
    info_table = soup.find("table", class_="infobox")
    if info_table:
        for row in info_table.find_all("tr"):
            th = row.find("th")
            td = row.find("td")
            if th and td:
                key = th.get_text(strip=True)
                val = td.get_text(" ", strip=True)
                infobox[key] = val

    # 提取正文内容（跳过导航、引用等）
    content = soup.find("div", class_="mw-parser-output")
    if not content:
        return {"sections": {}, "infobox": infobox}

    sections = {}
    current_section = "概述"
    current_text = []
    overview_text = []

    # 提取开头概述（第一个标题之前的内容）
    for elem in content.children:
        # 找第一个 h2 之前的段落
        if elem.name in ("h2", "h3", "h4"):
            # 保存之前的内容
            if overview_text:
                sections["概述"] = " ".join(overview_text)[:MAX_OVERVIEW_LENGTH]
                overview_text = []
            break
        if elem.name == "p" or elem.name == "div":
            t = elem.get_text(" ", strip=True)
            if len(t) > 20:
                overview_text.append(t)

    if overview_text:
        sections["概述"] = " ".join(overview_text)[:MAX_OVERVIEW_LENGTH]

    # 提取各 section
    for elem in content.children:
        if elem.name in ("h2", "h3", "h4"):
            if current_text and current_section:
                matched = None
                for key, pattern in SECTION_PATTERNS.items():
                    if pattern.search(current_section):
                        matched = key
                        break
                if matched:
                    sections[matched] = " ".join(current_text)[:MAX_SECTION_LENGTH]
                else:
                    sections[current_section] = " ".join(current_text)[:MAX_SECTION_LENGTH]
            current_section = elem.get_text(strip=True).replace("[编辑]", "").strip()
            current_text = []
            continue

        if elem.name in ("p", "ul", "ol", "div"):
            t = elem.get_text(" ", strip=True)
            if len(t) > 10:
                current_text.append(t)

    # 保存最后一个 section
    if current_text and current_section:
        matched = None
        for key, pattern in SECTION_PATTERNS.items():
            if pattern.search(current_section):
                matched = key
                break
        if matched:
            sections[matched] = " ".join(current_text)[:MAX_SECTION_LENGTH]
        else:
            sections[current_section] = " ".join(current_text)[:MAX_SECTION_LENGTH]

    return {
        "sections": sections,
        "infobox": infobox,
    }


def wiki_to_drug_entry(drug_name: str, parsed: dict) -> dict:
    """转换为标准 drug entry JSON"""
    infobox = parsed.get("infobox", {})
    sections = parsed.get("sections", {})

    return {
        "drug_id": f"WIKI-{drug_name[:30]}",
        "name": drug_name,
        "name_en": infobox.get("其他名称", ""),
        "pinyin": "",
        "category": infer_category(drug_name, sections),
        "subcategory": "",
        "sections": sections,
        "source": {
            "version": "维基百科",
            "volume": f"https://zh.wikipedia.org/wiki/{quote(drug_name)}",
            "page": "线上条目",
        },
    }


def infer_category(name: str, sections: dict) -> str:
    """推断药品分类（基于名称和正文内容综合判断）"""
    all_text = " ".join(sections.values()) + name

    # 生物制品特征词（优先级最高）
    bio_markers = ["疫苗", "毒素", "抗血清", "免疫球蛋白", "单克隆抗体", "重组", "基因工程"]
    if any(k in all_text for k in bio_markers):
        return "生物制品"

    # 中药特征：多个特征同时出现才判定为中药，减少误判
    tcm_name_markers = ["丸", "散", "汤", "丹", "膏", "颗粒", "胶囊"]
    tcm_content_markers = ["中药", "本草", "性味", "归经", "炮制", "饮片"]
    tcm_herb_chars = ["草", "花", "叶", "根", "皮", "参", "芪", "苓", "术", "芍", "芷"]

    # 检查药品名是否含多个中药特征
    name_tcm_score = sum(1 for k in tcm_name_markers if k in name)
    name_tcm_score += sum(1 for c in tcm_herb_chars if c in name)

    # 检查正文是否含中药特征
    content_tcm = any(k in all_text for k in tcm_content_markers)

    # 名称含中药剂型 + 至少 1 个其他特征 → 中药
    if name_tcm_score >= 2 or (name_tcm_score >= 1 and content_tcm):
        return "中药"

    return "化学药"


def main():
    # 读取药品索引
    index_path = Path(__file__).resolve().parent.parent / "data" / "drug_index.json"
    with open(index_path, "r") as f:
        drug_list = json.load(f)

    # 输出目录
    output_dir = Path(__file__).resolve().parent.parent / "data" / "wiki_drugs"
    output_dir.mkdir(parents=True, exist_ok=True)

    # 已采集的跳过（断点续采）
    already = set(f.stem for f in output_dir.glob("*.json"))

    # 只采集化学药 + 生物制品（跳过中药、通则需要更复杂的解析）
    candidates = [d for d in drug_list if d["volume"] in (2, 3)]
    
    print(f"📋 候选药品: {len(candidates)} 个（二部化学药 + 三部生物制品）")
    print(f"   已采集: {len(already)} 个")
    print(f"   待采集: {len(candidates) - len(already)} 个")
    print(f"   频率控制: {REQUEST_INTERVAL}s/次, 每{BATCH_SIZE}次暂停{BATCH_PAUSE}s")
    print(f"   预计耗时: ~{(len(candidates) - len(already)) * (REQUEST_INTERVAL + 1) / 60:.0f} 分钟\n")

    client = httpx.Client(timeout=20, follow_redirects=True)
    count = 0
    success = 0
    failed = 0

    for drug in candidates:
        name = drug["name"]
        safe_name = name.replace("/", "_").replace(":", "_")

        if safe_name in already:
            continue

        count += 1
        status, html = fetch_wiki_page(name, client)
        time.sleep(REQUEST_INTERVAL)

        if status != 200 or not html:
            print(f"  [{count}] ❌ {name} (HTTP {status})")
            failed += 1
        else:
            try:
                parsed = parse_drug_sections(html, name)
                entry = wiki_to_drug_entry(name, parsed)

                # 至少有一个 section 才算有效
                sections_count = len(entry["sections"])
                if sections_count < 1:
                    print(f"  [{count}] ⚠️  {name} (无有效内容)")
                    failed += 1
                    continue

                out_path = output_dir / f"{safe_name}.json"
                with open(out_path, "w", encoding="utf-8") as f:
                    json.dump(entry, f, ensure_ascii=False, indent=2)

                print(f"  [{count}] ✅ {name} ({sections_count} sections)")
                success += 1
            except Exception as e:
                print(f"  [{count}] ❌ {name} (解析失败: {str(e)[:50]})")
                failed += 1

        # 批次间暂停
        if count % BATCH_SIZE == 0 and count > 0:
            print(f"\n⏸  已处理 {count} 个，暂停 {BATCH_PAUSE}s...\n")
            time.sleep(BATCH_PAUSE)

    client.close()
    print(f"\n{'='*50}")
    print(f"采集完成: 成功 {success}, 失败 {failed}, 总计 {count}")
    print(f"输出目录: {output_dir}")
    print(f"{'='*50}")


if __name__ == "__main__":
    main()