freepik-api/app/routers/utilities.py

import logging
from datetime import datetime, timezone

from fastapi import APIRouter, HTTPException

from app.schemas.common import TaskResponse, TaskStatus
from app.schemas.utilities import (
    ClassificationResponse,
    IconRequest,
    RemoveBackgroundRequest,
    RemoveBackgroundResponse,
)
from app.services import freepik_client, task_tracker
from app.services.freepik_client import _extract_task_id

logger = logging.getLogger(__name__)

router = APIRouter(prefix='/api/v1/util', tags=['utilities'])


@router.post('/remove-background', response_model=RemoveBackgroundResponse)
async def remove_background(request: RemoveBackgroundRequest):
    """Remove background from an image. Takes a URL, returns result URLs."""
    result = await freepik_client.remove_background(request.image_url)
    return RemoveBackgroundResponse(**result)


@router.post('/classify', response_model=ClassificationResponse)
async def classify_image(request: dict):
    """Classify whether an image is AI-generated."""
    image = request.get('image')
    if not image:
        raise HTTPException(status_code=400, detail='image field is required')

    result = await freepik_client.classify_image(image)
    return result


@router.post('/audio-isolate', response_model=TaskResponse)
async def audio_isolate(request: dict):
    """Isolate audio tracks from an audio file."""
    audio = request.get('audio')
    if not audio:
        raise HTTPException(status_code=400, detail='audio field is required')

    result = await freepik_client.isolate_audio(audio)
    freepik_task_id = _extract_task_id(result.data)
    if not freepik_task_id:
        raise HTTPException(status_code=502, detail='No task_id in Freepik response')

    internal_id = task_tracker.submit(freepik_task_id, result.status_path, {'operation': 'audio-isolate'})
    return TaskResponse(
        task_id=internal_id,
        status=TaskStatus.pending,
        created_at=datetime.now(timezone.utc),
    )


# Icon generation lives under /api/v1/generate/ but is simple enough to keep here
icon_router = APIRouter(prefix='/api/v1/generate', tags=['utilities'])


@icon_router.post('/icon', response_model=TaskResponse)
async def generate_icon(request: IconRequest):
    result = await freepik_client.generate_icon(
        prompt=request.prompt,
        color=request.color,
        shape=request.shape,
        style=request.style,
    )
    freepik_task_id = _extract_task_id(result.data)
    if not freepik_task_id:
        raise HTTPException(status_code=502, detail='No task_id in Freepik response')

    internal_id = task_tracker.submit(freepik_task_id, result.status_path, {'operation': 'icon'})
    return TaskResponse(
        task_id=internal_id,
        status=TaskStatus.pending,
        created_at=datetime.now(timezone.utc),
    )