ai-swap/backend/api/ai_swap_bg_api.py

"""
AI换脸换装API接口
提供RESTful API接口，支持换脸换装、历史记录查询等功能
"""

from fastapi import FastAPI, HTTPException, Depends, UploadFile, File, Form, Body, APIRouter
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import os
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any, Union
import uvicorn
from datetime import datetime, timedelta
from pydantic import BaseModel
from sqlalchemy import func, and_

from backend.services.ai_swap_bg_service import AISwapBgService, process_swap_bg_with_record
from backend.modules.database.operations import DatabaseOperations
from backend.utils.logger_config import setup_logger
from backend.api.auth_api import router as auth_router
from backend.api.user_material_api import router as user_material_router
from backend.api.user_template_api import router as user_template_router
from backend.modules.database.models import ProcessRecord
from backend.api.auto_post_api import router as auto_post_router


logger = setup_logger(__name__)

router = APIRouter()

# 创建服务实例
ai_swap_bg_service = AISwapBgService()
db_ops = DatabaseOperations()

# 启动任务队列服务
from backend.services.task_queue_service import get_task_queue_service
task_queue_service = get_task_queue_service()
task_queue_service.start()

# ==================== 数据模型 ====================

class SwapBgRequest(BaseModel):
    user_id: int = Field(..., description="用户ID")
    image_id: int = Field(..., description="图片ID")
    prompt: str = Field(..., description="提示词", max_length=500)
    quantity: Optional[int] = Field(1, ge=1, le=10, description="每组生成数量")
    
    
class SwapBgResponse(BaseModel):
    success: bool
    task_id: Optional[str] = None
    process_record_id: Optional[int] = None
    result_image_id: Optional[int] = None
    copywriter_text: Optional[str] = None
    history_prompt: Optional[str] = None
    error: Optional[str] = None
    error_type: Optional[str] = None

# ==================== 中间件和依赖 ====================

async def verify_user(user_id: int):
    """验证用户是否存在且有效"""
    user = db_ops.get_user_by_id(user_id)
    if not user:
        raise HTTPException(status_code=404, detail=f"用户ID {user_id} 不存在")
    if not user.get("is_active", False):
        raise HTTPException(status_code=403, detail="用户账户已被禁用")
    return user

# ==================== API路由 ====================
@router.post("/api/v1/swap-bg", response_model=SwapBgResponse, tags=["AI换背景"])
async def process_swap_bg(request: SwapBgRequest):
    """
    执行AI换背景（异步任务）
    
    - **user_id**: 用户ID
    - **image_id**: 图片ID
    - **prompt**: 提示词
    - **quantity**: 每组生成数量（可选）
    """
    try:
        # 验证用户
        await verify_user(request.user_id)
        
        # 提交任务到队列（非阻塞）
        task_id = ai_swap_bg_service.submit_swap_bg_task(
            user_id=request.user_id,
            image_id=request.image_id,
            prompt=request.prompt,
            quantity=request.quantity or 1
        )
        
        logger.info(f"用户 {request.user_id} 换背景任务已提交，任务ID: {task_id}")
        return SwapBgResponse(
            success=True,
            task_id=task_id,
            process_record_id=None,
            result_image_id=None,
            copywriter_text=None,
            history_prompt=None
        )
            
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"API处理异常: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {str(e)}")