freepik-api/app/services/freepik_client.py

import asyncio
import logging
from typing import Any, Optional

import httpx

from app.config import settings

logger = logging.getLogger(__name__)

_client: Optional[httpx.AsyncClient] = None


def get_client() -> httpx.AsyncClient:
    if _client is None:
        raise RuntimeError('Freepik client not initialized')
    return _client


async def create_client() -> httpx.AsyncClient:
    global _client
    _client = httpx.AsyncClient(
        base_url=settings.freepik_base_url,
        headers={
            'x-freepik-api-key': settings.freepik_api_key,
            'Accept': 'application/json',
        },
        timeout=httpx.Timeout(60.0, connect=10.0),
    )
    return _client


async def close_client():
    global _client
    if _client:
        await _client.aclose()
        _client = None


async def _request(
    method: str,
    path: str,
    *,
    json: Optional[dict] = None,
    data: Optional[dict] = None,
    files: Optional[dict] = None,
    params: Optional[dict] = None,
    max_retries: int = 3,
) -> dict[str, Any]:
    """Make an authenticated request to Freepik API with retry on 429."""
    client = get_client()
    for attempt in range(max_retries):
        kwargs: dict[str, Any] = {}
        if json is not None:
            kwargs['json'] = json
        if data is not None:
            kwargs['data'] = data
        if files is not None:
            kwargs['files'] = files
        if params is not None:
            kwargs['params'] = params

        response = await client.request(method, path, **kwargs)

        if response.status_code == 429:
            retry_after = int(response.headers.get('Retry-After', 2 ** attempt))
            logger.warning(f'Rate limited, retrying in {retry_after}s (attempt {attempt + 1})')
            await asyncio.sleep(retry_after)
            continue

        response.raise_for_status()
        if response.headers.get('content-type', '').startswith('application/json'):
            return response.json()
        return {'raw': response.content}

    raise httpx.HTTPStatusError(
        'Rate limit exceeded after retries',
        request=response.request,
        response=response,
    )


async def _request_raw(
    method: str,
    path: str,
    *,
    json: Optional[dict] = None,
    data: Optional[dict] = None,
    files: Optional[dict] = None,
    max_retries: int = 3,
) -> bytes:
    """Make a request and return raw bytes (for binary responses)."""
    client = get_client()
    for attempt in range(max_retries):
        kwargs: dict[str, Any] = {}
        if json is not None:
            kwargs['json'] = json
        if data is not None:
            kwargs['data'] = data
        if files is not None:
            kwargs['files'] = files

        response = await client.request(method, path, **kwargs)

        if response.status_code == 429:
            retry_after = int(response.headers.get('Retry-After', 2 ** attempt))
            logger.warning(f'Rate limited, retrying in {retry_after}s (attempt {attempt + 1})')
            await asyncio.sleep(retry_after)
            continue

        response.raise_for_status()
        return response.content

    raise httpx.HTTPStatusError(
        'Rate limit exceeded after retries',
        request=response.request,
        response=response,
    )


def _strip_none(d: dict) -> dict:
    """Remove None values from a dict for clean API payloads."""
    return {k: v for k, v in d.items() if v is not None}


# ---------------------------------------------------------------------------
# Image generation
# ---------------------------------------------------------------------------

async def generate_mystic(
    prompt: str,
    negative_prompt: Optional[str] = None,
    resolution: Optional[str] = None,
    styling: Optional[dict] = None,
    seed: Optional[int] = None,
    num_images: Optional[int] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'negative_prompt': negative_prompt,
        'resolution': resolution,
        'styling': styling,
        'seed': seed,
        'num_images': num_images,
    })
    return await _request('POST', '/v1/ai/mystic', json=payload)


async def generate_flux_dev(
    prompt: str,
    image: Optional[str] = None,
    guidance_scale: Optional[float] = None,
    num_images: Optional[int] = None,
    seed: Optional[int] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'image': image,
        'guidance_scale': guidance_scale,
        'num_images': num_images,
        'seed': seed,
    })
    return await _request('POST', '/v1/ai/text-to-image/flux-dev', json=payload)


async def generate_flux_pro(
    prompt: str,
    image: Optional[str] = None,
    guidance_scale: Optional[float] = None,
    seed: Optional[int] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'image': image,
        'guidance_scale': guidance_scale,
        'seed': seed,
    })
    return await _request('POST', '/v1/ai/text-to-image/flux-pro-1.1', json=payload)


