celery 容灾升级

video-gen-api/app/services/generation/billing_service.py

@@ -7,7 +7,7 @@ from typing import Any, Mapping
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.enums.credit_record import CreditRecordBillingScene, CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
+from app.enums.credit_record import CreditRecordChargeKind, CreditRecordOwnerType, CreditRecordSourceModule
 from app.models.credit_record import CreditRecord
 from app.models.generation_record import GenerationRecord
 from app.services.generation.media_reference_service import calculate_media_reference_usage
@@ -22,6 +22,7 @@ from app.services.credit_record_meta_service import (
     build_shot_video_analysis_meta,
 )
 from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits, deduct_credits
+from app.utils.id_gen import generate_id
 
 
 CHARGE_TEXT_PROMPT = CreditRecordChargeKind.TEXT_PROMPT.value
@@ -416,12 +417,44 @@ async def charge_shot_video_analysis_usage(
         charge_kind=CHARGE_VIDEO_ANALYSIS,
         action="charge",
     )
+    usage_snapshot = dict(usage)
+    token_usage_result = await db.execute(
+        select(TokenUsage)
+        .where(
+            TokenUsage.owner_type == owner_type,
+            TokenUsage.owner_id == owner_id,
+            TokenUsage.biz_key == biz_key,
+        )
+        .order_by(TokenUsage.created_at.asc())
+        .limit(1)
+    )
+    token_usage = token_usage_result.scalar_one_or_none()
+    if token_usage is None:
+        token_usage = TokenUsage(
+            id=generate_id(),
+            model_config_id=usage_snapshot.get("model_config_id"),
+            user_id=user_id,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=_safe_int(
+                usage_snapshot.get("total_tokens"),
+                input_tokens + output_tokens,
+            ),
+            owner_type=owner_type,
+            owner_id=owner_id,
+            biz_key=biz_key,
+            source_module=CreditRecordSourceModule.SHOT_REPLICATE.value,
+            source_step_code="video_analysis",
+        )
+        db.add(token_usage)
+        await db.flush()
+    usage_snapshot["token_usage_id"] = token_usage.id
     record_meta = await build_shot_video_analysis_meta(
         db,
         owner_type=owner_type,
         owner_id=owner_id,
         attempt_no=attempt_no,
-        usage=usage,
+        usage=usage_snapshot,
         billing_scene=billing_scene,
         source_project_id=source_project_id,
         source_step_id=source_step_id,