Tool_Agent/tools/local/veo_video_gen/main.py

import os
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from veo_client import VEOClient

# 加载 .env 文件
load_dotenv()


app = FastAPI(title="VEO Video Generation API")


class VideoRequest(BaseModel):
    prompt: str = Field(..., description="视频生成的文本描述")
    model: str = Field(default="veo-3.1-fast", description="模型名称，如 veo-3.1-fast / veo-3.1 / veo-3.1-landscape-fast")
    timeout: int = Field(default=600, description="等待超时时间（秒）")


class VideoResponse(BaseModel):
    video_id: str
    url: str
    resolution: str
    model: str
    status: str


@app.post("/generate", response_model=VideoResponse)
async def generate_video(request: VideoRequest):
    api_key = os.getenv("APIYI_KEY")
    if not api_key:
        raise HTTPException(status_code=500, detail="APIYI_KEY 环境变量未设置")

    try:
        client = VEOClient(api_key)
        video_id = client.create_video(request.prompt, request.model)
        result = client.wait_for_completion(video_id, timeout=request.timeout)

        return VideoResponse(
            video_id=result["id"],
            url=result["video_url"],
            resolution=result.get("size", "unknown"),
            model=result.get("model", request.model),
            status="completed"
        )
    except TimeoutError:
        raise HTTPException(status_code=408, detail="视频生成超时")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
async def health():
    return {"status": "ok"}


if __name__ == "__main__":
    import uvicorn
    import sys

    # 从命令行参数获取端口，默认 5000
    port = 5000
    if "--port" in sys.argv:
        try:
            port_idx = sys.argv.index("--port")
            port = int(sys.argv[port_idx + 1])
        except (IndexError, ValueError):
            pass

    uvicorn.run(app, host="0.0.0.0", port=port)