AIyaodianzhushou/data-pipeline/ingest.py

"""
DashScope 文本嵌入 + PG 入库脚本
将药典条目 chunk → 向量化 → 写入 pgvector
"""
import json
import asyncio
import os
import sys
from pathlib import Path

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

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "backend"))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

os.chdir(Path(__file__).resolve().parent.parent)

def load_env():
    env_file = Path(__file__).resolve().parent.parent / ".env"
    if env_file.exists():
        for line in open(env_file):
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, _, val = line.partition("=")
                os.environ.setdefault(key.strip(), val.strip())

load_env()

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"

# 从环境变量读取数据库连接信息
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}"

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 {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") is not None:
        raise RuntimeError(f"Embedding API error: {data.get('message', data)}")

    # text-embedding-v3 返回结构: output.embeddings[].embedding
    return [item["embedding"] for item in data["output"]["embeddings"]]


async def ingest_drugs(json_path: str):
    with open(json_path, "r", encoding="utf-8") as f:
        drugs = json.load(f)

    engine = create_async_engine(DB_URL)
    chunks_to_embed = []
    chunk_rows = []

    print(f"📦 加载 {len(drugs)} 个药品条目")

    for drug in drugs:
        source = f"{drug['source']['version']} {drug['source']['volume']} P{drug['source']['page']}"
        for section_key, section_text in drug["sections"].items():
            if not section_text or len(section_text.strip()) < 10:
                continue
            content = f"【{drug['name']} - {section_key}】\n{section_text}\n\n来源：{source}"
            chunks_to_embed.append(content)
            chunk_rows.append({
                "drug_id": drug["drug_id"],
                "section": section_key,
                "content": content,
                "source": source,
                "drug_name": drug["name"],
                "category": drug["category"],
            })

    print(f"✂️  切出 {len(chunks_to_embed)} 个 chunk，正在向量化...")

    # 批量向量化（每批 10 个）
    all_vectors = []
    for i in range(0, len(chunks_to_embed), 10):
        batch = chunks_to_embed[i:i+10]
        vectors = await get_embeddings(batch, text_type="document")
        all_vectors.extend(vectors)
        n = min(i + 10, len(chunks_to_embed))
        print(f"  向量化进度: {n}/{len(chunks_to_embed)}")

    print(f"✅ 向量化完成，维度: {len(all_vectors[0]) if all_vectors else 'N/A'}")

    # 写入 PG
    count = 0
    async with engine.begin() as conn:
        for idx, (row, vec) in enumerate(zip(chunk_rows, all_vectors)):
            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)
                """),
                {
                    "drug_id": row["drug_id"],
                    "section": row["section"],
                    "content": row["content"],
                    "source": row["source"],
                    "chunk_index": idx,
                    "embedding": json.dumps(vec),
                    "vec": f"[{','.join(str(v) for v in vec)}]",
                },
            )
            count += 1

    # 同时写入 drugs 表
    async with engine.begin() as conn:
        for drug in drugs:
            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": drug["drug_id"],
                    "name": drug["name"],
                    "name_en": drug.get("name_en"),
                    "pinyin": drug.get("pinyin"),
                    "category": drug.get("category"),
                    "subcategory": drug.get("subcategory"),
                    "sections": json.dumps(drug.get("sections", {}), ensure_ascii=False),
                    "source_version": drug["source"]["version"],
                    "source_volume": drug["source"]["volume"],
                    "source_page": drug["source"]["page"],
                },
            )

    await engine.dispose()
    print(f"\n🎉 入库完成！药品 {len(drugs)} 个，chunk {count} 条")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        json_path = sys.argv[1]
    else:
        json_path = os.path.join(os.path.dirname(__file__), "data", "sample_drugs.json")

    if not os.path.exists(json_path):
        print(f"❌ 找不到数据文件: {json_path}")
        sys.exit(1)
    asyncio.run(ingest_drugs(json_path))