生成记录接口兼容图片类型|CHAT对话生成视频/图片任务API+celery异步模块完成
This commit is contained in:
+1
-1
@@ -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
|
||||
|
||||
@@ -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/
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
@@ -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,9 +1055,13 @@ async def admin_generate_video(
|
||||
raise HTTPException(status_code=404, detail="记录不存在")
|
||||
|
||||
record, project_name = row
|
||||
if record.status not in ("prompt_optimized", "failed"):
|
||||
raise HTTPException(status_code=400, detail="当前状态不允许生成视频")
|
||||
type_str = "视频" if record.gen_type == GenerationType.video else "图片"
|
||||
|
||||
if record.status not in ("prompt_optimized", "failed"):
|
||||
raise HTTPException(status_code=400, detail=f"当前状态不允许生成{type_str}")
|
||||
|
||||
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:
|
||||
@@ -1073,7 +1086,6 @@ async def admin_generate_video(
|
||||
|
||||
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)
|
||||
@@ -1084,5 +1096,27 @@ async def admin_generate_video(
|
||||
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}
|
||||
|
||||
@@ -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"],
|
||||
|
||||
@@ -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="任务已重新投递",
|
||||
)
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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 倒序排列",
|
||||
)
|
||||
@@ -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],
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user