From b8eda53c0b2944850f7ef2edd506a506ce72c514 Mon Sep 17 00:00:00 2001 From: GinHa <15201596918@163.com> Date: Wed, 27 May 2026 13:11:55 +0800 Subject: [PATCH] =?UTF-8?q?=E7=94=9F=E6=88=90=E8=AE=B0=E5=BD=95=E6=8E=A5?= =?UTF-8?q?=E5=8F=A3=E5=85=BC=E5=AE=B9=E5=9B=BE=E7=89=87=E7=B1=BB=E5=9E=8B?= =?UTF-8?q?|CHAT=E5=AF=B9=E8=AF=9D=E7=94=9F=E6=88=90=E8=A7=86=E9=A2=91/?= =?UTF-8?q?=E5=9B=BE=E7=89=87=E4=BB=BB=E5=8A=A1API+celery=E5=BC=82?= =?UTF-8?q?=E6=AD=A5=E6=A8=A1=E5=9D=97=E5=AE=8C=E6=88=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- video-gen-api/.env | 2 +- video-gen-api/.gitignore | 162 ++++++ video-gen-api/alembic/env.py | 3 + ...93c7f_add_chatapi_async_generation_mode.py | 149 ++++++ video-gen-api/app/api/v1/__init__.py | 2 + video-gen-api/app/api/v1/admin.py | 104 ++-- video-gen-api/app/api/v1/generation.py | 6 +- video-gen-api/app/api/v1/generation_ai.py | 395 ++++++++++++++ video-gen-api/app/config.py | 26 + video-gen-api/app/models/__init__.py | 4 + video-gen-api/app/models/base.py | 34 +- .../app/models/chat_generation_task.py | 64 +++ .../app/models/chat_generation_task_event.py | 25 + .../app/models/chat_provider_call_log.py | 35 ++ video-gen-api/app/schemas/generation_ai.py | 418 +++++++++++++++ .../app/services/generation_ai_service.py | 495 ++++++++++++++++++ .../services/generation_billing_service.py | 321 ++++++++++++ .../services/generation_download_service.py | 31 ++ .../app/services/generation_log_service.py | 134 +++++ .../app/services/generation_prompt_service.py | 184 +++++++ .../services/generation_provider_service.py | 161 ++++++ video-gen-api/app/services/image_gen.py | 4 +- video-gen-api/app/services/provider_limit.py | 59 +++ video-gen-api/app/services/video_gen.py | 4 +- video-gen-api/app/tasks/__init__.py | 11 + video-gen-api/app/tasks/async_runner.py | 48 ++ video-gen-api/app/tasks/celery_app.py | 68 ++- .../app/tasks/generation_create_tasks.py | 113 ++++ .../app/tasks/generation_download_tasks.py | 176 +++++++ .../app/tasks/generation_poll_tasks.py | 199 +++++++ 30 files changed, 3392 insertions(+), 45 deletions(-) create mode 100644 video-gen-api/.gitignore create mode 100644 video-gen-api/alembic/versions/c310d7193c7f_add_chatapi_async_generation_mode.py create mode 100644 video-gen-api/app/api/v1/generation_ai.py create mode 100644 video-gen-api/app/models/chat_generation_task.py create mode 100644 video-gen-api/app/models/chat_generation_task_event.py create mode 100644 video-gen-api/app/models/chat_provider_call_log.py create mode 100644 video-gen-api/app/schemas/generation_ai.py create mode 100644 video-gen-api/app/services/generation_ai_service.py create mode 100644 video-gen-api/app/services/generation_billing_service.py create mode 100644 video-gen-api/app/services/generation_download_service.py create mode 100644 video-gen-api/app/services/generation_log_service.py create mode 100644 video-gen-api/app/services/generation_prompt_service.py create mode 100644 video-gen-api/app/services/generation_provider_service.py create mode 100644 video-gen-api/app/services/provider_limit.py create mode 100644 video-gen-api/app/tasks/async_runner.py create mode 100644 video-gen-api/app/tasks/generation_create_tasks.py create mode 100644 video-gen-api/app/tasks/generation_download_tasks.py create mode 100644 video-gen-api/app/tasks/generation_poll_tasks.py diff --git a/video-gen-api/.env b/video-gen-api/.env index a6db5e50..4c3bba2f 100644 --- a/video-gen-api/.env +++ b/video-gen-api/.env @@ -8,7 +8,7 @@ SECRET_KEY=local-dev-secret-key-not-for-production DATABASE_URL=postgresql+asyncpg://videogen:7k33pnXdPL62Yyb4@180.184.42.66:5432/videogen # Redis (leave empty to disable - rate limiting and captcha will use in-memory fallback) -REDIS_URL= +REDIS_URL=redis://localhost:6379/0 # JWT JWT_ALGORITHM=HS256 diff --git a/video-gen-api/.gitignore b/video-gen-api/.gitignore new file mode 100644 index 00000000..35e593ed --- /dev/null +++ b/video-gen-api/.gitignore @@ -0,0 +1,162 @@ +# ========================= +# Python cache +# ========================= +__pycache__/ +*.py[cod] +*$py.class +*.pyo +*.pyd + +# ========================= +# Python virtual environments +# ========================= +.venv/ +venv/ +env/ +ENV/ +.python-version + +# ========================= +# Python build / package +# ========================= +build/ +dist/ +*.egg-info/ +.eggs/ +pip-wheel-metadata/ +wheelhouse/ + +# ========================= +# Test / coverage / cache +# ========================= +.pytest_cache/ +.coverage +.coverage.* +htmlcov/ +.tox/ +.nox/ +.cache/ +coverage.xml +*.cover +*.py,cover + +# ========================= +# Type checker / linter cache +# ========================= +.mypy_cache/ +.pyre/ +.pytype/ +.ruff_cache/ +.dmypy.json +dmypy.json + +# ========================= +# FastAPI / Uvicorn / runtime files +# ========================= +*.log +logs/ +log/ +tmp/ +temp/ +.pid +*.pid + +# ========================= +# Environment config +# ========================= +.env +.env.* +!.env.example + +# ========================= +# Database local files +# ========================= +*.db +*.sqlite +*.sqlite3 +videogen.db + +# ========================= +# Alembic +# 注意：不要忽略 alembic/versions/ +# 迁移文件需要提交到 Git +# ========================= + +# ========================= +# Local storage / uploaded files / generated videos +# ========================= +storage/ +uploads/ +media/ +videos/ +outputs/ + +# 如果你想保留空目录，可以手动加 .gitkeep +!storage/.gitkeep +!uploads/.gitkeep +!media/.gitkeep +!videos/.gitkeep +!outputs/.gitkeep + +# ========================= +# Encryption / keys / certs +# ========================= +*.pem +*.key +*.crt +*.csr +*.p12 +*.pfx +private_key* +public_key* +rsa_private* +rsa_public* + +# ========================= +# Docker local override +# ========================= +docker-compose.override.yml +docker-compose.local.yml + +# ========================= +# Celery / Redis / worker runtime +# ========================= +celerybeat-schedule +celerybeat.pid +*.dump + +# ========================= +# IDE / Editor +# ========================= +.vscode/ +.idea/ +*.swp +*.swo +*.swn +*.bak +*.tmp + +# ========================= +# OS files +# ========================= +.DS_Store +Thumbs.db +desktop.ini + +# ========================= +# Jupyter +# ========================= +.ipynb_checkpoints/ + +# ========================= +# Node / frontend optional +# 如果后面有前端目录可以保留 +# ========================= +node_modules/ +npm-debug.log* +yarn-debug.log* +yarn-error.log* +pnpm-debug.log* +.vite/ +.next/ +.nuxt/ \ No newline at end of file diff --git a/video-gen-api/alembic/env.py b/video-gen-api/alembic/env.py index 111000a8..957311d6 100644 --- a/video-gen-api/alembic/env.py +++ b/video-gen-api/alembic/env.py @@ -29,6 +29,9 @@ from app.models.payment_order import PaymentOrder # noqa: F401 from app.models.recharge_package import RechargePackage # noqa: F401 from app.models.video_engine import VideoEngine # noqa: F401 from app.models.image_engine import ImageEngine # noqa: F401 +from app.models.chat_generation_task import ChatGenerationTask # noqa: F401 +from app.models.chat_generation_task_event import ChatGenerationTaskEvent # noqa: F401 +from app.models.chat_provider_call_log import ChatProviderCallLog # noqa: F401 config = context.config diff --git a/video-gen-api/alembic/versions/c310d7193c7f_add_chatapi_async_generation_mode.py b/video-gen-api/alembic/versions/c310d7193c7f_add_chatapi_async_generation_mode.py new file mode 100644 index 00000000..05303322 --- /dev/null +++ b/video-gen-api/alembic/versions/c310d7193c7f_add_chatapi_async_generation_mode.py @@ -0,0 +1,149 @@ +"""add chatapi async generation mode + +Revision ID: c310d7193c7f +Revises: 66b40a9e2d63 +Create Date: 2026-05-26 14:52:20.327046 +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'c310d7193c7f' +down_revision: Union[str, None] = '66b40a9e2d63' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +def upgrade() -> None: + op.create_table( + "chat_generation_tasks", + sa.Column("id", sa.String(length=32), primary_key=True), + sa.Column("user_id", sa.String(length=32), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False), + sa.Column("original_prompt", sa.Text(), nullable=False), + sa.Column("optimized_prompt", sa.Text(), nullable=True), + sa.Column("gen_type", sa.String(length=16), server_default="video", nullable=False), + sa.Column("duration", sa.Integer(), nullable=True), + sa.Column("aspect_ratio", sa.String(length=8), nullable=True), + sa.Column("resolution", sa.String(length=8), nullable=True), + sa.Column("image_size", sa.String(length=16), nullable=True), + sa.Column("image_proportion", sa.String(length=8), nullable=True), + sa.Column("image_px", sa.String(length=16), nullable=True), + sa.Column("status", sa.String(length=32), server_default="generating", nullable=False), + sa.Column("pipeline_stage", sa.String(length=32), nullable=True), + sa.Column("generation_mode", sa.String(length=32), server_default="chatapi_async", nullable=False), + sa.Column("media_references", sa.Text(), nullable=True), + sa.Column("provider_task_id", sa.String(length=128), nullable=True), + sa.Column("seedance_task_id", sa.String(length=128), nullable=True), + sa.Column("remote_result_url", sa.Text(), nullable=True), + sa.Column("image_url", sa.String(length=512), nullable=True), + sa.Column("video_url", sa.String(length=512), nullable=True), + sa.Column("engine_id", sa.String(length=32), nullable=True), + sa.Column("engine_snapshot_json", sa.Text(), nullable=True), + sa.Column("provider_response_json", sa.Text(), nullable=True), + sa.Column("credits_cost", sa.Float(), server_default="0", nullable=False), + sa.Column("text_credits_cost", sa.Float(), server_default="0", nullable=False), + sa.Column("text_tokens_used", sa.Integer(), server_default="0", nullable=False), + sa.Column("video_tokens_used", sa.Integer(), server_default="0", nullable=False), + sa.Column("image_tokens_used", sa.Integer(), server_default="0", nullable=False), + sa.Column("retry_count", sa.Integer(), server_default="0", nullable=False), + sa.Column("poll_count", sa.Integer(), server_default="0", nullable=False), + sa.Column("last_poll_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("deadline_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("generated_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("idempotency_key", sa.String(length=64), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=True), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=True), + ) + op.create_index("ix_chat_generation_tasks_user_id", "chat_generation_tasks", ["user_id"]) + op.create_index("ix_chat_generation_tasks_gen_type", "chat_generation_tasks", ["gen_type"]) + op.create_index("ix_chat_generation_tasks_status", "chat_generation_tasks", ["status"]) + op.create_index("ix_chat_generation_tasks_pipeline_stage", "chat_generation_tasks", ["pipeline_stage"]) + op.create_index("ix_chat_generation_tasks_generation_mode", "chat_generation_tasks", ["generation_mode"]) + op.create_index("ix_chat_generation_tasks_provider_task_id", "chat_generation_tasks", ["provider_task_id"]) + op.create_index("ix_chat_generation_tasks_seedance_task_id", "chat_generation_tasks", ["seedance_task_id"]) + op.create_index("ix_chat_generation_tasks_engine_id", "chat_generation_tasks", ["engine_id"]) + op.create_index("ix_chat_generation_tasks_idempotency_key", "chat_generation_tasks", ["idempotency_key"]) + + op.create_table( + "chat_generation_task_events", + sa.Column("id", sa.String(length=32), primary_key=True), + sa.Column("task_id", sa.String(length=32), sa.ForeignKey("chat_generation_tasks.id", ondelete="CASCADE"), nullable=False), + sa.Column("generation_mode", sa.String(length=32), nullable=True), + sa.Column("event_type", sa.String(length=64), nullable=False), + sa.Column("from_status", sa.String(length=32), nullable=True), + sa.Column("to_status", sa.String(length=32), nullable=True), + sa.Column("from_stage", sa.String(length=32), nullable=True), + sa.Column("to_stage", sa.String(length=32), nullable=True), + sa.Column("message", sa.String(length=512), nullable=True), + sa.Column("detail_json", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=True), + ) + op.create_index("ix_chat_generation_task_events_task_id", "chat_generation_task_events", ["task_id"]) + op.create_index("ix_chat_generation_task_events_generation_mode", "chat_generation_task_events", ["generation_mode"]) + op.create_index("ix_chat_generation_task_events_event_type", "chat_generation_task_events", ["event_type"]) + op.create_index("ix_chat_generation_task_events_created_at", "chat_generation_task_events", ["created_at"]) + + op.create_table( + "chat_provider_call_logs", + sa.Column("id", sa.String(length=32), primary_key=True), + sa.Column("task_id", sa.String(length=32), sa.ForeignKey("chat_generation_tasks.id", ondelete="CASCADE"), nullable=False), + sa.Column("generation_mode", sa.String(length=32), nullable=True), + sa.Column("provider", sa.String(length=64), nullable=True), + sa.Column("api_type", sa.String(length=64), nullable=False), + sa.Column("model", sa.String(length=128), nullable=True), + sa.Column("engine_id", sa.String(length=32), nullable=True), + sa.Column("status", sa.String(length=32), nullable=False), + sa.Column("latency_ms", sa.Integer(), nullable=True), + sa.Column("http_status", sa.Integer(), nullable=True), + sa.Column("provider_task_id", sa.String(length=128), nullable=True), + sa.Column("request_hash", sa.String(length=64), nullable=True), + sa.Column("response_hash", sa.String(length=64), nullable=True), + sa.Column("request_excerpt", sa.Text(), nullable=True), + sa.Column("response_excerpt", sa.Text(), nullable=True), + sa.Column("prompt_tokens", sa.Integer(), server_default="0", nullable=False), + sa.Column("completion_tokens", sa.Integer(), server_default="0", nullable=False), + sa.Column("total_tokens", sa.Integer(), server_default="0", nullable=False), + sa.Column("error_code", sa.String(length=128), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=True), + ) + op.create_index("ix_chat_provider_call_logs_task_id", "chat_provider_call_logs", ["task_id"]) + op.create_index("ix_chat_provider_call_logs_generation_mode", "chat_provider_call_logs", ["generation_mode"]) + op.create_index("ix_chat_provider_call_logs_provider", "chat_provider_call_logs", ["provider"]) + op.create_index("ix_chat_provider_call_logs_api_type", "chat_provider_call_logs", ["api_type"]) + op.create_index("ix_chat_provider_call_logs_engine_id", "chat_provider_call_logs", ["engine_id"]) + op.create_index("ix_chat_provider_call_logs_status", "chat_provider_call_logs", ["status"]) + op.create_index("ix_chat_provider_call_logs_provider_task_id", "chat_provider_call_logs", ["provider_task_id"]) + op.create_index("ix_chat_provider_call_logs_created_at", "chat_provider_call_logs", ["created_at"]) + + +def downgrade() -> None: + op.drop_index("ix_chat_provider_call_logs_created_at", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_provider_task_id", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_status", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_engine_id", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_api_type", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_provider", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_generation_mode", table_name="chat_provider_call_logs") + op.drop_index("ix_chat_provider_call_logs_task_id", table_name="chat_provider_call_logs") + op.drop_table("chat_provider_call_logs") + + op.drop_index("ix_chat_generation_task_events_created_at", table_name="chat_generation_task_events") + op.drop_index("ix_chat_generation_task_events_event_type", table_name="chat_generation_task_events") + op.drop_index("ix_chat_generation_task_events_generation_mode", table_name="chat_generation_task_events") + op.drop_index("ix_chat_generation_task_events_task_id", table_name="chat_generation_task_events") + op.drop_table("chat_generation_task_events") + + op.drop_index("ix_chat_generation_tasks_idempotency_key", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_engine_id", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_seedance_task_id", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_provider_task_id", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_generation_mode", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_pipeline_stage", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_status", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_gen_type", table_name="chat_generation_tasks") + op.drop_index("ix_chat_generation_tasks_user_id", table_name="chat_generation_tasks") + op.drop_table("chat_generation_tasks") \ No newline at end of file diff --git a/video-gen-api/app/api/v1/__init__.py b/video-gen-api/app/api/v1/__init__.py index 6a8d82f8..68447e54 100644 --- a/video-gen-api/app/api/v1/__init__.py +++ b/video-gen-api/app/api/v1/__init__.py @@ -14,6 +14,7 @@ from app.api.v1.menu_configs import router as menu_configs_router from app.api.v1.recharge_packages import router as recharge_packages_router from app.api.v1.video_engines import router as video_engines_router from app.api.v1.image_engines import router as image_engines_router +from app.api.v1.generation_ai import router as generation_ai_router api_router = APIRouter() api_router.include_router(auth_router) @@ -30,3 +31,4 @@ api_router.include_router(menu_configs_router) api_router.include_router(recharge_packages_router) api_router.include_router(video_engines_router) api_router.include_router(image_engines_router) +api_router.include_router(generation_ai_router) diff --git a/video-gen-api/app/api/v1/admin.py b/video-gen-api/app/api/v1/admin.py index d4360386..ed751537 100644 --- a/video-gen-api/app/api/v1/admin.py +++ b/video-gen-api/app/api/v1/admin.py @@ -43,6 +43,7 @@ from app.services.notification import create_notification from app.services.auth import hash_password, verify_password from app.services.operation_log import log_operation from app.utils.id_gen import generate_id +from app.schemas.generation import GenerationType, ASPECT_RATIOS, RESOLUTIONS CST = timezone(timedelta(hours=8)) @@ -991,6 +992,14 @@ async def admin_list_generation_records( "error_message": record.error_message, "created_at": _iso(record.created_at), "generated_at": _iso(record.generated_at), + + # append img param + "gen_type": record.gen_type, + "image_size": record.image_size or '', + "image_url": record.image_url or '', + "image_tokens_used": record.image_tokens_used or 0, + "image_proportion": record.image_proportion or '', + "image_px": record.image_px or '', }) return {"total": total, "items": items} @@ -1033,8 +1042,8 @@ async def admin_generate_video( ): """Admin trigger video generation for a record with specified params.""" from app.models.project import Project - from app.services.credits import calc_video_credits, deduct_credits - from app.schemas.generation import ASPECT_RATIOS, RESOLUTIONS + from app.services.credits import calc_video_credits, deduct_credits, calc_image_credits + from app.services.video_queue import task_queue result = await db.execute( select(GenerationRecord, Project.name) @@ -1046,43 +1055,68 @@ async def admin_generate_video( raise HTTPException(status_code=404, detail="记录不存在") record, project_name = row + type_str = "视频" if record.gen_type == GenerationType.video else "图片" + if record.status not in ("prompt_optimized", "failed"): - raise HTTPException(status_code=400, detail="当前状态不允许生成视频") + raise HTTPException(status_code=400, detail=f"当前状态不允许生成{type_str}") - aspect_ratio = body.get("aspect_ratio", "16:9") - resolution = body.get("resolution", "720p") - if aspect_ratio not in ASPECT_RATIOS: - raise HTTPException(status_code=400, detail="不支持的画面比例") - if resolution not in RESOLUTIONS: - raise HTTPException(status_code=400, detail="不支持的分辨率") + if record.gen_type == GenerationType.video: + # Video Generation + aspect_ratio = body.get("aspect_ratio", "16:9") + resolution = body.get("resolution", "720p") + if aspect_ratio not in ASPECT_RATIOS: + raise HTTPException(status_code=400, detail="不支持的画面比例") + if resolution not in RESOLUTIONS: + raise HTTPException(status_code=400, detail="不支持的分辨率") - duration = record.duration or 5 - video_credits = await calc_video_credits(db, duration, resolution) - await deduct_credits( - db, record.user_id, video_credits, - f"视频生成(管理后台) - {project_name}", - related_id=record_id, - ) + duration = record.duration or 5 + video_credits = await calc_video_credits(db, duration, resolution) + await deduct_credits( + db, record.user_id, video_credits, + f"视频生成(管理后台) - {project_name}", + related_id=record_id, + ) - record.aspect_ratio = aspect_ratio - record.resolution = resolution - record.credits_cost = (record.credits_cost or 0) + video_credits - record.status = "generating" - record.error_message = None - await db.flush() - - try: - from app.services.video_gen import get_active_engine, submit_video_task - from app.services.video_queue import task_queue - - engine = await get_active_engine(db) - task_id = await submit_video_task(db, engine, record) - record.seedance_task_id = task_id - await db.flush() - await task_queue.enqueue(record_id) - except Exception as e: - record.status = "failed" - record.error_message = str(e) + record.aspect_ratio = aspect_ratio + record.resolution = resolution + record.credits_cost = (record.credits_cost or 0) + video_credits + record.status = "generating" + record.error_message = None await db.flush() + try: + from app.services.video_gen import get_active_engine, submit_video_task + + engine = await get_active_engine(db) + task_id = await submit_video_task(db, engine, record) + record.seedance_task_id = task_id + await db.flush() + await task_queue.enqueue(record_id) + except Exception as e: + record.status = "failed" + record.error_message = str(e) + await db.flush() + elif record.gen_type == GenerationType.image: + # Image generation + + post_image_size = body.get("image_size", "") + image_credits = await calc_image_credits(db, post_image_size or record.image_size or "2K") + await deduct_credits( + db, record.user_id, image_credits, + f"图片生成 - {project_name}", + related_id=record_id, + ) + + try: + record.image_size = post_image_size or record.image_size or "2K" + record.credits_cost = round(image_credits, 2) + record.status = "generating" + record.error_message = None + await db.flush() + await task_queue.enqueue(record_id) + except Exception as e: + record.status = "failed" + record.error_message = str(e) + await db.flush() + return {"message": "ok", "record_id": record_id} diff --git a/video-gen-api/app/api/v1/generation.py b/video-gen-api/app/api/v1/generation.py index 35326b48..94725ac0 100644 --- a/video-gen-api/app/api/v1/generation.py +++ b/video-gen-api/app/api/v1/generation.py @@ -201,7 +201,11 @@ async def optimize( # record.error_message = extract_error_message(e, "提示词") # await db.flush() # await db.commit() - raise HTTPException(status_code=502, detail=f"AI模型调用失败: {extract_error_message(e, "提示词")}") + error_message = extract_error_message(e, "提示词") + raise HTTPException( + status_code=502, + detail=f"AI模型调用失败: {error_message}" + ) text_credits = await calc_text_credits( db, token_usage["input_tokens"], token_usage["output_tokens"], diff --git a/video-gen-api/app/api/v1/generation_ai.py b/video-gen-api/app/api/v1/generation_ai.py new file mode 100644 index 00000000..b30ea560 --- /dev/null +++ b/video-gen-api/app/api/v1/generation_ai.py @@ -0,0 +1,395 @@ +from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.dependencies import get_current_user, get_db +from app.models.chat_generation_task import ChatGenerationTask +from app.models.user import User +from app.schemas.generation_ai import ( + GenerationAIHistoryDayItemsOut, + GenerationAIHistoryGroupedOut, + GenerationAIRetryOut, + GenerationAITaskCreate, + GenerationAITaskListOut, + GenerationAITaskOut, +) +from app.services.generation_ai_service import ( + create_async_generation_task, + list_async_generation_tasks, + list_generation_history_day_items, + list_generation_history_grouped_days, + record_to_out, +) +from app.services.generation_log_service import log_task_event +from app.tasks.celery_app import celery_app + +router = APIRouter( + prefix="/generation-ai", + tags=["generation-ai"], +) + + +@router.post( + "/tasks", + response_model=GenerationAITaskOut, + summary="创建AI图片/视频生成任务", + description=( + "创建一个项目无关的AI生成任务。" + "该接口用于 Chat 风格的图片/视频生成，不再绑定 project_id。" + "创建成功后会写入 chat_generation_tasks 表，并投递 Celery 异步任务。" + "支持 image 图片生成和 video 视频生成。" + "建议前端传入 idempotency_key，用于防止按钮连点、网络重试导致重复创建任务和重复扣费。" + ), + responses={ + 200: { + "description": "任务创建成功，返回任务详情", + }, + 400: { + "description": "请求参数错误，例如 gen_type 不支持、图片尺寸不支持、视频参数不支持等", + }, + 401: { + "description": "未登录或 Token 无效", + }, + 503: { + "description": "Celery 未启用或消息队列未配置", + }, + }, +) +async def create_task( + req: GenerationAITaskCreate = Body( + ..., + description="AI生成任务创建参数。gen_type=image 时使用图片参数；gen_type=video 时使用视频参数", + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + if celery_app is None: + raise HTTPException(status_code=503, detail="Celery未启用：请配置 REDIS_URL 或 CELERY_BROKER_URL 后启动 worker") + + task = await create_async_generation_task(db, current_user, req) + await db.commit() + + await log_task_event( + task, + event_type="TASK_CREATED", + to_status="generating", + to_stage="queued", + detail={"gen_type": task.gen_type}, + ) + + from app.tasks.generation_create_tasks import chatapi_create_generation_task + + chatapi_create_generation_task.delay(task.id) + return record_to_out(task) + + +@router.get( + "/tasks", + response_model=GenerationAITaskListOut, + summary="获取AI生成任务列表", + description=( + "分页获取当前登录用户的AI生成任务列表。" + "可按生成类型 gen_type 和任务状态 status 过滤。" + "该接口返回的是普通任务列表，不按日期分组。" + "如果前端需要按生成日期分组展示历史记录，请使用 /generation-ai/history 接口。" + ), + responses={ + 200: { + "description": "查询成功，返回任务总数和当前分页任务列表", + }, + 401: { + "description": "未登录或 Token 无效", + }, + }, +) +async def list_tasks( + gen_type: str | None = Query( + None, + description="生成类型筛选：image=图片任务，video=视频任务；为空表示不过滤生成类型", + examples=["image"], + ), + status: str | None = Query( + None, + description=( + "任务状态筛选。常见值：generating=生成中，completed=已完成，failed=失败；" + "为空表示不过滤状态" + ), + examples=["completed"], + ), + page: int = Query( + 1, + ge=1, + description="分页页码，从1开始", + examples=[1], + ), + page_size: int = Query( + 20, + ge=1, + le=100, + description="每页返回数量，范围 1~100", + examples=[20], + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + total, items = await list_async_generation_tasks( + db, + current_user.id, + gen_type, + status, + page, + page_size, + ) + return GenerationAITaskListOut(total=total, items=[record_to_out(i) for i in items]) + + +@router.get( + "/history", + response_model=GenerationAIHistoryGroupedOut, + summary="获取AI生成历史日期分组", + description=( + "按生成完成日期倒序返回当前用户的AI生成历史记录。" + "该接口只返回生成成功的任务，即 status=completed 且 generated_at 不为空的数据。" + "必须通过 gen_type 区分图片和视频。" + "分页对象是生成日期，不是单条记录。" + "每页最多返回10个生成日期分组，每个日期分组内最多返回该日期下倒序前10条生成记录。" + "如果某一天 total 大于10，前端可调用 /generation-ai/history/{generated_date} 加载该日期下的后续分页数据。" + ), + responses={ + 200: { + "description": "查询成功，返回按生成日期分组的历史记录", + }, + 400: { + "description": "参数错误，例如 gen_type 不是 image 或 video", + }, + 401: { + "description": "未登录或 Token 无效", + }, + }, +) +async def list_history_grouped_days( + gen_type: str = Query( + ..., + description="生成类型：image=图片历史，video=视频历史", + examples=["image"], + ), + page: int = Query( + 1, + ge=1, + description="日期分组分页页码，从1开始。注意：这里分页的是生成日期，不是单条生成记录", + examples=[1], + ), + page_size: int = Query( + 10, + ge=1, + le=10, + description="每页返回的生成日期数量，范围 1~10，最大只能获取10天", + examples=[10], + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + return await list_generation_history_grouped_days( + db=db, + user_id=current_user.id, + gen_type=gen_type, + page=page, + page_size=page_size, + ) + + +@router.get( + "/history/{generated_date}", + response_model=GenerationAIHistoryDayItemsOut, + summary="获取指定日期下的AI生成历史分页", + description=( + "获取某一个生成日期下的生成成功记录分页。" + "该接口用于前端在历史分组列表中继续加载某一天的后续记录。" + "例如 /history 接口中某一天 total=18，但 items 只返回前10条，" + "则前端可以调用本接口 page=2&page_size=10 获取该日期下剩余记录。" + "该接口同样必须通过 gen_type 区分图片和视频。" + ), + responses={ + 200: { + "description": "查询成功，返回指定日期下的历史记录分页", + }, + 400: { + "description": "参数错误，例如 generated_date 格式不是 YYYY-MM-DD，或 gen_type 不合法", + }, + 401: { + "description": "未登录或 Token 无效", + }, + }, +) +async def list_history_day_items( + generated_date: str = Path( + ..., + description="生成日期，格式：YYYY-MM-DD，例如：2026-05-27", + examples=["2026-05-27"], + ), + gen_type: str = Query( + ..., + description="生成类型：image=图片历史，video=视频历史", + examples=["image"], + ), + page: int = Query( + 1, + ge=1, + description="当前日期下的记录分页页码，从1开始", + examples=[1], + ), + page_size: int = Query( + 10, + ge=1, + le=100, + description="当前日期下每页返回的生成记录数量，范围 1~100", + examples=[10], + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + return await list_generation_history_day_items( + db=db, + user_id=current_user.id, + gen_type=gen_type, + generated_date=generated_date, + page=page, + page_size=page_size, + ) + + +@router.get( + "/tasks/{task_id}", + response_model=GenerationAITaskOut, + summary="获取AI生成任务详情", + description=( + "根据任务ID获取当前登录用户的AI生成任务详情。" + "只能查询当前用户自己的任务，且只查询 generation_mode=chatapi_async 的任务。" + "如果任务不存在或不属于当前用户，返回404。" + ), + responses={ + 200: { + "description": "查询成功，返回任务详情", + }, + 401: { + "description": "未登录或 Token 无效", + }, + 404: { + "description": "任务不存在，或任务不属于当前用户", + }, + }, +) +async def get_task( + task_id: str = Path( + ..., + description="AI生成任务ID", + examples=["0019e0a44895b6d837d"], + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + result = await db.execute( + select(ChatGenerationTask).where( + ChatGenerationTask.id == task_id, + ChatGenerationTask.user_id == current_user.id, + ChatGenerationTask.generation_mode == "chatapi_async", + ) + ) + task = result.scalar_one_or_none() + if not task: + raise HTTPException(status_code=404, detail="任务不存在") + return record_to_out(task) + + +@router.post( + "/tasks/{task_id}/retry", + response_model=GenerationAIRetryOut, + summary="重试失败的AI生成任务", + description=( + "重新投递一个失败的AI生成任务。" + "只有 status=failed 的任务允许重试。" + "重试次数超过3次后不允许继续重试。" + "后端会根据任务当前数据自动判断应从哪个流水线阶段恢复，" + "例如重新优化提示词、重新创建第三方任务、继续轮询远程结果或重新下载结果。" + ), + responses={ + 200: { + "description": "任务重新投递成功", + }, + 400: { + "description": "任务状态不允许重试，或已超过最大重试次数", + }, + 401: { + "description": "未登录或 Token 无效", + }, + 404: { + "description": "任务不存在，或任务不属于当前用户", + }, + 503: { + "description": "Celery 未启用或消息队列未配置", + }, + }, +) +async def retry_task( + task_id: str = Path( + ..., + description="需要重试的AI生成任务ID", + examples=["0019e0a44895b6d837d"], + ), + current_user: User = Depends(get_current_user), + db: AsyncSession = Depends(get_db), +): + if celery_app is None: + raise HTTPException(status_code=503, detail="Celery未启用：请配置 REDIS_URL 或 CELERY_BROKER_URL 后启动 worker") + + result = await db.execute( + select(ChatGenerationTask).where( + ChatGenerationTask.id == task_id, + ChatGenerationTask.user_id == current_user.id, + ChatGenerationTask.generation_mode == "chatapi_async", + ) + ) + task = result.scalar_one_or_none() + if not task: + raise HTTPException(status_code=404, detail="任务不存在") + if task.status != "failed": + raise HTTPException(status_code=400, detail="只有失败任务可以重试") + + task.status = "generating" + task.error_message = None + task.retry_count = (task.retry_count or 0) + 1 + if task.retry_count > 3: + raise HTTPException(status_code=400, detail="任务已超过最大重试次数") + + # Resume from the earliest missing stage. + if not task.optimized_prompt: + task.pipeline_stage = "queued" + from app.tasks.generation_create_tasks import chatapi_create_generation_task + + chatapi_create_generation_task.delay(task.id) + elif not task.seedance_task_id and not task.remote_result_url: + task.pipeline_stage = "creating_provider_task" + from app.tasks.generation_create_tasks import chatapi_create_generation_task + + chatapi_create_generation_task.delay(task.id) + elif task.seedance_task_id and not task.remote_result_url: + task.pipeline_stage = "waiting_remote" + from app.tasks.generation_poll_tasks import poll_generation_task + + poll_generation_task.delay(task.id) + elif task.remote_result_url and not (task.image_url or task.video_url): + task.pipeline_stage = "result_ready" + from app.tasks.generation_download_tasks import download_generation_result_task + + download_generation_result_task.delay(task.id) + else: + task.status = "completed" + task.pipeline_stage = "done" + + await db.commit() + return GenerationAIRetryOut( + id=task.id, + status=task.status, + pipeline_stage=task.pipeline_stage, + message="任务已重新投递", + ) \ No newline at end of file diff --git a/video-gen-api/app/config.py b/video-gen-api/app/config.py index 07e2d41d..8be32f24 100644 --- a/video-gen-api/app/config.py +++ b/video-gen-api/app/config.py @@ -58,5 +58,31 @@ class Settings(BaseSettings): RATE_LIMIT_ENABLED: bool = True + # ChatAPI async generation pipeline settings + CELERY_BROKER_URL: str = "" + CELERY_RESULT_BACKEND: str = "" + CHATAPI_REQUEST_TIMEOUT_SECONDS: int = 120 + CHATAPI_VIDEO_FPS: float = 0.5 + CHATAPI_ASYNC_MAX_RETRIES: int = 3 + CHATAPI_ASYNC_RETRY_BACKOFF_SECONDS: int = 30 + CHATAPI_ASYNC_POLL_INTERVAL_SECONDS: int = 30 + CHATAPI_ASYNC_IMAGE_DEADLINE_MINUTES: int = 10 + CHATAPI_ASYNC_VIDEO_DEADLINE_MINUTES: int = 30 + + # Distributed provider concurrency limits. 0 means disabled/no-op. + ARK_CHAT_PROMPT_MAX_CONCURRENCY: int = 20 + ARK_IMAGE_CREATE_MAX_CONCURRENCY: int = 10 + ARK_VIDEO_CREATE_MAX_CONCURRENCY: int = 10 + ARK_IMAGE_POLL_MAX_CONCURRENCY: int = 50 + ARK_VIDEO_POLL_MAX_CONCURRENCY: int = 50 + RESULT_DOWNLOAD_MAX_CONCURRENCY: int = 10 + PROVIDER_LIMIT_WAIT_TIMEOUT_SECONDS: float = 30.0 + PROVIDER_LIMIT_TOKEN_TTL_SECONDS: int = 600 + + CELERY_DB_POOL_SIZE: int = 1 + CELERY_DB_MAX_OVERFLOW: int = 1 + CELERY_DB_POOL_TIMEOUT: int = 30 + CELERY_DB_POOL_RECYCLE: int = 1800 + settings = Settings() diff --git a/video-gen-api/app/models/__init__.py b/video-gen-api/app/models/__init__.py index e35c472b..6534ea90 100644 --- a/video-gen-api/app/models/__init__.py +++ b/video-gen-api/app/models/__init__.py @@ -15,6 +15,9 @@ from app.models.credit_ratio import CreditRatio from app.models.menu_config import MenuConfig from app.models.recharge_package import RechargePackage from app.models.operation_log import OperationLog +from app.models.chat_generation_task import ChatGenerationTask +from app.models.chat_generation_task_event import ChatGenerationTaskEvent +from app.models.chat_provider_call_log import ChatProviderCallLog __all__ = [ "Base", "TimestampMixin", "engine", "async_session", @@ -23,4 +26,5 @@ __all__ = [ "ModelConfig", "SystemConfig", "Notification", "PaymentOrder", "TokenUsage", "IndustryConfig", "VideoEngine", "CreditRatio", "MenuConfig", "RechargePackage", "OperationLog", + "ChatGenerationTask", "ChatGenerationTaskEvent", "ChatProviderCallLog", ] diff --git a/video-gen-api/app/models/base.py b/video-gen-api/app/models/base.py index 02798394..7c5f9c49 100644 --- a/video-gen-api/app/models/base.py +++ b/video-gen-api/app/models/base.py @@ -1,4 +1,5 @@ from datetime import datetime +import os from sqlalchemy import DateTime, func from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine @@ -6,8 +7,35 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from app.config import settings -engine = create_async_engine(settings.DATABASE_URL, echo=settings.DEBUG) -async_session = async_sessionmaker(engine, expire_on_commit=False) + +def _is_celery_process() -> bool: + argv = " ".join(os.sys.argv).lower() + return "celery" in argv + + +engine_kwargs = { + "echo": settings.DEBUG, + "pool_pre_ping": True, +} + +# SQLite 本地调试时不要乱塞 pool_size/max_overflow，PostgreSQL/asyncpg 才建议配置 +if _is_celery_process() and not settings.DATABASE_URL.startswith("sqlite"): + engine_kwargs.update( + pool_size=settings.CELERY_DB_POOL_SIZE, + max_overflow=settings.CELERY_DB_MAX_OVERFLOW, + pool_timeout=settings.CELERY_DB_POOL_TIMEOUT, + pool_recycle=settings.CELERY_DB_POOL_RECYCLE, + ) + +engine = create_async_engine(settings.DATABASE_URL, **engine_kwargs) + +async_session = async_sessionmaker( + engine, + expire_on_commit=False, + autoflush=False, +) + +AsyncSessionLocal = async_session class Base(AsyncAttrs, DeclarativeBase): @@ -29,4 +57,4 @@ async def init_database() -> None: async def close_database() -> None: - await engine.dispose() + await engine.dispose() \ No newline at end of file diff --git a/video-gen-api/app/models/chat_generation_task.py b/video-gen-api/app/models/chat_generation_task.py new file mode 100644 index 00000000..a36696ec --- /dev/null +++ b/video-gen-api/app/models/chat_generation_task.py @@ -0,0 +1,64 @@ +from datetime import datetime + +from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, Text +from sqlalchemy.orm import Mapped, mapped_column + +from app.models.base import Base, TimestampMixin + + +class ChatGenerationTask(Base, TimestampMixin): + """Project-independent AI chat/image/video generation task. + + This table is intentionally NOT linked to projects. It is used by the + /generation-ai Celery pipeline so chat-style generation does not touch the + legacy generation_records -> projects foreign-key chain. + """ + + __tablename__ = "chat_generation_tasks" + + id: Mapped[str] = mapped_column(String(32), primary_key=True) + user_id: Mapped[str] = mapped_column( + String(32), ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False + ) + + original_prompt: Mapped[str] = mapped_column(Text, nullable=False) + optimized_prompt: Mapped[str | None] = mapped_column(Text, nullable=True) + gen_type: Mapped[str] = mapped_column(String(16), default="video", index=True) + + duration: Mapped[int | None] = mapped_column(Integer, nullable=True) + aspect_ratio: Mapped[str | None] = mapped_column(String(8), nullable=True) + resolution: Mapped[str | None] = mapped_column(String(8), nullable=True) + image_size: Mapped[str | None] = mapped_column(String(16), nullable=True) + image_proportion: Mapped[str | None] = mapped_column(String(8), nullable=True) + image_px: Mapped[str | None] = mapped_column(String(16), nullable=True) + + status: Mapped[str] = mapped_column(String(32), default="generating", index=True) + pipeline_stage: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True) + generation_mode: Mapped[str] = mapped_column(String(32), default="chatapi_async", index=True) + + media_references: Mapped[str | None] = mapped_column(Text, nullable=True) + provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True) + # Kept as alias-compatible storage for existing frontend/service naming. + seedance_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True) + + remote_result_url: Mapped[str | None] = mapped_column(Text, nullable=True) + image_url: Mapped[str | None] = mapped_column(String(512), nullable=True) + video_url: Mapped[str | None] = mapped_column(String(512), nullable=True) + + engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True) + engine_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True) + provider_response_json: Mapped[str | None] = mapped_column(Text, nullable=True) + + credits_cost: Mapped[float] = mapped_column(Float, default=0.0) + text_credits_cost: Mapped[float] = mapped_column(Float, default=0.0) + text_tokens_used: Mapped[int] = mapped_column(Integer, default=0) + video_tokens_used: Mapped[int] = mapped_column(Integer, default=0) + image_tokens_used: Mapped[int] = mapped_column(Integer, default=0) + + retry_count: Mapped[int] = mapped_column(Integer, default=0) + poll_count: Mapped[int] = mapped_column(Integer, default=0) + last_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + generated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + error_message: Mapped[str | None] = mapped_column(Text, nullable=True) + idempotency_key: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True) diff --git a/video-gen-api/app/models/chat_generation_task_event.py b/video-gen-api/app/models/chat_generation_task_event.py new file mode 100644 index 00000000..fd32ac11 --- /dev/null +++ b/video-gen-api/app/models/chat_generation_task_event.py @@ -0,0 +1,25 @@ +from sqlalchemy import DateTime, ForeignKey, String, Text, func +from sqlalchemy.orm import Mapped, mapped_column + +from app.models.base import Base +from app.utils.id_gen import generate_id + + +class ChatGenerationTaskEvent(Base): + """Append-only event log for project-independent chat generation tasks.""" + + __tablename__ = "chat_generation_task_events" + + id: Mapped[str] = mapped_column(String(32), primary_key=True, default=generate_id) + task_id: Mapped[str] = mapped_column( + String(32), ForeignKey("chat_generation_tasks.id", ondelete="CASCADE"), index=True + ) + generation_mode: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True) + event_type: Mapped[str] = mapped_column(String(64), index=True) + from_status: Mapped[str | None] = mapped_column(String(32), nullable=True) + to_status: Mapped[str | None] = mapped_column(String(32), nullable=True) + from_stage: Mapped[str | None] = mapped_column(String(32), nullable=True) + to_stage: Mapped[str | None] = mapped_column(String(32), nullable=True) + message: Mapped[str | None] = mapped_column(String(512), nullable=True) + detail_json: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at = mapped_column(DateTime(timezone=True), server_default=func.now(), index=True) diff --git a/video-gen-api/app/models/chat_provider_call_log.py b/video-gen-api/app/models/chat_provider_call_log.py new file mode 100644 index 00000000..e83f15f8 --- /dev/null +++ b/video-gen-api/app/models/chat_provider_call_log.py @@ -0,0 +1,35 @@ +from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func +from sqlalchemy.orm import Mapped, mapped_column + +from app.models.base import Base +from app.utils.id_gen import generate_id + + +class ChatProviderCallLog(Base): + """Provider call audit log for chat_generation_tasks.""" + + __tablename__ = "chat_provider_call_logs" + + id: Mapped[str] = mapped_column(String(32), primary_key=True, default=generate_id) + task_id: Mapped[str] = mapped_column( + String(32), ForeignKey("chat_generation_tasks.id", ondelete="CASCADE"), index=True + ) + generation_mode: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True) + provider: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True) + api_type: Mapped[str] = mapped_column(String(64), index=True) + model: Mapped[str | None] = mapped_column(String(128), nullable=True) + engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True) + status: Mapped[str] = mapped_column(String(32), index=True) + latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True) + http_status: Mapped[int | None] = mapped_column(Integer, nullable=True) + provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True) + request_hash: Mapped[str | None] = mapped_column(String(64), nullable=True) + response_hash: Mapped[str | None] = mapped_column(String(64), nullable=True) + request_excerpt: Mapped[str | None] = mapped_column(Text, nullable=True) + response_excerpt: Mapped[str | None] = mapped_column(Text, nullable=True) + prompt_tokens: Mapped[int] = mapped_column(Integer, default=0) + completion_tokens: Mapped[int] = mapped_column(Integer, default=0) + total_tokens: Mapped[int] = mapped_column(Integer, default=0) + error_code: Mapped[str | None] = mapped_column(String(128), nullable=True) + error_message: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at = mapped_column(DateTime(timezone=True), server_default=func.now(), index=True) diff --git a/video-gen-api/app/schemas/generation_ai.py b/video-gen-api/app/schemas/generation_ai.py new file mode 100644 index 00000000..a76456dc --- /dev/null +++ b/video-gen-api/app/schemas/generation_ai.py @@ -0,0 +1,418 @@ +from pydantic import BaseModel, ConfigDict, Field + +from app.schemas.common import NaiveDatetimeOptional + + +class GenerationAIReference(BaseModel): + """AI生成任务参考素材。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "url": "https://example.com/reference.png", + "type": "image", + "name": "参考图.png", + } + } + ) + + url: str = Field( + ..., + description="参考素材地址，可以是图片地址或视频地址", + examples=["https://example.com/reference.png"], + ) + type: str = Field( + ..., + description="参考素材类型：image=图片，video=视频", + examples=["image"], + ) + name: str | None = Field( + None, + description="参考素材名称，前端展示用，可为空", + examples=["参考图.png"], + ) + + +class GenerationAITaskCreate(BaseModel): + """创建AI图片/视频生成任务请求体。""" + + model_config = ConfigDict( + json_schema_extra={ + "examples": [ + { + "project_id": None, + "gen_type": "image", + "original_prompt": "生成一张赛博朋克风格的城市夜景", + "engine_id": None, + "media_references": [ + { + "url": "https://example.com/reference.png", + "type": "image", + "name": "参考图.png", + } + ], + "idempotency_key": "frontend-submit-uuid-001", + "image_size": "2K", + "image_proportion": "1:1", + "image_px": "2048x2048", + "duration": None, + "aspect_ratio": None, + "resolution": None, + }, + { + "project_id": None, + "gen_type": "video", + "original_prompt": "生成一段海边日落的电影感视频", + "engine_id": None, + "media_references": None, + "idempotency_key": "frontend-submit-uuid-002", + "image_size": None, + "image_proportion": None, + "image_px": None, + "duration": 4, + "aspect_ratio": "16:9", + "resolution": "480p", + }, + ] + } + ) + + # 新 chat 生成任务不绑定 project_id。为了兼容旧前端误传，保留可选字段但后端不使用。 + project_id: str | None = Field( + None, + description="兼容旧前端字段。当前 /generation-ai 任务不绑定项目，后端不使用该字段，可传 null 或不传", + examples=[None], + ) + gen_type: str = Field( + ..., + description="生成类型：image=图片生成，video=视频生成", + examples=["image"], + ) + original_prompt: str = Field( + ..., + min_length=1, + max_length=5000, + description="用户原始提示词，不能为空，最长5000字符", + examples=["生成一张赛博朋克风格的城市夜景"], + ) + engine_id: str | None = Field( + None, + description="图片/视频引擎ID；为空则使用当前激活且优先级最高的引擎", + examples=[None], + ) + media_references: list[GenerationAIReference] | None = Field( + None, + description="参考素材列表。可以传图片/视频参考素材；为空表示不使用参考素材", + ) + idempotency_key: str | None = Field( + None, + max_length=64, + description=( + "幂等键，用于防止前端重复提交、网络重试导致重复创建任务和重复扣费。" + "同一用户、同一 idempotency_key、同一 generation_mode 下重复请求会返回已有任务。" + "建议前端每次点击生成时生成 UUID；同一次请求失败重试时复用同一个 UUID。" + ), + examples=["frontend-submit-uuid-001"], + ) + + # image params + image_size: str | None = Field( + None, + description="图片分辨率档位，例如：1K、2K。仅图片生成或视频首帧参数需要使用；为空则使用引擎默认值", + examples=["2K"], + ) + image_proportion: str | None = Field( + None, + description="图片比例，例如：1:1、16:9、9:16。仅图片生成或视频首帧参数需要使用；为空则使用默认值", + examples=["1:1"], + ) + image_px: str | None = Field( + None, + description="图片像素尺寸，例如：2048x2048。为空时后端根据 image_size 和 image_proportion 自动匹配", + examples=["2048x2048"], + ) + + # video params + duration: int | None = Field( + None, + description="视频时长，单位秒。仅视频生成使用；为空则使用默认时长", + examples=[4], + ) + aspect_ratio: str | None = Field( + None, + description="视频比例，例如：16:9、9:16、1:1。仅视频生成使用；为空则使用默认比例", + examples=["16:9"], + ) + resolution: str | None = Field( + None, + description="视频分辨率，例如：480p、720p、1080p。仅视频生成使用；为空则使用默认分辨率", + examples=["480p"], + ) + + +class GenerationAITaskOut(BaseModel): + """AI生成任务详情响应体。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "id": "0019e0a44895b6d837d", + "project_id": None, + "gen_type": "image", + "generation_mode": "chatapi_async", + "pipeline_stage": "done", + "status": "completed", + "original_prompt": "生成一张赛博朋克风格的城市夜景", + "optimized_prompt": None, + "duration": None, + "aspect_ratio": None, + "resolution": None, + "image_size": "2K", + "image_proportion": "1:1", + "image_px": "2048x2048", + "media_references": None, + "provider_task_id": "provider_task_xxx", + "seedance_task_id": "provider_task_xxx", + "remote_result_url": None, + "image_url": "https://example.com/result.png", + "video_url": None, + "engine_id": "engine_xxx", + "engine_snapshot": { + "engine_type": "image", + "id": "engine_xxx", + "name": "图片生成引擎", + "provider": "provider_name", + "model_name": "model_name", + "supported_models": [], + "default_size": "2K", + "selected_size": "2K", + "selected_proportion": "1:1", + "selected_px": "2048x2048", + }, + "credits_cost": 10.0, + "text_credits_cost": 0.0, + "text_tokens_used": 0, + "image_tokens_used": 0, + "video_tokens_used": 0, + "retry_count": 0, + "poll_count": 3, + "error_message": None, + "created_at": "2026-05-27T10:12:00", + "generated_at": "2026-05-27T10:15:30", + } + } + ) + + id: str = Field(..., description="生成任务ID") + project_id: str | None = Field( + None, + description="项目ID。当前 /generation-ai 任务不绑定项目，通常为 null", + ) + gen_type: str = Field(..., description="生成类型：image=图片，video=视频") + generation_mode: str | None = Field( + None, + description="生成模式。当前异步Chat生成任务一般为 chatapi_async", + ) + pipeline_stage: str | None = Field( + None, + description=( + "任务流水线阶段，例如：queued=已入队，creating_provider_task=创建第三方任务中，" + "waiting_remote=等待第三方生成，result_ready=远程结果已就绪，downloading=下载中，done=完成" + ), + ) + status: str = Field( + ..., + description="任务状态，例如：generating=生成中，completed=已完成，failed=失败", + ) + original_prompt: str = Field(..., description="用户原始提示词") + optimized_prompt: str | None = Field(None, description="优化后的提示词，可能为空") + duration: int | None = Field(None, description="视频时长，单位秒。图片任务通常为空") + aspect_ratio: str | None = Field(None, description="视频比例，例如 16:9。图片任务通常为空") + resolution: str | None = Field(None, description="视频分辨率，例如 480p。图片任务通常为空") + image_size: str | None = Field(None, description="图片分辨率档位，例如 2K") + image_proportion: str | None = Field(None, description="图片比例，例如 1:1") + image_px: str | None = Field(None, description="图片像素尺寸，例如 2048x2048") + media_references: list[dict] | None = Field( + None, + description="参考素材列表。由创建任务时传入的 media_references 转换而来", + ) + provider_task_id: str | None = Field(None, description="第三方服务商任务ID") + seedance_task_id: str | None = Field( + None, + description="兼容旧前端/旧服务命名的第三方任务ID字段", + ) + remote_result_url: str | None = Field( + None, + description="第三方远程结果地址。当前接口可能隐藏或为空，最终展示优先使用 image_url/video_url", + ) + image_url: str | None = Field(None, description="最终图片地址。图片任务完成后通常有值") + video_url: str | None = Field(None, description="最终视频地址。视频任务完成后通常有值") + engine_id: str | None = Field(None, description="本次任务使用的生成引擎ID") + engine_snapshot: dict | None = Field( + None, + description="生成任务创建时的引擎快照，用于前端展示当时使用的模型、比例、尺寸等信息", + ) + credits_cost: float = Field(0.0, description="本次任务总消耗积分") + text_credits_cost: float = Field(0.0, description="文本优化或文本处理消耗积分") + text_tokens_used: int = Field(0, description="文本 token 使用量") + image_tokens_used: int = Field(0, description="图片 token 使用量") + video_tokens_used: int = Field(0, description="视频 token 使用量") + retry_count: int = Field(0, description="任务重试次数") + poll_count: int = Field(0, description="轮询第三方任务状态次数") + error_message: str | None = Field(None, description="错误信息。成功任务一般为 null") + created_at: NaiveDatetimeOptional = Field(None, description="任务创建时间") + generated_at: NaiveDatetimeOptional = Field(None, description="任务生成完成时间") + + +class GenerationAITaskListOut(BaseModel): + """AI生成任务列表响应体。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "total": 1, + "items": [ + { + "id": "0019e0a44895b6d837d", + "project_id": None, + "gen_type": "image", + "generation_mode": "chatapi_async", + "pipeline_stage": "done", + "status": "completed", + "original_prompt": "生成一张赛博朋克风格的城市夜景", + "optimized_prompt": None, + "duration": None, + "aspect_ratio": None, + "resolution": None, + "image_size": "2K", + "image_proportion": "1:1", + "image_px": "2048x2048", + "media_references": None, + "provider_task_id": "provider_task_xxx", + "seedance_task_id": "provider_task_xxx", + "remote_result_url": None, + "image_url": "https://example.com/result.png", + "video_url": None, + "engine_id": "engine_xxx", + "engine_snapshot": {}, + "credits_cost": 10.0, + "text_credits_cost": 0.0, + "text_tokens_used": 0, + "image_tokens_used": 0, + "video_tokens_used": 0, + "retry_count": 0, + "poll_count": 3, + "error_message": None, + "created_at": "2026-05-27T10:12:00", + "generated_at": "2026-05-27T10:15:30", + } + ], + } + } + ) + + total: int = Field(..., description="符合筛选条件的任务总数") + items: list[GenerationAITaskOut] = Field( + default_factory=list, + description="当前分页的任务列表", + ) + + +class GenerationAIRetryOut(BaseModel): + """AI生成任务重试响应体。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "id": "0019e0a44895b6d837d", + "status": "generating", + "pipeline_stage": "queued", + "message": "任务已重新投递", + } + } + ) + + id: str = Field(..., description="被重试的任务ID") + status: str = Field(..., description="重试后的任务状态") + pipeline_stage: str | None = Field(None, description="重试后的任务流水线阶段") + message: str = Field(..., description="操作结果提示信息") + + +class GenerationAIHistoryDayGroupOut(BaseModel): + """AI生成历史按天分组响应项。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "generated_date": "2026-05-27", + "total": 18, + "items": [], + } + } + ) + + generated_date: str = Field(..., description="生成日期，格式：YYYY-MM-DD") + total: int = Field(..., description="当前生成日期下的生成成功记录总数") + items: list[GenerationAITaskOut] = Field( + default_factory=list, + description="当前生成日期下倒序前10条生成记录详情", + ) + + +class GenerationAIHistoryGroupedOut(BaseModel): + """AI生成历史日期分组列表响应体。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "total_days": 2, + "page": 1, + "page_size": 10, + "groups": [ + { + "generated_date": "2026-05-27", + "total": 18, + "items": [], + }, + { + "generated_date": "2026-05-26", + "total": 6, + "items": [], + }, + ], + } + } + ) + + total_days: int = Field(..., description="当前生成类型下，用户一共有多少个生成日期分组") + page: int = Field(..., description="当前日期分组分页页码") + page_size: int = Field(..., description="当前每页返回的日期分组数量，最大10") + groups: list[GenerationAIHistoryDayGroupOut] = Field( + default_factory=list, + description="按生成日期倒序排列的历史记录分组列表", + ) + + +class GenerationAIHistoryDayItemsOut(BaseModel): + """指定日期下AI生成历史分页响应体。""" + + model_config = ConfigDict( + json_schema_extra={ + "example": { + "generated_date": "2026-05-27", + "total": 18, + "page": 2, + "page_size": 10, + "items": [], + } + } + ) + + generated_date: str = Field(..., description="当前查询的生成日期，格式：YYYY-MM-DD") + total: int = Field(..., description="当前日期下的生成成功记录总数") + page: int = Field(..., description="当前日期下的记录分页页码") + page_size: int = Field(..., description="当前日期下每页返回的生成记录数量") + items: list[GenerationAITaskOut] = Field( + default_factory=list, + description="当前日期下的生成记录详情列表，按 generated_at 倒序排列", + ) \ No newline at end of file diff --git a/video-gen-api/app/services/generation_ai_service.py b/video-gen-api/app/services/generation_ai_service.py new file mode 100644 index 00000000..ddf56a55 --- /dev/null +++ b/video-gen-api/app/services/generation_ai_service.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +import json +from datetime import datetime, timedelta, timezone, date +from typing import Any + +from fastapi import HTTPException +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.config import settings +from app.models.chat_generation_task import ChatGenerationTask +from app.models.image_engine import ImageEngine +from app.models.user import User +from app.models.video_engine import VideoEngine +from app.schemas.generation_ai import GenerationAITaskCreate, GenerationAITaskOut +from app.services.generation_billing_service import charge_generation_media_by_params +from app.utils.id_gen import generate_id + +IMAGE_DEFAULT_SIZE = "2K" +IMAGE_DEFAULT_PROPORTION = "1:1" +IMAGE_DEFAULT_PX = "2048x2048" +VIDEO_DEFAULT_DURATION = 4 +VIDEO_DEFAULT_RATIO = "16:9" +VIDEO_DEFAULT_RESOLUTION = "480p" + +HISTORY_DAY_PAGE_SIZE_MAX = 10 +HISTORY_GROUP_ITEM_LIMIT = 10 + +def normalize_px(value: str | None) -> str | None: + if not value: + return value + return value.replace("×", "x").replace("X", "x").replace("×x", "x").replace("x×", "x") + + +def _json(data: Any) -> str | None: + if data is None: + return None + return json.dumps(data, ensure_ascii=False, default=str) + + +def _parse_json(text: str | None): + if not text: + return None + try: + return json.loads(text) + except Exception: + return None + + +async def _get_image_engine(db: AsyncSession, engine_id: str | None) -> ImageEngine: + query = select(ImageEngine).where(ImageEngine.is_active == True) + if engine_id: + query = query.where(ImageEngine.id == engine_id) + else: + query = query.order_by(ImageEngine.priority.desc()).limit(1) + result = await db.execute(query) + engine = result.scalar_one_or_none() + if not engine: + raise HTTPException(status_code=400, detail="没有可用的图片引擎") + return engine + + +async def _get_video_engine(db: AsyncSession, engine_id: str | None) -> VideoEngine: + query = select(VideoEngine).where(VideoEngine.is_active == True) + if engine_id: + query = query.where(VideoEngine.id == engine_id) + else: + query = query.order_by(VideoEngine.priority.desc()).limit(1) + result = await db.execute(query) + engine = result.scalar_one_or_none() + if not engine: + raise HTTPException(status_code=400, detail="没有可用的视频引擎") + return engine + + +def _image_supported_sizes(engine: ImageEngine) -> dict: + try: + data = json.loads(engine.supported_sizes or "{}") + return data if isinstance(data, dict) else {} + except Exception: + return {} + + +def _parse_list(value: str | None, fallback: list): + try: + parsed = json.loads(value or "") + return parsed if isinstance(parsed, list) else fallback + except Exception: + return fallback + + +def _build_image_snapshot(engine: ImageEngine, size: str, proportion: str, px: str) -> dict: + return { + "engine_type": "image", + "id": engine.id, + "name": engine.name, + "provider": engine.provider, + "api_base": engine.api_base, + "api_key_masked": "****" if engine.api_key else "", + "model_name": engine.model_name, + "generate_url": engine.generate_url, + "supported_models": _parse_list(engine.supported_models, []), + "default_size": engine.default_size, + "selected_size": size, + "selected_proportion": proportion, + "selected_px": px, + } + + +def _build_video_snapshot(engine: VideoEngine, ratio: str, resolution: str, duration: int) -> dict: + return { + "engine_type": "video", + "id": engine.id, + "name": engine.name, + "provider": engine.provider, + "api_base": engine.api_base, + "api_key_masked": "****" if engine.api_key else "", + "model_name": engine.model_name, + "generate_url": engine.generate_url, + "query_url": engine.query_url, + "supported_ratios": _parse_list(engine.supported_ratios, []), + "supported_resolutions": _parse_list(engine.supported_resolutions, []), + "supported_durations": _parse_list(engine.supported_durations, []), + "max_duration": engine.max_duration, + "selected_ratio": ratio, + "selected_resolution": resolution, + "selected_duration": duration, + } + + +async def create_async_generation_task(db: AsyncSession, current_user: User, req: GenerationAITaskCreate) -> ChatGenerationTask: + """Create a project-independent chat generation task. + + Important: this writes chat_generation_tasks, not generation_records, so chat + image/video generation no longer needs or validates a project_id. + """ + gen_type = req.gen_type.lower().strip() + if gen_type not in ("image", "video"): + raise HTTPException(status_code=400, detail="gen_type 仅支持 image 或 video") + + if req.idempotency_key: + result = await db.execute( + select(ChatGenerationTask).where( + ChatGenerationTask.user_id == current_user.id, + ChatGenerationTask.idempotency_key == req.idempotency_key, + ChatGenerationTask.generation_mode == "chatapi_async", + ).order_by(ChatGenerationTask.created_at.desc()).limit(1) + ) + existing = result.scalar_one_or_none() + if existing: + return existing + + refs = [r.model_dump() for r in (req.media_references or [])] + now = datetime.now(timezone.utc) + task_id = generate_id() + + if gen_type == "image": + engine = await _get_image_engine(db, req.engine_id) + sizes = _image_supported_sizes(engine) + size = req.image_size or engine.default_size or IMAGE_DEFAULT_SIZE + proportion = req.image_proportion or IMAGE_DEFAULT_PROPORTION + px = normalize_px(req.image_px) + if sizes: + if size not in sizes: + raise HTTPException(status_code=400, detail=f"图片分辨率档位不支持: {size}") + if proportion not in sizes.get(size, {}): + raise HTTPException(status_code=400, detail=f"图片比例不支持: {proportion}") + px = px or normalize_px((sizes.get(size) or {}).get(proportion)) + px = px or IMAGE_DEFAULT_PX + media_billing = await charge_generation_media_by_params( + db, + user_id=current_user.id, + record_id=task_id, + gen_type="image", + image_size=size, + project_name="AI生成任务", + description_prefix="Chat任务", + ) + snapshot = _build_image_snapshot(engine, size, proportion, px) + task = ChatGenerationTask( + id=task_id, + user_id=current_user.id, + original_prompt=req.original_prompt, + gen_type="image", + image_size=size, + image_proportion=proportion, + image_px=px, + status="generating", + generation_mode="chatapi_async", + pipeline_stage="queued", + engine_id=engine.id, + engine_snapshot_json=_json(snapshot), + media_references=_json(refs) if refs else None, + credits_cost=round(media_billing.total_charged, 2), + idempotency_key=req.idempotency_key, + deadline_at=now + timedelta(minutes=settings.CHATAPI_ASYNC_IMAGE_DEADLINE_MINUTES), + ) + else: + engine = await _get_video_engine(db, req.engine_id) + ratio = req.aspect_ratio or VIDEO_DEFAULT_RATIO + resolution = req.resolution or VIDEO_DEFAULT_RESOLUTION + duration = req.duration or VIDEO_DEFAULT_DURATION + ratios = _parse_list(engine.supported_ratios, []) + resolutions = _parse_list(engine.supported_resolutions, []) + durations = _parse_list(engine.supported_durations, []) + if ratios and ratio not in ratios: + raise HTTPException(status_code=400, detail=f"视频比例不支持: {ratio}") + if resolutions and resolution not in resolutions: + raise HTTPException(status_code=400, detail=f"视频分辨率不支持: {resolution}") + if durations and duration not in durations: + raise HTTPException(status_code=400, detail=f"视频时长不支持: {duration}") + if engine.max_duration and duration > engine.max_duration: + raise HTTPException(status_code=400, detail=f"视频时长不能超过 {engine.max_duration} 秒") + media_billing = await charge_generation_media_by_params( + db, + user_id=current_user.id, + record_id=task_id, + gen_type="video", + duration=duration, + resolution=resolution, + project_name="AI生成任务", + description_prefix="Chat任务", + ) + snapshot = _build_video_snapshot(engine, ratio, resolution, duration) + task = ChatGenerationTask( + id=task_id, + user_id=current_user.id, + original_prompt=req.original_prompt, + gen_type="video", + duration=duration, + aspect_ratio=ratio, + resolution=resolution, + image_size=req.image_size or IMAGE_DEFAULT_SIZE, + image_proportion=req.image_proportion or IMAGE_DEFAULT_PROPORTION, + image_px=normalize_px(req.image_px) or IMAGE_DEFAULT_PX, + status="generating", + generation_mode="chatapi_async", + pipeline_stage="queued", + engine_id=engine.id, + engine_snapshot_json=_json(snapshot), + media_references=_json(refs) if refs else None, + credits_cost=round(media_billing.total_charged, 2), + idempotency_key=req.idempotency_key, + deadline_at=now + timedelta(minutes=settings.CHATAPI_ASYNC_VIDEO_DEADLINE_MINUTES), + ) + + db.add(task) + await db.flush() + return task + + +def record_to_out(task: ChatGenerationTask) -> GenerationAITaskOut: + refs = _parse_json(task.media_references) + snapshot = engine_snapshot_out(_parse_json(task.engine_snapshot_json)) + return GenerationAITaskOut( + id=task.id, + # project_id=None, + gen_type=task.gen_type, + generation_mode=task.generation_mode, + pipeline_stage=task.pipeline_stage, + status=task.status, + original_prompt=task.original_prompt, + # optimized_prompt=task.optimized_prompt, + duration=task.duration, + aspect_ratio=task.aspect_ratio, + resolution=task.resolution, + image_size=task.image_size, + image_proportion=task.image_proportion, + image_px=task.image_px, + media_references=refs, + provider_task_id=task.provider_task_id, + seedance_task_id=task.seedance_task_id, + # remote_result_url=task.remote_result_url, + image_url=task.image_url, + video_url=task.video_url, + engine_id=task.engine_id, + engine_snapshot=snapshot, + credits_cost=task.credits_cost or 0.0, + text_credits_cost=task.text_credits_cost or 0.0, + text_tokens_used=task.text_tokens_used or 0, + image_tokens_used=task.image_tokens_used or 0, + video_tokens_used=task.video_tokens_used or 0, + retry_count=task.retry_count or 0, + poll_count=task.poll_count or 0, + error_message=task.error_message, + created_at=task.created_at, + generated_at=task.generated_at, + ) + + +def engine_snapshot_out(snapshot: dict) -> dict: + """ + 从完整的 engine_snapshot 中过滤出需要返回的字段 + """ + if not snapshot: + return {} + + return { + "engine_type": snapshot.get("engine_type"), + "id": snapshot.get("id"), + "name": snapshot.get("name"), + "provider": snapshot.get("provider"), + # "api_base": snapshot.get("api_base"), + # "api_key_masked": snapshot.get("api_key_masked"), + "model_name": snapshot.get("model_name"), + # "generate_url": snapshot.get("generate_url"), + "supported_models": snapshot.get("supported_models", []), + "default_size": snapshot.get("default_size"), + "selected_size": snapshot.get("selected_size"), + "selected_proportion": snapshot.get("selected_proportion"), + "selected_px": snapshot.get("selected_px") + } + +async def list_async_generation_tasks( + db: AsyncSession, + user_id: str, + gen_type: str | None, + status: str | None, + page: int, + page_size: int, +): + query = select(ChatGenerationTask).where( + ChatGenerationTask.user_id == user_id, + ChatGenerationTask.generation_mode == "chatapi_async", + ) + if gen_type: + query = query.where(ChatGenerationTask.gen_type == gen_type) + if status: + query = query.where(ChatGenerationTask.status == status) + count_query = select(func.count()).select_from(query.subquery()) + total = (await db.execute(count_query)).scalar_one() + result = await db.execute( + query.order_by(ChatGenerationTask.created_at.desc()) + .offset((page - 1) * page_size) + .limit(page_size) + ) + return total, list(result.scalars().all()) + +def _normalize_history_gen_type(gen_type: str | None) -> str: + value = (gen_type or "").lower().strip() + if value not in ("image", "video"): + raise HTTPException(status_code=400, detail="gen_type 仅支持 image 或 video") + return value + + +def _history_day_to_str(value) -> str: + if isinstance(value, datetime): + return value.date().strftime("%Y-%m-%d") + if isinstance(value, date): + return value.strftime("%Y-%m-%d") + return str(value)[:10] + + +def _parse_history_date(value: str) -> date: + try: + return datetime.strptime(value, "%Y-%m-%d").date() + except ValueError: + raise HTTPException(status_code=400, detail="generated_date 格式必须是 YYYY-MM-DD") + + +def _history_base_filters(user_id: str, gen_type: str): + return [ + ChatGenerationTask.user_id == user_id, + ChatGenerationTask.generation_mode == "chatapi_async", + ChatGenerationTask.status == "completed", + ChatGenerationTask.gen_type == gen_type, + ChatGenerationTask.generated_at.is_not(None), + ] + + +async def list_generation_history_grouped_days( + db: AsyncSession, + user_id: str, + gen_type: str, + page: int, + page_size: int, +): + """ + 按生成日期倒序返回历史记录分组。 + + - 每页最多返回 10 个生成日期 + - 每个日期分组内最多返回倒序前 10 条任务 + - 只返回 completed 成功任务 + """ + gen_type = _normalize_history_gen_type(gen_type) + page = max(page, 1) + page_size = min(max(page_size, 1), HISTORY_DAY_PAGE_SIZE_MAX) + + filters = _history_base_filters(user_id, gen_type) + day_expr = func.date(ChatGenerationTask.generated_at).label("generated_date") + + days_subquery = ( + select(day_expr) + .where(*filters) + .group_by(day_expr) + .subquery() + ) + + total_days = ( + await db.execute(select(func.count()).select_from(days_subquery)) + ).scalar_one() + + day_rows_result = await db.execute( + select( + day_expr, + func.count(ChatGenerationTask.id).label("total"), + ) + .where(*filters) + .group_by(day_expr) + .order_by(day_expr.desc()) + .offset((page - 1) * page_size) + .limit(page_size) + ) + day_rows = day_rows_result.all() + + groups = [] + for generated_day, day_total in day_rows: + item_result = await db.execute( + select(ChatGenerationTask) + .where( + *filters, + func.date(ChatGenerationTask.generated_at) == generated_day, + ) + .order_by(ChatGenerationTask.generated_at.desc(), ChatGenerationTask.created_at.desc()) + .limit(HISTORY_GROUP_ITEM_LIMIT) + ) + tasks = list(item_result.scalars().all()) + + groups.append( + { + "generated_date": _history_day_to_str(generated_day), + "total": int(day_total or 0), + "items": [record_to_out(task) for task in tasks], + } + ) + + return { + "total_days": int(total_days or 0), + "page": page, + "page_size": page_size, + "groups": groups, + } + + +async def list_generation_history_day_items( + db: AsyncSession, + user_id: str, + gen_type: str, + generated_date: str, + page: int, + page_size: int, +): + """ + 获取指定生成日期下的历史记录分页。 + + 用于前端点击某一天后，继续加载该日期下的第 2 页、第 3 页数据。 + """ + gen_type = _normalize_history_gen_type(gen_type) + target_day = _parse_history_date(generated_date) + page = max(page, 1) + page_size = min(max(page_size, 1), 100) + + filters = _history_base_filters(user_id, gen_type) + day_expr = func.date(ChatGenerationTask.generated_at) + + total = ( + await db.execute( + select(func.count(ChatGenerationTask.id)).where( + *filters, + day_expr == target_day, + ) + ) + ).scalar_one() + + result = await db.execute( + select(ChatGenerationTask) + .where( + *filters, + day_expr == target_day, + ) + .order_by(ChatGenerationTask.generated_at.desc(), ChatGenerationTask.created_at.desc()) + .offset((page - 1) * page_size) + .limit(page_size) + ) + + tasks = list(result.scalars().all()) + + return { + "generated_date": target_day.strftime("%Y-%m-%d"), + "total": int(total or 0), + "page": page, + "page_size": page_size, + "items": [record_to_out(task) for task in tasks], + } \ No newline at end of file diff --git a/video-gen-api/app/services/generation_billing_service.py b/video-gen-api/app/services/generation_billing_service.py new file mode 100644 index 00000000..daf77743 --- /dev/null +++ b/video-gen-api/app/services/generation_billing_service.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +from dataclasses import dataclass, asdict +from typing import Any, Mapping + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.models.credit_record import CreditRecord +from app.models.generation_record import GenerationRecord +from app.models.system_config import SystemConfig +from app.models.user import User +from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits +from app.utils.exceptions import InsufficientCreditsError +from app.utils.id_gen import generate_id + + +CHARGE_TEXT_PROMPT = "CHATAPI_TEXT_PROMPT" +CHARGE_FILE_PARSE = "CHATAPI_FILE_PARSE" +CHARGE_VISION_INPUT = "CHATAPI_VISION_INPUT" +CHARGE_MEDIA_IMAGE = "CHATAPI_MEDIA_IMAGE" +CHARGE_MEDIA_VIDEO = "CHATAPI_MEDIA_VIDEO" + + +@dataclass +class BillingItem: + charge_key: str + amount: float + charged: bool + skipped_reason: str | None = None + + +@dataclass +class BillingSummary: + record_id: str + user_id: str + items: list[BillingItem] + + @property + def total_charged(self) -> float: + return round(sum(item.amount for item in self.items if item.charged), 2) + + def get_amount(self, charge_key: str) -> float: + return round(sum(item.amount for item in self.items if item.charge_key == charge_key and item.charged), 2) + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["total_charged"] = self.total_charged + return data + + +def _round2(value: float | int | None) -> float: + return round(float(value or 0), 2) + + +def _safe_int(value: Any, default: int = 0) -> int: + try: + if value is None: + return default + return int(value) + except Exception: + return default + + +async def _get_config_float_or_none(db: AsyncSession, key: str) -> float | None: + result = await db.execute(select(SystemConfig).where(SystemConfig.key == key).limit(1)) + config = result.scalar_one_or_none() + if not config: + return None + try: + return float(config.value) + except Exception: + return None + + +async def _calc_optional_token_credits(db: AsyncSession, tokens: int, config_key: str) -> float: + """Calculate optional token billing. Missing config means do not charge. + + This prevents double-charging existing projects where uploaded file/OCR/vision + content is already included in the LLM provider's input_tokens. + """ + tokens = _safe_int(tokens) + if tokens <= 0: + return 0.0 + rate = await _get_config_float_or_none(db, config_key) + if rate is None: + return 0.0 + return round(tokens * rate / 1000, 2) + + +def _legacy_description_keywords(charge_key: str) -> list[str]: + # Compatibility with old patch/original project records that were inserted + # before this safe billing service added [CHARGE_KEY] prefixes. + if charge_key == CHARGE_TEXT_PROMPT: + return ["ChatAPI提示词整理", "提示词优化"] + if charge_key == CHARGE_MEDIA_IMAGE: + return ["ChatAPI异步图片生成", "图片生成"] + if charge_key == CHARGE_MEDIA_VIDEO: + return ["ChatAPI异步视频生成", "视频生成"] + if charge_key == CHARGE_FILE_PARSE: + return ["文件解析Token"] + if charge_key == CHARGE_VISION_INPUT: + return ["图片理解Token"] + return [] + + +async def _find_existing_charge(db: AsyncSession, user_id: str, related_id: str, charge_key: str) -> CreditRecord | None: + base = ( + select(CreditRecord) + .where(CreditRecord.user_id == user_id) + .where(CreditRecord.related_id == related_id) + .where(CreditRecord.type == "consume") + ) + + result = await db.execute(base.where(CreditRecord.description.like(f"[{charge_key}]%")).limit(1)) + existing = result.scalar_one_or_none() + if existing: + return existing + + for keyword in _legacy_description_keywords(charge_key): + result = await db.execute(base.where(CreditRecord.description.like(f"%{keyword}%")).limit(1)) + existing = result.scalar_one_or_none() + if existing: + return existing + return None + + +async def deduct_credits_locked_once( + db: AsyncSession, + *, + user_id: str, + amount: float, + description: str, + related_id: str, + charge_key: str, +) -> BillingItem: + """Deduct credits with row lock and idempotency. + + - User row is locked by SELECT ... FOR UPDATE, so concurrent deductions for + the same user are serialized in PostgreSQL/MySQL. + - CreditRecord description prefix + related_id is used as an idempotency key + without changing existing table structures. + """ + amount = _round2(amount) + if amount <= 0: + return BillingItem(charge_key=charge_key, amount=0.0, charged=False, skipped_reason="amount_lte_zero") + + result = await db.execute(select(User).where(User.id == user_id).with_for_update()) + user = result.scalar_one_or_none() + if not user: + raise ValueError("User not found") + + existing_charge = await _find_existing_charge(db, user_id, related_id, charge_key) + if existing_charge: + return BillingItem( + charge_key=charge_key, + amount=abs(_round2(existing_charge.amount)), + charged=False, + skipped_reason="already_charged", + ) + + if float(user.credits or 0) < amount: + raise InsufficientCreditsError() + + user.credits = round(float(user.credits or 0) - amount, 2) + db.add( + CreditRecord( + id=generate_id(), + user_id=user_id, + type="consume", + amount=-amount, + balance_after=user.credits, + description=f"[{charge_key}] {description}", + related_id=related_id, + ) + ) + await db.flush() + return BillingItem(charge_key=charge_key, amount=amount, charged=True) + + +async def charge_chatapi_prompt_usage( + db: AsyncSession, + *, + record: GenerationRecord, + usage: Mapping[str, Any], + project_name: str | None = None, +) -> BillingSummary: + """Charge ChatAPI prompt optimization and optional uploaded-file/vision tokens. + + file_parse_credits / vision_input_credits are optional and disabled unless + SystemConfig contains these keys: + - file_parse_credits_per_1000_tokens + - vision_input_credits_per_1000_tokens + """ + project_name = project_name or "AI生成任务" + items: list[BillingItem] = [] + + input_tokens = _safe_int(usage.get("input_tokens")) + output_tokens = _safe_int(usage.get("output_tokens")) + text_credits = await calc_text_credits(db, input_tokens, output_tokens) + items.append( + await deduct_credits_locked_once( + db, + user_id=record.user_id, + amount=text_credits, + description=f"ChatAPI提示词整理 - {project_name}", + related_id=record.id, + charge_key=CHARGE_TEXT_PROMPT, + ) + ) + + file_tokens = usage.get("file_parse_tokens") or usage.get("file_tokens") or usage.get("document_tokens") or 0 + file_parse_credits = await _calc_optional_token_credits( + db, + _safe_int(file_tokens), + "file_parse_credits_per_1000_tokens", + ) + items.append( + await deduct_credits_locked_once( + db, + user_id=record.user_id, + amount=file_parse_credits, + description=f"文件解析Token - {project_name}", + related_id=record.id, + charge_key=CHARGE_FILE_PARSE, + ) + ) + + vision_tokens = usage.get("vision_input_tokens") or usage.get("image_input_tokens") or usage.get("image_tokens") or 0 + vision_input_credits = await _calc_optional_token_credits( + db, + _safe_int(vision_tokens), + "vision_input_credits_per_1000_tokens", + ) + items.append( + await deduct_credits_locked_once( + db, + user_id=record.user_id, + amount=vision_input_credits, + description=f"图片理解Token - {project_name}", + related_id=record.id, + charge_key=CHARGE_VISION_INPUT, + ) + ) + + if hasattr(record, "text_credits_cost"): + # Store expected text-side cost, even when this task is a retry and the + # actual CreditRecord was already written by an earlier attempt. + record.text_credits_cost = round(text_credits + file_parse_credits + vision_input_credits, 2) + if hasattr(record, "text_tokens_used"): + record.text_tokens_used = _safe_int(usage.get("total_tokens"), input_tokens + output_tokens) + + return BillingSummary(record_id=record.id, user_id=record.user_id, items=items) + + +async def charge_generation_media_by_params( + db: AsyncSession, + *, + user_id: str, + record_id: str, + gen_type: str, + image_size: str | None = None, + duration: int | None = None, + resolution: str | None = None, + project_name: str | None = None, + description_prefix: str = "ChatAPI异步", +) -> BillingSummary: + """Charge image/video generation fee safely before creating provider task.""" + project_name = project_name or "AI生成任务" + gen_type = (gen_type or "").lower().strip() + items: list[BillingItem] = [] + + if gen_type == "image": + size = image_size or "2K" + amount = await calc_image_credits(db, size) + items.append( + await deduct_credits_locked_once( + db, + user_id=user_id, + amount=amount, + description=f"{description_prefix}图片生成 - {project_name}", + related_id=record_id, + charge_key=CHARGE_MEDIA_IMAGE, + ) + ) + elif gen_type == "video": + amount = await calc_video_credits(db, duration or 5, resolution or "720p") + items.append( + await deduct_credits_locked_once( + db, + user_id=user_id, + amount=amount, + description=f"{description_prefix}视频生成 - {project_name}", + related_id=record_id, + charge_key=CHARGE_MEDIA_VIDEO, + ) + ) + else: + raise ValueError(f"不支持的生成类型: {gen_type}") + + return BillingSummary(record_id=record_id, user_id=user_id, items=items) + + +async def charge_generation_media_for_record( + db: AsyncSession, + *, + record: GenerationRecord, + project_name: str | None = None, + description_prefix: str = "ChatAPI异步", +) -> BillingSummary: + return await charge_generation_media_by_params( + db, + user_id=record.user_id, + record_id=record.id, + gen_type=record.gen_type, + image_size=record.image_size, + duration=record.duration, + resolution=record.resolution, + project_name=project_name, + description_prefix=description_prefix, + ) diff --git a/video-gen-api/app/services/generation_download_service.py b/video-gen-api/app/services/generation_download_service.py new file mode 100644 index 00000000..0e2ffb9e --- /dev/null +++ b/video-gen-api/app/services/generation_download_service.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import os +from datetime import datetime + +from app.config import settings +from app.models.chat_generation_task import ChatGenerationTask +from app.services.image_gen import download_image +from app.services.provider_limit import provider_limit +from app.services.video_gen import download_video + + +async def download_generation_result(record: ChatGenerationTask) -> str: + if not record.remote_result_url: + raise ValueError("缺少远程结果URL") + + date_dir = datetime.now().strftime("%Y/%m/%d") + if record.gen_type == "image": + dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir) + os.makedirs(dest_dir, exist_ok=True) + dest = os.path.join(dest_dir, f"{record.id}.png") + async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY): + await download_image(record.remote_result_url, dest) + return f"/images/{date_dir}/{record.id}.png" + + dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir) + os.makedirs(dest_dir, exist_ok=True) + dest = os.path.join(dest_dir, f"{record.id}.mp4") + async with provider_limit("result_download", settings.RESULT_DOWNLOAD_MAX_CONCURRENCY): + await download_video(record.remote_result_url, dest) + return f"/videos/{date_dir}/{record.id}.mp4" diff --git a/video-gen-api/app/services/generation_log_service.py b/video-gen-api/app/services/generation_log_service.py new file mode 100644 index 00000000..0a0d1698 --- /dev/null +++ b/video-gen-api/app/services/generation_log_service.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import hashlib +import json +from typing import Any + +from app.models.base import async_session +from app.models.chat_generation_task_event import ChatGenerationTaskEvent +from app.models.chat_provider_call_log import ChatProviderCallLog +from app.utils.id_gen import generate_id + +MAX_EXCERPT_CHARS = 2000 + + +def _safe_json(data: Any) -> str | None: + if data is None: + return None + try: + return json.dumps(data, ensure_ascii=False, default=str) + except Exception: + return str(data) + + +def _excerpt(data: Any, limit: int = MAX_EXCERPT_CHARS) -> str | None: + text = _safe_json(data) + if text is None: + return None + # Avoid storing secrets in logs. + text = text.replace("Authorization", "Authorization-REDACTED") + text = text.replace("api_key", "api_key_REDACTED") + if len(text) > limit: + return text[:limit] + "...[truncated]" + return text + + +def _hash(data: Any) -> str | None: + text = _safe_json(data) + if text is None: + return None + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +async def log_task_event( + task: Any | None = None, + *, + record: Any | None = None, + task_id: str | None = None, + record_id: str | None = None, + event_type: str, + from_status: str | None = None, + to_status: str | None = None, + from_stage: str | None = None, + to_stage: str | None = None, + message: str | None = None, + detail: Any = None, +) -> None: + """Write task event in a separate transaction; failure must not affect main flow.""" + try: + obj = task or record + tid = task_id or record_id or (obj.id if obj else None) + if not tid: + return + async with async_session() as db: + db.add(ChatGenerationTaskEvent( + id=generate_id(), + task_id=tid, + generation_mode=getattr(obj, "generation_mode", "chatapi_async"), + event_type=event_type, + from_status=from_status, + to_status=to_status, + from_stage=from_stage, + to_stage=to_stage, + message=message, + detail_json=_excerpt(detail), + )) + await db.commit() + except Exception: + return + + +async def log_provider_call( + task: Any | None = None, + *, + record: Any | None = None, + task_id: str | None = None, + record_id: str | None = None, + provider: str | None, + api_type: str, + model: str | None = None, + engine_id: str | None = None, + status: str, + latency_ms: int | None = None, + http_status: int | None = None, + provider_task_id: str | None = None, + request_data: Any = None, + response_data: Any = None, + prompt_tokens: int = 0, + completion_tokens: int = 0, + total_tokens: int = 0, + error_code: str | None = None, + error_message: str | None = None, +) -> None: + """Write provider call log in a separate transaction; failure must not affect main flow.""" + try: + obj = task or record + tid = task_id or record_id or (obj.id if obj else None) + if not tid: + return + async with async_session() as db: + db.add(ChatProviderCallLog( + id=generate_id(), + task_id=tid, + generation_mode=getattr(obj, "generation_mode", "chatapi_async"), + provider=provider, + api_type=api_type, + model=model, + engine_id=engine_id, + status=status, + latency_ms=latency_ms, + http_status=http_status, + provider_task_id=provider_task_id, + request_hash=_hash(request_data), + response_hash=_hash(response_data), + request_excerpt=_excerpt(request_data), + response_excerpt=_excerpt(response_data), + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=total_tokens or 0, + error_code=error_code, + error_message=error_message, + )) + await db.commit() + except Exception: + return diff --git a/video-gen-api/app/services/generation_prompt_service.py b/video-gen-api/app/services/generation_prompt_service.py new file mode 100644 index 00000000..25c7408f --- /dev/null +++ b/video-gen-api/app/services/generation_prompt_service.py @@ -0,0 +1,184 @@ +from __future__ import annotations + +import json +import mimetypes +import os +import time +from typing import Any + +import httpx +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.config import settings +from app.models.chat_generation_task import ChatGenerationTask +from app.models.model_config import ModelConfig +from app.models.token_usage import TokenUsage +from app.services.generation_log_service import log_provider_call +from app.services.provider_limit import provider_limit +from app.utils.id_gen import generate_id + + +def _absolute_url(url: str) -> str: + if url.startswith("http://") or url.startswith("https://") or url.startswith("data:"): + return url + base = settings.BASE_URL.rstrip("/") + return f"{base}/{url.lstrip('/')}" + + +def _load_refs(record: ChatGenerationTask) -> list[dict]: + if not record.media_references: + return [] + try: + data = json.loads(record.media_references) + return data if isinstance(data, list) else [] + except Exception: + return [] + + +def _build_user_content(record: ChatGenerationTask) -> list[dict[str, Any]]: + if record.gen_type == "image": + params = f"图片参数：分辨率档位={record.image_size or '2K'}，比例={record.image_proportion or '1:1'}，像素={record.image_px or '2048x2048'}" + else: + params = f"视频参数：时长={record.duration or 4}秒，比例={record.aspect_ratio or '16:9'}，分辨率={record.resolution or '480p'}" + + text = ( + f"生成类型：{record.gen_type}\n" + f"{params}\n" + f"用户描述：{record.original_prompt}\n\n" + "请只输出最终可直接用于图片/视频生成模型的 prompt，不要说你已经生成了图片或视频。" + ) + parts: list[dict[str, Any]] = [{"type": "text", "text": text}] + for ref in _load_refs(record): + ref_type = ref.get("type") + ref_url = ref.get("url") or "" + if not ref_url: + continue + url = _absolute_url(ref_url) + if ref_type == "image": + parts.append({"type": "image_url", "image_url": {"url": url}}) + elif ref_type == "video": + parts.append({"type": "video_url", "video_url": {"url": url, "fps": settings.CHATAPI_VIDEO_FPS}}) + return parts + + +async def _get_model_config(db: AsyncSession) -> ModelConfig: + result = await db.execute( + select(ModelConfig) + .where(ModelConfig.is_active == True) + .order_by(ModelConfig.priority.desc()) + .limit(1) + ) + config = result.scalar_one_or_none() + if not config: + raise ValueError("没有可用的ChatAPI模型配置") + if config.provider == "mock": + return config + if not config.api_base or not config.api_key or not config.model_name: + raise ValueError("ChatAPI模型配置不完整") + return config + + +async def build_prompt_with_chatapi(db: AsyncSession, record: ChatGenerationTask) -> tuple[str, dict]: + """Call ChatAPI once with current request params and attachments. No history context.""" + config = await _get_model_config(db) + if config.provider == "mock": + return record.original_prompt, {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + + system_prompt = ( + "你是图片/视频生成提示词整理助手。你的职责是根据用户文字、上传图片/视频和生成参数，" + "整理最终可直接用于生成模型的 prompt。不要声称你已经生成图片或视频，不要调用工具。" + "输出中文为主，内容具体、可执行，保留用户关键要求。" + ) + request_data = { + "model": config.model_name, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": _build_user_content(record)}, + ], + "max_tokens": config.max_tokens, + "temperature": config.temperature, + } + started = time.perf_counter() + async with provider_limit("ark_chat_prompt", settings.ARK_CHAT_PROMPT_MAX_CONCURRENCY): + async with httpx.AsyncClient(timeout=settings.CHATAPI_REQUEST_TIMEOUT_SECONDS) as client: + try: + response = await client.post( + f"{config.api_base.rstrip('/')}/chat/completions", + headers={ + "Authorization": f"Bearer {config.api_key}", + "Content-Type": "application/json", + }, + json=request_data, + ) + latency_ms = int((time.perf_counter() - started) * 1000) + if response.status_code >= 400: + await log_provider_call( + record, + provider=config.provider, + api_type="chat_prompt", + model=config.model_name, + engine_id=record.engine_id, + status="failed", + latency_ms=latency_ms, + http_status=response.status_code, + request_data=request_data, + response_data=response.text, + error_message=response.text[:1000], + ) + raise RuntimeError(f"ChatAPI HTTP {response.status_code}: {response.text}") + data = response.json() + except Exception as exc: + latency_ms = int((time.perf_counter() - started) * 1000) + await log_provider_call( + record, + provider=config.provider, + api_type="chat_prompt", + model=config.model_name, + engine_id=record.engine_id, + status="failed", + latency_ms=latency_ms, + request_data=request_data, + response_data=None, + error_message=str(exc), + ) + raise + + usage = data.get("usage", {}) or {} + input_tokens = int(usage.get("prompt_tokens", 0) or 0) + output_tokens = int(usage.get("completion_tokens", 0) or 0) + total_tokens = int(usage.get("total_tokens", input_tokens + output_tokens) or 0) + content = data.get("choices", [{}])[0].get("message", {}).get("content", "").strip() + if not content: + raise RuntimeError("ChatAPI未返回有效prompt") + + db.add(TokenUsage( + id=generate_id(), + model_config_id=config.id, + user_id=record.user_id, + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + )) + await db.flush() + + await log_provider_call( + record, + provider=config.provider, + api_type="chat_prompt", + model=config.model_name, + engine_id=record.engine_id, + status="success", + latency_ms=int((time.perf_counter() - started) * 1000), + http_status=200, + request_data=request_data, + response_data=data, + prompt_tokens=input_tokens, + completion_tokens=output_tokens, + total_tokens=total_tokens, + ) + return content, { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + } diff --git a/video-gen-api/app/services/generation_provider_service.py b/video-gen-api/app/services/generation_provider_service.py new file mode 100644 index 00000000..9dd53fb1 --- /dev/null +++ b/video-gen-api/app/services/generation_provider_service.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import asyncio +import json +import time +from types import SimpleNamespace +from typing import Any + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.config import settings +from app.models.chat_generation_task import ChatGenerationTask +from app.models.image_engine import ImageEngine +from app.models.video_engine import VideoEngine +from app.services.generation_log_service import log_provider_call +from app.services.image_gen import poll_image_task_status, submit_image_task +from app.services.provider_limit import provider_limit +from app.services.video_gen import poll_task_status, submit_video_task + + +def _loads(data: str | None) -> dict: + if not data: + return {} + try: + obj = json.loads(data) + return obj if isinstance(obj, dict) else {} + except Exception: + return {} + + +async def get_runtime_engine(db: AsyncSession, task: ChatGenerationTask) -> Any: + """Use frozen snapshot for historical params, current DB row only for secret api_key.""" + snapshot = _loads(task.engine_snapshot_json) + if not task.engine_id: + raise ValueError("缺少 engine_id") + if task.gen_type == "image": + result = await db.execute(select(ImageEngine).where(ImageEngine.id == task.engine_id)) + else: + result = await db.execute(select(VideoEngine).where(VideoEngine.id == task.engine_id)) + engine = result.scalar_one_or_none() + if not engine: + raise ValueError("引擎不存在或已删除") + return SimpleNamespace( + id=task.engine_id, + name=snapshot.get("name") or engine.name, + provider=snapshot.get("provider") or engine.provider, + api_base=snapshot.get("api_base") or engine.api_base, + api_key=engine.api_key, + model_name=snapshot.get("model_name") or engine.model_name, + generate_url=snapshot.get("generate_url") or getattr(engine, "generate_url", ""), + query_url=snapshot.get("query_url") or getattr(engine, "query_url", ""), + default_size=snapshot.get("default_size") or getattr(engine, "default_size", "2K"), + ) + + +async def create_provider_task(db: AsyncSession, task: ChatGenerationTask) -> dict: + if task.gen_type == "video": + return await _create_video_task(db, task) + if task.gen_type == "image": + return await _create_image_sync_task(db, task) + raise ValueError(f"不支持的生成类型: {task.gen_type}") + + +async def _create_video_task(db: AsyncSession, task: ChatGenerationTask) -> dict: + """Create video provider task through the original Ark SDK async task API.""" + engine = await get_runtime_engine(db, task) + started = time.perf_counter() + async with provider_limit("ark_video_create", settings.ARK_VIDEO_CREATE_MAX_CONCURRENCY): + try: + provider_task_id = await submit_video_task(db, engine, task) + response = {"task_id": provider_task_id} + await log_provider_call( + task, + provider=engine.provider, + api_type="video_create", + model=engine.model_name, + engine_id=task.engine_id, + status="success", + latency_ms=int((time.perf_counter() - started) * 1000), + provider_task_id=provider_task_id, + response_data=response, + ) + return {"task_id": provider_task_id, "response_data": response} + except Exception as exc: + await log_provider_call( + task, + provider=engine.provider, + api_type="video_create", + model=engine.model_name, + engine_id=task.engine_id, + status="failed", + latency_ms=int((time.perf_counter() - started) * 1000), + error_message=str(exc), + ) + raise + + +async def _create_image_sync_task(db: AsyncSession, task: ChatGenerationTask) -> dict: + """Run the original synchronous image generation SDK under Celery control. + + The legacy image SDK returns a final remote image URL immediately. We do + NOT use image_generation.tasks.create here, so image generation stays aligned + with the old working flow while no longer blocking the FastAPI request. + """ + engine = await get_runtime_engine(db, task) + started = time.perf_counter() + async with provider_limit("ark_image_sync_create", settings.ARK_IMAGE_CREATE_MAX_CONCURRENCY): + try: + result = await asyncio.to_thread(submit_image_task, db, engine, task) + if result.get("error"): + raise RuntimeError(result.get("error")) + response_data = _try_json(result.get("response_data")) or result + await log_provider_call( + task, + provider=engine.provider, + api_type="image_sync_create", + model=engine.model_name, + engine_id=task.engine_id, + status="success", + latency_ms=int((time.perf_counter() - started) * 1000), + provider_task_id=None, + response_data=response_data, + ) + return { + "task_id": None, + "remote_result_url": result.get("image_url"), + "image_tokens": result.get("image_tokens", 0) or 0, + "response_data": response_data, + } + except Exception as exc: + await log_provider_call( + task, + provider=engine.provider, + api_type="image_sync_create", + model=engine.model_name, + engine_id=task.engine_id, + status="failed", + latency_ms=int((time.perf_counter() - started) * 1000), + error_message=str(exc), + ) + raise + + +def _try_json(text: Any) -> Any: + if not isinstance(text, str): + return text + try: + return json.loads(text) + except Exception: + return None + + +async def poll_provider_task(db: AsyncSession, task: ChatGenerationTask) -> dict: + engine = await get_runtime_engine(db, task) + task_id = task.seedance_task_id or task.provider_task_id + if task.gen_type == "video": + async with provider_limit("ark_video_poll", settings.ARK_VIDEO_POLL_MAX_CONCURRENCY): + return await poll_task_status(engine, task_id) + async with provider_limit("ark_image_poll", settings.ARK_IMAGE_POLL_MAX_CONCURRENCY): + return await poll_image_task_status(engine, task_id) diff --git a/video-gen-api/app/services/image_gen.py b/video-gen-api/app/services/image_gen.py index b2d5aabf..47e3c3c8 100644 --- a/video-gen-api/app/services/image_gen.py +++ b/video-gen-api/app/services/image_gen.py @@ -95,7 +95,9 @@ def _resolve_url(url: str) -> str: # return f"data:{mime};base64,{b64}" - return settings.BASE_URL + url + if url.startswith(("http://", "https://", "data:")): + return url + return f"{settings.BASE_URL.rstrip('/')}/{url.lstrip('/')}" def submit_image_task( diff --git a/video-gen-api/app/services/provider_limit.py b/video-gen-api/app/services/provider_limit.py new file mode 100644 index 00000000..6eb5e299 --- /dev/null +++ b/video-gen-api/app/services/provider_limit.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import asyncio +import time +from contextlib import asynccontextmanager +from uuid import uuid4 + +from app.config import settings + + +class ProviderLimitTimeout(TimeoutError): + pass + + +@asynccontextmanager +async def provider_limit(name: str, limit: int | None = None, wait_timeout: float | None = None, ttl: int | None = None): + """Best-effort Redis distributed semaphore. + + If Redis is disabled, it becomes a no-op. This keeps local development simple. + """ + limit = limit or 0 + wait_timeout = wait_timeout if wait_timeout is not None else settings.PROVIDER_LIMIT_WAIT_TIMEOUT_SECONDS + ttl = ttl or settings.PROVIDER_LIMIT_TOKEN_TTL_SECONDS + + if limit <= 0 or not settings.REDIS_URL: + yield + return + + import redis.asyncio as redis + + client = redis.from_url(settings.REDIS_URL, encoding="utf-8", decode_responses=True) + key = f"provider_limit:{name}" + token = str(uuid4()) + acquired = False + deadline = time.monotonic() + wait_timeout + try: + while time.monotonic() < deadline: + now = time.time() + pipe = client.pipeline() + pipe.zremrangebyscore(key, 0, now - ttl) + pipe.zcard(key) + _, count = await pipe.execute() + if count < limit: + added = await client.zadd(key, {token: now}, nx=True) + await client.expire(key, ttl) + if added: + acquired = True + break + await asyncio.sleep(0.2) + if not acquired: + raise ProviderLimitTimeout(f"provider limit exceeded: {name}") + yield + finally: + if acquired: + try: + await client.zrem(key, token) + except Exception: + pass + await client.aclose() diff --git a/video-gen-api/app/services/video_gen.py b/video-gen-api/app/services/video_gen.py index b3a41706..43b1e7f5 100644 --- a/video-gen-api/app/services/video_gen.py +++ b/video-gen-api/app/services/video_gen.py @@ -96,7 +96,9 @@ def _resolve_url(url: str) -> str: # with open(file_path, "rb") as f: # b64 = base64.b64encode(f.read()).decode() # return f"data:{mime};base64,{b64}" - return settings.BASE_URL + url + if url.startswith(("http://", "https://", "data:")): + return url + return f"{settings.BASE_URL.rstrip('/')}/{url.lstrip('/')}" async def submit_video_task( diff --git a/video-gen-api/app/tasks/__init__.py b/video-gen-api/app/tasks/__init__.py index e69de29b..f86ca85d 100644 --- a/video-gen-api/app/tasks/__init__.py +++ b/video-gen-api/app/tasks/__init__.py @@ -0,0 +1,11 @@ +"""Celery task module imports. + +Celery autodiscover imports ``app.tasks``; importing task modules here ensures +custom named tasks are registered when workers start. +""" + +try: + from app.tasks import generation_create_tasks, generation_poll_tasks, generation_download_tasks # noqa: F401 +except Exception: + # Keep application importable even when optional Celery dependencies/config are absent. + pass diff --git a/video-gen-api/app/tasks/async_runner.py b/video-gen-api/app/tasks/async_runner.py new file mode 100644 index 00000000..7d6f1df8 --- /dev/null +++ b/video-gen-api/app/tasks/async_runner.py @@ -0,0 +1,48 @@ +import asyncio +import os +import threading +from typing import Awaitable, TypeVar + +T = TypeVar("T") + +_thread_local = threading.local() + + +def _get_or_create_loop() -> asyncio.AbstractEventLoop: + """ + 给当前进程/线程维护一个长期 event loop。 + + Linux prefork： + 每个 Celery 子进程通常单线程跑任务，这里相当于每个子进程一个长期 loop。 + + Windows -P threads： + 每个线程一个 loop，但注意 asyncpg pool 仍不适合跨线程共享； + Windows threads 模式建议继续用 NullPool 或只做本地调试。 + """ + pid = os.getpid() + loop = getattr(_thread_local, "loop", None) + loop_pid = getattr(_thread_local, "pid", None) + + if loop is None or loop.is_closed() or loop_pid != pid: + loop = asyncio.new_event_loop() + _thread_local.loop = loop + _thread_local.pid = pid + + return loop + + +def run_async(coro: Awaitable[T]) -> T: + """ + Celery 同步 task 调用异步协程的统一入口。 + 不使用 asyncio.run()，避免每个 task 结束时关闭 event loop。 + """ + loop = _get_or_create_loop() + return loop.run_until_complete(coro) + + +def close_loop() -> None: + loop = getattr(_thread_local, "loop", None) + if loop is not None and not loop.is_closed(): + loop.close() + _thread_local.loop = None + _thread_local.pid = None \ No newline at end of file diff --git a/video-gen-api/app/tasks/celery_app.py b/video-gen-api/app/tasks/celery_app.py index f43570c0..b7150959 100644 --- a/video-gen-api/app/tasks/celery_app.py +++ b/video-gen-api/app/tasks/celery_app.py @@ -1,17 +1,79 @@ from celery import Celery from app.config import settings -if settings.REDIS_URL: +from celery.signals import worker_process_init, worker_process_shutdown + +from app.tasks.async_runner import run_async, close_loop +from app.models.base import engine + + +def _derive_redis_db(url: str, db_no: int) -> str: + if not url: + return url + import re + if re.search(r"/\d+$", url): + return re.sub(r"/\d+$", f"/{db_no}", url) + return url.rstrip("/") + f"/{db_no}" + + +broker_url = settings.CELERY_BROKER_URL or (_derive_redis_db(settings.REDIS_URL, 1) if settings.REDIS_URL else "") +backend_url = settings.CELERY_RESULT_BACKEND or (_derive_redis_db(settings.REDIS_URL, 2) if settings.REDIS_URL else "") + +if broker_url: celery_app = Celery("videogen") celery_app.conf.update( - broker_url=settings.REDIS_URL.replace("/0", "/1"), - result_backend=settings.REDIS_URL.replace("/0", "/2"), + broker_url=broker_url, + result_backend=backend_url or broker_url, task_serializer="json", accept_content=["json"], + result_serializer="json", + timezone="Asia/Shanghai", + enable_utc=True, task_soft_time_limit=600, task_time_limit=900, + task_acks_late=True, + task_reject_on_worker_lost=True, + task_track_started=True, worker_prefetch_multiplier=1, + broker_transport_options={ + "visibility_timeout": 3600, + }, + task_routes={ + "generation.chatapi_create_generation_task": {"queue": "gen_chatapi_create"}, + "generation.poll_generation_task": {"queue": "gen_provider_poll"}, + "generation.download_generation_result_task": {"queue": "gen_result_download"}, + # "generation.reconcile_generation_billing_task": {"queue": "gen_billing"}, + "app.tasks.cleanup.*": {"queue": "default"}, + }, ) celery_app.autodiscover_tasks(["app.tasks"]) else: celery_app = None + + +@worker_process_init.connect +def on_worker_process_init(**kwargs): + """ + Linux prefork 子进程启动后执行。 + + 目的： + 1. 丢弃 fork 前可能继承的连接池状态。 + 2. 后续任务会在当前子进程自己的长期 event loop 上重新建连接池。 + """ + try: + run_async(engine.dispose()) + except Exception: + pass + + +@worker_process_shutdown.connect +def on_worker_process_shutdown(**kwargs): + """ + 子进程退出前关闭连接池和 event loop。 + """ + try: + run_async(engine.dispose()) + except Exception: + pass + finally: + close_loop() \ No newline at end of file diff --git a/video-gen-api/app/tasks/generation_create_tasks.py b/video-gen-api/app/tasks/generation_create_tasks.py new file mode 100644 index 00000000..51c7108f --- /dev/null +++ b/video-gen-api/app/tasks/generation_create_tasks.py @@ -0,0 +1,113 @@ +from app.tasks.async_runner import run_async +import json +from datetime import datetime, timezone + +from sqlalchemy import select + +from app.models.base import async_session +from app.models.chat_generation_task import ChatGenerationTask +from app.services.error_codes import extract_error_message +from app.services.generation_billing_service import charge_chatapi_prompt_usage +from app.services.generation_log_service import log_task_event +from app.services.generation_prompt_service import build_prompt_with_chatapi +from app.services.generation_provider_service import create_provider_task +from app.tasks.celery_app import celery_app + + +async def _run(task_id: str): + async with async_session() as db: + result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id)) + task = result.scalar_one_or_none() + if not task or task.generation_mode != "chatapi_async": + return + if task.status != "generating": + return + if task.deadline_at and datetime.now(timezone.utc) > task.deadline_at: + task.status = "failed" + task.pipeline_stage = "timeout" + task.error_message = "任务超时" + await db.commit() + await log_task_event(task, event_type="TASK_TIMEOUT", to_status="failed", to_stage="timeout") + return + if task.pipeline_stage not in ("queued", "preparing", "creating_provider_task"): + return + + try: + if not task.optimized_prompt: + old_stage = task.pipeline_stage + task.pipeline_stage = "preparing" + await db.commit() + await log_task_event(task, event_type="CHATAPI_START", from_stage=old_stage, to_stage="preparing") + + optimized, usage = await build_prompt_with_chatapi(db, task) + await charge_chatapi_prompt_usage(db, record=task, usage=usage, project_name="AI生成任务") + task.optimized_prompt = optimized + task.text_tokens_used = usage["total_tokens"] + await db.commit() + await log_task_event(task, event_type="CHATAPI_SUCCESS", to_stage="preparing", detail=usage) + + if task.seedance_task_id or task.provider_task_id: + task.pipeline_stage = "waiting_remote" + await db.commit() + elif task.remote_result_url: + task.pipeline_stage = "result_ready" + await db.commit() + from app.tasks.generation_download_tasks import download_generation_result_task + download_generation_result_task.delay(task.id) + return + else: + old_stage = task.pipeline_stage + task.pipeline_stage = "creating_provider_task" + await db.commit() + await log_task_event(task, event_type="PROVIDER_CREATE_START", from_stage=old_stage, to_stage="creating_provider_task") + + created = await create_provider_task(db, task) + provider_task_id = created.get("task_id") + if provider_task_id: + task.provider_task_id = provider_task_id + task.seedance_task_id = provider_task_id + task.remote_result_url = created.get("remote_result_url") or task.remote_result_url + if task.gen_type == "image": + task.image_tokens_used = created.get("image_tokens", task.image_tokens_used or 0) or 0 + task.provider_response_json = json.dumps(created.get("response_data") or {}, ensure_ascii=False, default=str) + if task.remote_result_url and not task.seedance_task_id: + # Sync image path: original SDK already returned final URL. + task.pipeline_stage = "result_ready" + else: + # Video path: provider returns task id, poll later. + task.pipeline_stage = "waiting_remote" + task.status = "generating" + await db.commit() + await log_task_event(task, event_type="PROVIDER_CREATE_SUCCESS", to_stage=task.pipeline_stage, detail=created) + + if task.pipeline_stage == "result_ready": + from app.tasks.generation_download_tasks import download_generation_result_task + download_generation_result_task.delay(task.id) + else: + from app.tasks.generation_poll_tasks import poll_generation_task + poll_generation_task.delay(task.id) + except Exception as exc: + try: + await db.rollback() + except Exception: + pass + result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id)) + task = result.scalar_one_or_none() + if task: + task.status = "failed" + task.error_message = extract_error_message(exc, "生成任务") if callable(extract_error_message) else str(exc) + await db.commit() + await log_task_event(task, event_type="TASK_FAILED", message=task.error_message) + + +if celery_app: + @celery_app.task(name="generation.chatapi_create_generation_task", bind=True, max_retries=3, default_retry_delay=30) + def chatapi_create_generation_task(self, task_id: str): + return run_async(_run(task_id)) +else: + class _DisabledTask: + def delay(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + def apply_async(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + chatapi_create_generation_task = _DisabledTask() diff --git a/video-gen-api/app/tasks/generation_download_tasks.py b/video-gen-api/app/tasks/generation_download_tasks.py new file mode 100644 index 00000000..81f4d174 --- /dev/null +++ b/video-gen-api/app/tasks/generation_download_tasks.py @@ -0,0 +1,176 @@ +from app.tasks.async_runner import run_async +from datetime import datetime, timezone, timedelta + +from sqlalchemy import select + +from app.models.base import async_session +from app.models.chat_generation_task import ChatGenerationTask +from app.services.error_codes import extract_error_message +from app.services.generation_download_service import download_generation_result +from app.services.generation_log_service import log_task_event +from app.tasks.celery_app import celery_app + + +# downloading 卡住多久后允许自动恢复。 +# 说明： +# - worker 在 pipeline_stage 改成 downloading 后，如果被 kill，任务可能永远停在 downloading。 +# - 这里允许超过该时间的 downloading 任务重新进入下载流程。 +# - 如果你的视频文件特别大，可以把这个时间调大，比如 20 * 60。 +DOWNLOAD_STUCK_SECONDS = 10 * 60 + + +def _to_aware_utc(dt): + """ + 把 datetime 统一转成 timezone-aware UTC，避免 offset-naive 和 offset-aware 比较报错。 + PostgreSQL / SQLite / 不同驱动返回的 updated_at 可能有时区，也可能没有。 + """ + if not dt: + return None + if dt.tzinfo is None: + return dt.replace(tzinfo=timezone.utc) + return dt.astimezone(timezone.utc) + + +def _is_recent_downloading(task: ChatGenerationTask) -> bool: + """ + 判断 downloading 是否仍然是较新的下载任务。 + + 返回 True： + - 说明可能有另一个 worker 刚进入下载，不要重复下载。 + + 返回 False： + - 说明 downloading 已经超过 DOWNLOAD_STUCK_SECONDS，认为可能卡死，可以恢复。 + """ + updated_at = _to_aware_utc(getattr(task, "updated_at", None)) + if not updated_at: + return False + + return datetime.now(timezone.utc) - updated_at < timedelta(seconds=DOWNLOAD_STUCK_SECONDS) + + +async def _reload_task(db, task_id: str) -> ChatGenerationTask | None: + """ + rollback 后重新查询任务对象。 + + 说明： + - SQLAlchemy rollback 后，当前 ORM 对象可能过期。 + - 继续访问旧 task 有概率触发异步懒加载异常。 + """ + result = await db.execute( + select(ChatGenerationTask).where(ChatGenerationTask.id == task_id) + ) + return result.scalar_one_or_none() + + +async def _run(task_id: str): + async with async_session() as db: + result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id)) + task = result.scalar_one_or_none() + if not task or task.generation_mode != "chatapi_async": + return + + if task.status != "generating": + return + + # 关键修改 3： + # 原来只允许 result_ready 进入下载。 + # 现在允许 downloading 恢复，但只有“卡住超过 DOWNLOAD_STUCK_SECONDS”的 downloading 才继续。 + if task.pipeline_stage == "downloading": + if _is_recent_downloading(task): + # downloading 很新，说明可能有 worker 正在下载，直接跳过，避免并发重复下载。 + return + + # downloading 已经很久没更新，认为 worker 可能挂了，允许恢复下载。 + await log_task_event( + task, + event_type="DOWNLOAD_STUCK_RECOVER", + message=f"downloading 超过 {DOWNLOAD_STUCK_SECONDS} 秒，重新进入下载流程", + ) + + elif task.pipeline_stage != "result_ready": + return + + try: + old_stage = task.pipeline_stage + + # 无论从 result_ready 进入，还是从 stuck downloading 恢复，都重新标记为 downloading。 + task.pipeline_stage = "downloading" + await db.commit() + + await log_task_event( + task, + event_type="DOWNLOAD_START", + from_stage=old_stage, + to_stage="downloading", + ) + + local_url = await download_generation_result(task) + + if task.gen_type == "image": + task.image_url = local_url + else: + task.video_url = local_url + + task.status = "completed" + task.pipeline_stage = "done" + task.generated_at = datetime.now(timezone.utc) + task.retry_count = 0 + await db.commit() + + await log_task_event( + task, + event_type="DOWNLOAD_SUCCESS", + to_status="completed", + to_stage="done", + ) + + except Exception as exc: + # 关键修改 2： + # 异常后先 rollback，再重新查询 task，不继续使用 rollback 前的旧 ORM 对象。 + try: + await db.rollback() + except Exception: + pass + + task = await _reload_task(db, task_id) + if not task: + return + + task.retry_count = (task.retry_count or 0) + 1 + + if task.retry_count > 3: + task.status = "failed" + task.pipeline_stage = "download_failed" + task.error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc) + await db.commit() + + await log_task_event( + task, + event_type="DOWNLOAD_FAILED", + message=task.error_message, + ) + else: + # 下载失败但未超过重试次数，改回 result_ready，等待下一次下载。 + # 这样不会卡死在 downloading。 + task.pipeline_stage = "result_ready" + await db.commit() + + download_generation_result_task.apply_async( + args=[task.id], + countdown=30 * task.retry_count, + ) + + +if celery_app: + @celery_app.task(name="generation.download_generation_result_task", bind=True, max_retries=3, default_retry_delay=30) + def download_generation_result_task(self, task_id: str): + return run_async(_run(task_id)) +else: + class _DisabledTask: + def delay(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + + def apply_async(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + + download_generation_result_task = _DisabledTask() \ No newline at end of file diff --git a/video-gen-api/app/tasks/generation_poll_tasks.py b/video-gen-api/app/tasks/generation_poll_tasks.py new file mode 100644 index 00000000..dc973175 --- /dev/null +++ b/video-gen-api/app/tasks/generation_poll_tasks.py @@ -0,0 +1,199 @@ +from app.tasks.async_runner import run_async +import json +from datetime import datetime, timezone + +from sqlalchemy import select + +from app.config import settings +from app.models.base import async_session +from app.models.chat_generation_task import ChatGenerationTask +from app.services.error_codes import extract_error_message +from app.services.generation_log_service import log_task_event, log_provider_call +from app.services.generation_provider_service import poll_provider_task +from app.tasks.celery_app import celery_app + + +def _is_success(status: str) -> bool: + return status in ("succeeded", "success", "completed", "done") + + +def _is_failed(status: str) -> bool: + return status in ("failed", "error", "canceled", "cancelled") + + +def _engine_snapshot(task: ChatGenerationTask) -> dict: + try: + return json.loads(task.engine_snapshot_json or "{}") + except Exception: + return {} + + +async def _reload_task(db, task_id: str) -> ChatGenerationTask | None: + """ + rollback 后重新查询任务对象。 + + 说明： + - SQLAlchemy rollback 后，当前 ORM 对象可能进入过期状态。 + - 后续继续访问旧 task.retry_count / task.id 等字段，有概率触发异步懒加载异常。 + - 所以 poll/download 的异常分支统一 rollback 后重新 select。 + """ + result = await db.execute( + select(ChatGenerationTask).where(ChatGenerationTask.id == task_id) + ) + return result.scalar_one_or_none() + + +async def _run(task_id: str): + async with async_session() as db: + result = await db.execute(select(ChatGenerationTask).where(ChatGenerationTask.id == task_id)) + task = result.scalar_one_or_none() + if not task or task.generation_mode != "chatapi_async": + return + + # 只处理正在生成，且处于远程等待/轮询中的任务。 + if task.status != "generating" or task.pipeline_stage not in ("waiting_remote", "polling"): + return + + if task.deadline_at and datetime.now(timezone.utc) > task.deadline_at: + task.status = "failed" + task.pipeline_stage = "timeout" + task.error_message = "任务轮询超时" + await db.commit() + await log_task_event(task, event_type="TASK_TIMEOUT", to_status="failed", to_stage="timeout") + return + + if not (task.seedance_task_id or task.provider_task_id): + task.status = "failed" + task.pipeline_stage = "failed" + task.error_message = "缺少外部任务ID" + await db.commit() + await log_task_event(task, event_type="POLL_FAILED", message=task.error_message) + return + + # 标记本次正在轮询。 + # 注意：pending 后会再改回 waiting_remote，避免任务长期卡在 polling。 + task.pipeline_stage = "polling" + task.poll_count = (task.poll_count or 0) + 1 + task.last_poll_at = datetime.now(timezone.utc) + await db.commit() + + try: + poll_result = await poll_provider_task(db, task) + status = poll_result.get("status") + response_data = poll_result.get("response_data") + + try: + provider_response = json.loads(response_data or "{}") + except Exception: + provider_response = {"raw": response_data} + + snapshot = _engine_snapshot(task) + + await log_provider_call( + task, + provider=snapshot.get("provider") or "ark", + api_type=f"{task.gen_type}_poll", + model=snapshot.get("model_name"), + engine_id=task.engine_id, + status="success", + provider_task_id=task.seedance_task_id or task.provider_task_id, + response_data=provider_response, + ) + + if _is_success(status): + if task.gen_type == "image": + task.remote_result_url = poll_result.get("image_url") + task.image_tokens_used = poll_result.get("image_tokens", 0) or 0 + else: + task.remote_result_url = poll_result.get("video_url") + task.video_tokens_used = poll_result.get("video_tokens", 0) or 0 + + task.provider_response_json = response_data + + if not task.remote_result_url: + task.status = "failed" + task.pipeline_stage = "failed" + task.error_message = "供应商任务成功但未返回结果URL" + await db.commit() + await log_task_event(task, event_type="POLL_FAILED", message=task.error_message) + return + + task.pipeline_stage = "result_ready" + task.retry_count = 0 + await db.commit() + + await log_task_event(task, event_type="POLL_SUCCESS", to_stage="result_ready") + + from app.tasks.generation_download_tasks import download_generation_result_task + download_generation_result_task.delay(task.id) + return + + if _is_failed(status): + task.status = "failed" + task.pipeline_stage = "failed" + task.error_message = poll_result.get("error") or f"供应商任务失败: {status}" + task.provider_response_json = response_data + await db.commit() + await log_task_event(task, event_type="POLL_FAILED", message=task.error_message, detail=poll_result) + return + + # 关键修改 1： + # 供应商仍在 pending / running 时，把阶段从 polling 改回 waiting_remote。 + # 这样数据库状态表示“等待下一次轮询”，不会长期停在 polling。 + # 同时可以降低重复 Celery 消息形成多条轮询链的概率。 + task.pipeline_stage = "waiting_remote" + task.retry_count = 0 + await db.commit() + + await log_task_event(task, event_type="POLL_PENDING", message=f"status={status}") + + poll_generation_task.apply_async( + args=[task.id], + countdown=settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS, + ) + + except Exception as exc: + # 关键修改 2： + # 异常后先 rollback，再重新查询 task，不继续使用 rollback 前的旧 ORM 对象。 + try: + await db.rollback() + except Exception: + pass + + task = await _reload_task(db, task_id) + if not task: + return + + task.retry_count = (task.retry_count or 0) + 1 + + if task.retry_count > settings.CHATAPI_ASYNC_MAX_RETRIES: + task.status = "failed" + task.pipeline_stage = "failed" + task.error_message = extract_error_message(exc, "轮询") if callable(extract_error_message) else str(exc) + await db.commit() + await log_task_event(task, event_type="POLL_FAILED", message=task.error_message) + else: + # 临时轮询异常时，不让任务停在 polling。 + # 回到 waiting_remote，等待下一次重试轮询。 + task.pipeline_stage = "waiting_remote" + await db.commit() + + poll_generation_task.apply_async( + args=[task.id], + countdown=settings.CHATAPI_ASYNC_RETRY_BACKOFF_SECONDS * task.retry_count, + ) + + +if celery_app: + @celery_app.task(name="generation.poll_generation_task", bind=True, max_retries=3, default_retry_delay=30) + def poll_generation_task(self, task_id: str): + return run_async(_run(task_id)) +else: + class _DisabledTask: + def delay(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + + def apply_async(self, *args, **kwargs): + raise RuntimeError("Celery is disabled") + + poll_generation_task = _DisabledTask() \ No newline at end of file