ai.allucanget.biz/backend/app/routers/generate.py

"""Generate router: text, image, video, and image-to-video generation."""
from fastapi import APIRouter, Depends, HTTPException, status

from backend.app.dependencies import get_current_user
from backend.app.models.ai import (
    ImageRequest,
    ImageResponse,
    ImageResult,
    TextRequest,
    TextResponse,
    VideoFromImageRequest,
    VideoRequest,
    VideoResponse,
)
from backend.app.services import openrouter

router = APIRouter(prefix="/generate", tags=["generate"])


@router.post("/text", response_model=TextResponse)
async def generate_text(
    body: TextRequest,
    _: dict = Depends(get_current_user),
) -> TextResponse:
    """Generate text from a prompt using a chat model."""
    messages = []
    if body.system_prompt:
        messages.append({"role": "system", "content": body.system_prompt})
    messages.append({"role": "user", "content": body.prompt})

    try:
        result = await openrouter.chat_completion(
            model=body.model,
            messages=messages,
            temperature=body.temperature,
            max_tokens=body.max_tokens,
        )
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")

    try:
        choice = result["choices"][0]
        return TextResponse(
            id=result["id"],
            model=result.get("model", body.model),
            content=choice["message"]["content"],
            usage=result.get("usage"),
        )
    except (KeyError, IndexError) as exc:
        raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY,
                            detail=f"Unexpected response format: {exc}")


@router.post("/image", response_model=ImageResponse)
async def generate_image(
    body: ImageRequest,
    _: dict = Depends(get_current_user),
) -> ImageResponse:
    """Generate images from a text prompt."""
    try:
        result = await openrouter.generate_image(
            model=body.model,
            prompt=body.prompt,
            n=body.n,
            size=body.size,
        )
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")

    try:
        images = [
            ImageResult(
                url=item.get("url"),
                b64_json=item.get("b64_json"),
                revised_prompt=item.get("revised_prompt"),
            )
            for item in result.get("data", [])
        ]
        return ImageResponse(
            id=result.get("id", ""),
            model=result.get("model", body.model),
            images=images,
        )
    except (KeyError, TypeError) as exc:
        raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY,
                            detail=f"Unexpected response format: {exc}")


@router.post("/video", response_model=VideoResponse)
async def generate_video(
    body: VideoRequest,
    _: dict = Depends(get_current_user),
) -> VideoResponse:
    """Generate a video from a text prompt."""
    try:
        result = await openrouter.generate_video(
            model=body.model,
            prompt=body.prompt,
            duration_seconds=body.duration_seconds,
            aspect_ratio=body.aspect_ratio,
        )
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")

    urls = result.get("unsigned_urls") or result.get("video_urls")
    return VideoResponse(
        id=result.get("id", ""),
        model=body.model,
        status=result.get("status", "queued"),
        polling_url=result.get("polling_url"),
        video_urls=urls,
        video_url=(urls or [None])[0],
        error=result.get("error"),
        metadata=result.get("metadata"),
    )


@router.post("/video/from-image", response_model=VideoResponse)
async def generate_video_from_image(
    body: VideoFromImageRequest,
    _: dict = Depends(get_current_user),
) -> VideoResponse:
    """Generate a video from an image and a text prompt."""
    try:
        result = await openrouter.generate_video_from_image(
            model=body.model,
            image_url=body.image_url,
            prompt=body.prompt,
            duration_seconds=body.duration_seconds,
            aspect_ratio=body.aspect_ratio,
        )
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")

    urls = result.get("unsigned_urls") or result.get("video_urls")
    return VideoResponse(
        id=result.get("id", ""),
        model=body.model,
        status=result.get("status", "queued"),
        polling_url=result.get("polling_url"),
        video_urls=urls,
        video_url=(urls or [None])[0],
        error=result.get("error"),
        metadata=result.get("metadata"),
    )


@router.get("/video/status", response_model=VideoResponse)
async def poll_video_status(
    polling_url: str,
    _: dict = Depends(get_current_user),
) -> VideoResponse:
    """Poll the status of a video generation job via its polling_url."""
    try:
        result = await openrouter.poll_video_status(polling_url)
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY, detail=f"OpenRouter error: {exc}")

    urls = result.get("unsigned_urls") or result.get("video_urls")
    return VideoResponse(
        id=result.get("id", ""),
        model=result.get("model", ""),
        status=result.get("status", "processing"),
        polling_url=result.get("polling_url"),
        video_urls=urls,
        video_url=(urls or [None])[0],
        error=result.get("error"),
        metadata=result.get("metadata"),
    )