dieya_risk_prediction_backend/scripts/backfill_embeddings.py

"""Generate and store anomaly-case embeddings after pgvector is available.

Usage:
    .venv/Scripts/python -m scripts.backfill_embeddings

The script is intentionally separate from application startup. Embedding 2,448
rows can take time and should never make the API unavailable.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import random
from typing import Any

import httpx
from app.core.config import get_settings
from app.database.session import SessionFactory
from sqlalchemy import text

logger = logging.getLogger("dieya_risk_prediction.embedding_backfill")


def _case_text(row: dict[str, Any]) -> str:
    labels = (
        ("产品", row.get("product_name")),
        ("品类", row.get("product_category")),
        ("包材", row.get("material_type")),
        ("供应商", row.get("supplier")),
        ("加工厂", row.get("factory")),
        ("阶段", row.get("stage_new")),
        ("异常类型", row.get("anomaly_type")),
        ("异常描述", row.get("description")),
        ("原因", row.get("root_cause")),
        ("解决方案", row.get("solution")),
    )
    return "；".join(f"{label}：{value}" for label, value in labels if value)


async def _embedding(
    client: httpx.AsyncClient,
    endpoint: str,
    api_key: str,
    texts: list[str],
    model: str,
) -> list[list[float]]:
    for attempt in range(4):
        response = await client.post(
            endpoint,
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json={"model": model, "input": texts},
        )
        if response.status_code not in {429, 500, 502, 503, 504}:
            response.raise_for_status()
            data = response.json().get("data") or []
            ordered = sorted(data, key=lambda item: item.get("index", 0))
            return [list(map(float, item["embedding"])) for item in ordered]
        if attempt == 3:
            response.raise_for_status()
        retry_after = response.headers.get("retry-after")
        try:
            delay = min(60.0, max(3.0, float(retry_after))) if retry_after else 3.0 * (2**attempt)
        except ValueError:
            delay = 3.0 * (2**attempt)
        delay += random.uniform(0, 1)
        logger.warning(
            "Embedding API 限流或暂时不可用 status=%s attempt=%s delay=%.1fs",
            response.status_code,
            attempt + 1,
            delay,
        )
        await asyncio.sleep(delay)
    raise RuntimeError("Embedding 请求重试次数已用尽")


async def main(batch_size: int) -> None:
    settings = get_settings()
    api_key = (
        settings.rag_embedding_api_key.get_secret_value()
        if settings.rag_embedding_api_key
        else settings.llm_api_key.get_secret_value() if settings.llm_api_key else ""
    )
    base_url = settings.rag_embedding_base_url or settings.llm_base_url
    if not api_key or not base_url:
        raise RuntimeError("未配置 Embeddings API Key 或 Base URL")
    endpoint = f"{base_url.rstrip('/')}/embeddings"
    if not base_url.rstrip("/").endswith("/v1"):
        endpoint = f"{base_url.rstrip('/')}/v1/embeddings"

    async with SessionFactory() as session:
        column_exists = await session.scalar(
            text(
                """
                SELECT EXISTS (
                    SELECT 1 FROM information_schema.columns
                    WHERE table_name = 'anomaly_cases'
                      AND column_name = 'embedding_vector'
                      AND udt_name = 'halfvec'
                )
                """
            )
        )
        if not column_exists:
            raise RuntimeError("当前数据库没有 embedding_vector，请先安装 pgvector 并执行 alembic upgrade head")
        rows = (
            await session.execute(
                text(
                    """
                    SELECT id, product_name, product_category, material_type, supplier,
                           factory, stage_new, anomaly_type, description, root_cause, solution
                    FROM anomaly_cases
                    WHERE embedding_vector IS NULL
                    ORDER BY id
                    """
                )
            )
        ).mappings().all()

    logger.info("待生成案例向量数量=%s", len(rows))
    print(f"待生成案例向量数量：{len(rows)}，批次大小：{batch_size}", flush=True)
    timeout = httpx.Timeout(
        settings.rag_embedding_timeout_seconds,
        connect=min(5, settings.rag_embedding_timeout_seconds),
    )
    async with httpx.AsyncClient(timeout=timeout) as client:
        success_count = 0
        failure_count = 0
        for start in range(0, len(rows), batch_size):
            batch = rows[start : start + batch_size]
            try:
                vectors = await _embedding(
                    client,
                    endpoint,
                    api_key,
                    [_case_text(dict(row)) for row in batch],
                    settings.rag_embedding_model,
                )
            except (httpx.HTTPError, ValueError, TypeError, KeyError) as exc:
                failure_count += len(batch)
                logger.exception("案例向量生成失败 case_id=%s", batch[0]["id"])
                print(f"案例 {batch[0]['id']} 失败：{exc.__class__.__name__}，已跳过", flush=True)
                continue
            if len(vectors) != len(batch):
                raise RuntimeError("Embeddings 返回数量与请求数量不一致")
            async with SessionFactory() as session:
                for row, vector in zip(batch, vectors, strict=True):
                    literal = "[" + ",".join(str(value) for value in vector) + "]"
                    await session.execute(
                        text(
                            "UPDATE anomaly_cases "
                            "SET embedding_vector = CAST(:embedding AS halfvec) "
                            "WHERE id = :id"
                        ),
                        {"embedding": literal, "id": row["id"]},
                )
                await session.commit()
            success_count += len(batch)
            logger.info("已完成案例向量回填 %s/%s", min(start + batch_size, len(rows)), len(rows))
            print(
                f"已完成：{min(start + batch_size, len(rows))}/{len(rows)}，"
                f"成功：{success_count}，失败：{failure_count}，当前案例：{batch[-1]['id']}",
                flush=True,
            )
            await asyncio.sleep(1.5)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch-size", type=int, default=1)
    args = parser.parse_args()
    if not 1 <= args.batch_size <= 8:
        raise SystemExit("--batch-size 必须在 1 到 8 之间")
    asyncio.run(main(args.batch_size))