192 lines
7.8 KiB
Python
192 lines
7.8 KiB
Python
"""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("canonical_supplier") or 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):
|
||
try:
|
||
response = await client.post(
|
||
endpoint,
|
||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||
json={"model": model, "input": texts},
|
||
)
|
||
except httpx.RequestError:
|
||
if attempt == 3:
|
||
raise
|
||
delay = 2.0 * (2**attempt) + random.uniform(0, 1)
|
||
logger.warning("Embedding 网络异常 attempt=%s delay=%.1fs", attempt + 1, delay)
|
||
await asyncio.sleep(delay)
|
||
continue
|
||
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, *, force: bool = False) -> 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 ac.id, ac.product_name, ac.product_category, ac.material_type, ac.supplier,
|
||
ac.factory, ac.stage_new, ac.anomaly_type, ac.description, ac.root_cause, ac.solution,
|
||
string_agg(DISTINCT sm.canonical_name, ' / ') AS canonical_supplier
|
||
FROM anomaly_cases ac
|
||
LEFT JOIN anomaly_case_suppliers acs ON acs.case_id = ac.id
|
||
LEFT JOIN supplier_master sm ON sm.id = acs.supplier_id
|
||
WHERE (:force OR ac.embedding_vector IS NULL)
|
||
GROUP BY ac.id
|
||
ORDER BY id
|
||
"""
|
||
),
|
||
{"force": force},
|
||
)
|
||
).mappings().all()
|
||
|
||
logger.info("待生成案例向量数量=%s", len(rows))
|
||
print(f"待生成案例向量数量：{len(rows)}，批次大小：{batch_size}", flush=True)
|
||
timeout = httpx.Timeout(
|
||
max(30, settings.rag_embedding_timeout_seconds),
|
||
connect=min(10, max(30, 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,
|
||
)
|
||
# The provider already returns 429 with Retry-After when throttled;
|
||
# keep only a small pacing gap between successful batches.
|
||
await asyncio.sleep(0.2)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--batch-size", type=int, default=1)
|
||
parser.add_argument("--force", action="store_true", help="recompute existing vectors after source normalization")
|
||
args = parser.parse_args()
|
||
if not 1 <= args.batch_size <= 64:
|
||
raise SystemExit("--batch-size 必须在 1 到 64 之间")
|
||
asyncio.run(main(args.batch_size, force=args.force))
|