async def generate_seedream(
    prompt: str,
    image: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
    num_images: Optional[int] = None,
    seed: Optional[int] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'image': image,
        'aspect_ratio': aspect_ratio,
        'num_images': num_images,
        'seed': seed,
    })
    return await _request('POST', '/v1/ai/text-to-image/seedream', json=payload)


# ---------------------------------------------------------------------------
# Video generation
# ---------------------------------------------------------------------------

async def generate_video_kling(
    image: str,
    prompt: Optional[str] = None,
    duration: Optional[str] = None,
    aspect_ratio: Optional[str] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'prompt': prompt,
        'duration': duration,
        'aspect_ratio': aspect_ratio,
    })
    return await _request('POST', '/v1/ai/image-to-video/kling', json=payload)


async def generate_video_minimax(
    prompt: str,
    first_frame_image: Optional[str] = None,
    subject_reference: Optional[str] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'first_frame_image': first_frame_image,
        'subject_reference': subject_reference,
    })
    return await _request('POST', '/v1/ai/image-to-video/minimax', json=payload)


async def generate_video_seedance(
    prompt: str,
    image: Optional[str] = None,
    duration: Optional[str] = None,
    resolution: Optional[str] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'image': image,
        'duration': duration,
        'resolution': resolution,
    })
    return await _request('POST', '/v1/ai/image-to-video/seedance', json=payload)


# ---------------------------------------------------------------------------
# Image editing
# ---------------------------------------------------------------------------

async def upscale_creative(
    image: str,
    prompt: Optional[str] = None,
    scale: Optional[int] = None,
    creativity: Optional[float] = None,
    resemblance: Optional[float] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'prompt': prompt,
        'scale': scale,
        'creativity': creativity,
        'resemblance': resemblance,
    })
    return await _request('POST', '/v1/ai/upscale/creative', json=payload)


async def upscale_precision(
    image: str,
    scale: Optional[int] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'scale': scale,
    })
    return await _request('POST', '/v1/ai/upscale/precision', json=payload)


async def relight_image(
    image: str,
    prompt: Optional[str] = None,
    light_source: Optional[str] = None,
    intensity: Optional[float] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'prompt': prompt,
        'light_source': light_source,
        'intensity': intensity,
    })
    return await _request('POST', '/v1/ai/relight', json=payload)


async def style_transfer(
    image: str,
    style_reference: str,
    strength: Optional[float] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'style_reference': style_reference,
        'strength': strength,
    })
    return await _request('POST', '/v1/ai/style-transfer', json=payload)


async def expand_image(
    image: str,
    prompt: Optional[str] = None,
    direction: Optional[str] = None,
) -> dict:
    payload = _strip_none({
        'image': image,
        'prompt': prompt,
        'direction': direction,
    })
    return await _request('POST', '/v1/ai/expand', json=payload)


async def inpaint(
    image: str,
    mask: str,
    prompt: str,
) -> dict:
    payload = {
        'image': image,
        'mask': mask,
        'prompt': prompt,
    }
    return await _request('POST', '/v1/ai/inpaint', json=payload)


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

async def remove_background(image: str) -> bytes:
    return await _request_raw('POST', '/v1/ai/remove-background', json={'image': image})


async def classify_image(image: str) -> dict:
    return await _request('POST', '/v1/ai/classifier', json={'image': image})


async def isolate_audio(audio: str) -> dict:
    return await _request('POST', '/v1/ai/audio-isolate', json={'audio': audio})


async def generate_icon(
    prompt: str,
    color: Optional[str] = None,
    shape: Optional[str] = None,
    style: Optional[str] = None,
) -> dict:
    payload = _strip_none({
        'prompt': prompt,
        'color': color,
        'shape': shape,
        'style': style,
    })
    return await _request('POST', '/v1/ai/icon', json=payload)


# ---------------------------------------------------------------------------
# Task management
# ---------------------------------------------------------------------------

async def get_task_status(task_id: str) -> dict:
    return await _request('GET', f'/v1/ai/tasks/{task_id}')


async def check_api_key() -> bool:
    """Verify the API key is valid by making a lightweight request."""
    try:
        client = get_client()
        response = await client.get('/v1/ai/tasks', params={'limit': 1})
        return response.status_code != 401
    except Exception:
        return False