shaguabijia-app-server/app/services/comparison_llm_backfill.py

"""Persist and repair comparison-record LLM token costs."""
from __future__ import annotations

import logging
import time
from datetime import UTC, datetime, timedelta

from sqlalchemy import select

from app.core.config import settings
from app.core.rewards import CN_TZ
from app.db.session import SessionLocal
from app.models.app_config import AppConfig
from app.models.comparison import ComparisonRecord
from app.services.llm_cost import compute_llm_cost, get_llm_prices
from app.services.pricebot_llm_calls import fetch_llm_calls

logger = logging.getLogger("shagua.comparison_llm_backfill")


def _utc_to_beijing_naive(value: datetime) -> datetime:
    """Convert a DB UTC timestamp to comparison_record's Beijing wall-clock."""
    if value.tzinfo is None:
        value = value.replace(tzinfo=UTC)
    return value.astimezone(CN_TZ).replace(tzinfo=None)


def _store_calls(record_id: int, trace_id: str, calls: list[dict]) -> bool:
    """Store calls and all derived fields atomically."""
    with SessionLocal() as db:
        rec = db.get(ComparisonRecord, record_id)
        if rec is None or rec.trace_id != trace_id:
            logger.warning(
                "LLM cost backfill record mismatch record_id=%s trace=%s",
                record_id,
                trace_id,
            )
            return False

        # Never recalculate a frozen historical cost with a newer price config.
        if rec.llm_cost_yuan is not None and rec.llm_calls:
            return False

        rec.llm_calls = calls
        rec.llm_call_count = len(calls)
        rec.retry_count = sum(1 for call in calls if call.get("error"))
        rec.input_tokens = sum(
            (call.get("usage") or {}).get("prompt_tokens") or 0 for call in calls
        )
        rec.output_tokens = sum(
            (call.get("usage") or {}).get("completion_tokens") or 0 for call in calls
        )
        rec.llm_cost_yuan, rec.llm_price_snapshot = compute_llm_cost(
            calls, get_llm_prices(db)
        )
        db.commit()
        logger.info(
            "LLM cost backfilled trace=%s calls=%d input_tokens=%d "
            "output_tokens=%d cost=%s",
            trace_id,
            len(calls),
            rec.input_tokens,
            rec.output_tokens,
            rec.llm_cost_yuan,
        )
        return True


def backfill_comparison_llm_cost(
    record_id: int,
    trace_id: str,
    *,
    attempts: int = 3,
    retry_delays: tuple[float, ...] = (1.0, 3.0),
) -> bool:
    """Fetch and persist one record, retrying short-lived upstream races."""
    if not settings.INTERNAL_API_SECRET or not trace_id:
        logger.warning(
            "LLM cost backfill skipped trace=%s: INTERNAL_API_SECRET is not configured",
            trace_id,
        )
        return False
    total_attempts = max(1, attempts)
    for attempt in range(total_attempts):
        calls = fetch_llm_calls(trace_id)
        if calls:
            try:
                return _store_calls(record_id, trace_id, calls)
            except Exception:  # noqa: BLE001 - background repair must stay alive
                logger.exception(
                    "LLM cost store failed trace=%s record_id=%s",
                    trace_id,
                    record_id,
                )
                return False

        if attempt + 1 < total_attempts:
            delay = retry_delays[min(attempt, len(retry_delays) - 1)] if retry_delays else 0
            if delay > 0:
                time.sleep(delay)

    logger.warning(
        "LLM cost backfill has no calls trace=%s record_id=%s attempts=%d",
        trace_id,
        record_id,
        total_attempts,
    )
    return False


def repair_missing_comparison_llm_costs(
    *,
    limit: int = 100,
    lookback_days: int = 30,
) -> dict[str, int]:
    """Repair a bounded batch of recent terminal records with missing cost."""
    cutoff = datetime.now(CN_TZ).replace(tzinfo=None) - timedelta(
        days=max(1, lookback_days)
    )
    with SessionLocal() as db:
        # app_config has no price history. Repricing a record from before the
        # current config became effective would fabricate a historical cost, so
        # only repair records at/after that timestamp.
        price_config_updated_at = db.execute(
            select(AppConfig.updated_at).where(AppConfig.key == "llm_token_price")
        ).scalar_one_or_none()
        date_conditions = [ComparisonRecord.created_at >= cutoff]
        if price_config_updated_at is not None:
            date_conditions.append(
                ComparisonRecord.created_at
                >= _utc_to_beijing_naive(price_config_updated_at)
            )
        candidates = list(
            db.execute(
                select(ComparisonRecord.id, ComparisonRecord.trace_id)
                .where(
                    *date_conditions,
                    ComparisonRecord.status.in_(("success", "failed", "cancelled")),
                    ComparisonRecord.llm_cost_yuan.is_(None),
                )
                .order_by(ComparisonRecord.created_at.desc(), ComparisonRecord.id.desc())
                .limit(max(1, limit))
            ).all()
        )

    repaired = 0
    for record_id, trace_id in candidates:
        if backfill_comparison_llm_cost(
            record_id, trace_id, attempts=1, retry_delays=()
        ):
            repaired += 1
    return {
        "candidates": len(candidates),
        "repaired": repaired,
        "unresolved": len(candidates) - repaired,
    }