Tool_Agent/tools/local/kuaishou_kling/main.py

"""
快手可灵AI工具 FastAPI 服务
提供统一的HTTP接口用于AI视频生成、AI图片生成、AI对口型等功能
"""
import asyncio
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any, List
from enum import Enum
import uvicorn

from kling_client import KlingClient, BizType, TaskStatus


app = FastAPI(
    title="快手可灵AI工具",
    description="支持AI视频生成、AI图片生成、AI对口型等功能",
    version="1.0.0"
)


class GenerateRequest(BaseModel):
    """生成请求模型"""
    biz_type: str = Field(..., description="业务类型: aiImage, aiVideo, aiLipSync")
    action: Optional[str] = Field(None, description="动作类型")
    prompt: Optional[str] = Field(None, description="生成内容的提示词")
    negative_prompt: Optional[str] = Field(None, description="不希望呈现的内容")
    cfg: str = Field("50", description="创意想象力与创意相关性比例")
    mode: Optional[str] = Field(None, description="生成模式: text2video, audio2video")
    image_url: Optional[str] = Field(None, description="参考图片地址")
    aspect_ratio: str = Field("16:9", description="长宽比: 9:16, 16:9, 1:1")
    task_id: Optional[str] = Field(None, description="查询任务状态时使用")
    cookie: Optional[str] = Field(None, description="认证Cookie")
    version: Optional[str] = Field(None, description="模型版本")
    image_count: int = Field(4, description="生成图片数量(1-4)")
    add_audio: bool = Field(False, description="是否自动添加音频")
    start_frame_image: Optional[str] = Field(None, description="首帧图片URL")
    end_frame_image: Optional[str] = Field(None, description="尾帧图片URL")
    video_id: Optional[str] = Field(None, description="视频ID(对口型用)")
    video_url: Optional[str] = Field(None, description="视频URL(对口型用)")
    text: Optional[str] = Field(None, description="对口型文本内容")
    voice_id: Optional[str] = Field(None, description="音色ID")
    voice_language: str = Field("zh", description="音色语种: zh, en")
    voice_speed: float = Field(1.0, description="语速")
    audio_type: Optional[str] = Field(None, description="音频类型: file, url")
    audio_file: Optional[str] = Field(None, description="音频文件路径")
    audio_url: Optional[str] = Field(None, description="音频URL")


class GenerateResponse(BaseModel):
    """生成响应模型"""
    task_id: Optional[str] = Field(None, description="任务ID")
    status: Optional[str] = Field(None, description="任务状态: process, finished, failed")
    result: Optional[Dict[str, Any]] = Field(None, description="生成结果")
    error: Optional[str] = Field(None, description="错误信息")


class StatusResponse(BaseModel):
    """状态查询响应模型"""
    task_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态: process, finished, failed")
    result: Optional[Dict[str, Any]] = Field(None, description="生成结果")
    error: Optional[str] = Field(None, description="错误信息")


@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "快手可灵AI工具",
        "version": "1.0.0",
        "endpoints": {
            "generate": "POST /generate - 创建生成任务",
            "status": "GET /status/{task_id} - 查询任务状态"
        }
    }


@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """
    创建生成任务
    
    支持三种业务类型：
    - aiImage: AI图片生成
    - aiVideo: AI视频生成
    - aiLipSync: AI对口型
    """
    try:
        client = KlingClient(cookie=request.cookie)
        
        # 根据业务类型调用不同的API
        if request.biz_type == "aiImage":
            if not request.prompt:
                raise HTTPException(status_code=400, detail="prompt is required for aiImage")
            
            result = await client.create_image_task(
                prompt=request.prompt,
                negative_prompt=request.negative_prompt,
                cfg=request.cfg,
                aspect_ratio=request.aspect_ratio,
                image_count=request.image_count,
                version=request.version
            )
            
        elif request.biz_type == "aiVideo":
            if not request.prompt:
                raise HTTPException(status_code=400, detail="prompt is required for aiVideo")
            
            result = await client.create_video_task(
                prompt=request.prompt,
                mode=request.mode or "text2video",
                image_url=request.image_url,
                aspect_ratio=request.aspect_ratio,
                add_audio=request.add_audio,
                start_frame_image=request.start_frame_image,
                end_frame_image=request.end_frame_image,
                version=request.version
            )
            
        elif request.biz_type == "aiLipSync":
            result = await client.create_lipsync_task(
                video_id=request.video_id,
                video_url=request.video_url,
                mode=request.mode or "text2video",
                text=request.text,
                voice_id=request.voice_id,
                voice_language=request.voice_language,
                voice_speed=request.voice_speed,
                audio_type=request.audio_type,
                audio_file=request.audio_file,
                audio_url=request.audio_url
            )
            
        else:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid biz_type: {request.biz_type}. Must be one of: aiImage, aiVideo, aiLipSync"
            )
        
        # 解析响应
        task_id = result.get("task_id") or result.get("data", {}).get("task_id")
        status = result.get("status", "process")
        
        return GenerateResponse(
            task_id=task_id,
            status=status,
            result=result.get("result") or result.get("data"),
            error=result.get("error")
        )
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")


@app.get("/status/{task_id}", response_model=StatusResponse)
async def get_status(
    task_id: str,
    biz_type: str = "aiImage",
    cookie: Optional[str] = None
):
    """
    查询任务状态
    
    Args:
        task_id: 任务ID
        biz_type: 业务类型 (aiImage, aiVideo, aiLipSync)
        cookie: 认证Cookie
    """
    try:
        # 验证biz_type
        if biz_type not in ["aiImage", "aiVideo", "aiLipSync"]:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid biz_type: {biz_type}. Must be one of: aiImage, aiVideo, aiLipSync"
            )
        
        client = KlingClient(cookie=cookie)
        biz_type_enum = BizType(biz_type)
        
        result = await client.query_task_status(task_id, biz_type_enum)
        
        return StatusResponse(
            task_id=task_id,
            status=result.get("status", "process"),
            result=result.get("result") or result.get("data"),
            error=result.get("error")
        )
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}")


@app.get("/health")
async def health():
    """健康检查"""
    return {"status": "healthy"}


if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
    uvicorn.run(app, host="0.0.0.0", port=port)