AIyaodianzhushou/data-pipeline/embed_only.py

#!/usr/bin/env python3
"""仅向量化：读取PG中无向量的chunk，批量embedding后回写"""
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

ROOT = Path(__file__).resolve().parent.parent
env_file = ROOT / ".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())

API_KEY = os.environ.get("QWEN_API_KEY", "")
EMBEDDING_URL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"

DB_PW = os.environ.get("POSTGRES_PASSWORD", "pharma2025")
DB_USER = os.environ.get("POSTGRES_USER", "postgres")
DB_HOST = os.environ.get("POSTGRES_HOST", "localhost")
DB_PORT = os.environ.get("POSTGRES_PORT", "5432")
DB_NAME = os.environ.get("POSTGRES_DB", "pharmacopoeia")
DB_URL = f"postgresql+asyncpg://{DB_USER}:{DB_PW}@{DB_HOST}:{DB_PORT}/{DB_NAME}"


async def embed_batch(texts: list[str]) -> list[list[float]]:
    async with httpx.AsyncClient(timeout=120) as client:
        resp = await client.post(
            EMBEDDING_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {API_KEY}",
            },
            json={
                "model": "text-embedding-v3",
                "input": {"texts": texts},
                "parameters": {"text_type": "document"},
            },
        )
        data = resp.json()

    if data.get("code") and data["code"] != "":
        raise RuntimeError(f"Embedding API error: {data.get('message', data)}")
    return [item["embedding"] for item in data["output"]["embeddings"]]


async def main():
    print("=" * 50)
    print("Embedding-only: 仅向量化 + 写回 PG")
    print("=" * 50)

    if not API_KEY:
        print("❌ QWEN_API_KEY not set")
        sys.exit(1)

    engine = create_async_engine(DB_URL)

    async with engine.begin() as conn:
        rows = (await conn.execute(text(
            "SELECT id, content FROM drug_chunks WHERE embedding IS NULL OR embedding::text = '[]' ORDER BY id"
        ))).fetchall()

    total = len(rows)
    print(f"\n📊 待处理 chunk: {total}")

    if total == 0:
        print("✅ 所有 chunk 已有向量，无需处理")
        await engine.dispose()
        return

    ids = [r[0] for r in rows]
    contents = [r[1] for r in rows]

    BATCH = 10
    count = 0

    async with engine.begin() as conn:
        for i in range(0, total, BATCH):
            batch_texts = contents[i : i + BATCH]
            batch_ids = ids[i : i + BATCH]

            vecs = await embed_batch(batch_texts)

            for cid, vec in zip(batch_ids, vecs):
                vec_str = json.dumps(vec)
                await conn.execute(
                    text("UPDATE drug_chunks SET embedding = CAST(:emb AS jsonb) WHERE id = :id"),
                    {"id": cid, "emb": vec_str},
                )

            count += len(batch_texts)
            print(f"  进度: {count}/{total}")

    await engine.dispose()
    print(f"\n🎉 向量化完成！{count} 个 chunk 已更新")


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