Initial commit: Freepik REST API
FastAPI async wrapper for Freepik cloud AI API supporting image generation (Mystic, Flux Dev/Pro, SeedReam), video generation (Kling, MiniMax, Seedance), image editing (upscale, relight, style transfer, expand, inpaint), and utilities (background removal, classifier, audio isolation). Includes async task tracking with polling, Docker containerization, and Gitea CI/CD workflow. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
179
app/services/task_tracker.py
Normal file
179
app/services/task_tracker.py
Normal file
@@ -0,0 +1,179 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from app.config import settings
|
||||
from app.schemas.common import TaskStatus
|
||||
from app.services import freepik_client
|
||||
from app.services.file_manager import download_result
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_tasks: dict[str, dict] = {}
|
||||
_poll_tasks: dict[str, asyncio.Task] = {}
|
||||
|
||||
|
||||
def submit(freepik_task_id: str, metadata: Optional[dict] = None) -> str:
|
||||
"""Register a Freepik task and start background polling."""
|
||||
internal_id = str(uuid.uuid4())
|
||||
now = datetime.now(timezone.utc)
|
||||
_tasks[internal_id] = {
|
||||
'task_id': internal_id,
|
||||
'freepik_task_id': freepik_task_id,
|
||||
'status': TaskStatus.pending,
|
||||
'created_at': now,
|
||||
'updated_at': now,
|
||||
'progress': None,
|
||||
'result_url': None,
|
||||
'local_path': None,
|
||||
'error': None,
|
||||
'metadata': metadata or {},
|
||||
}
|
||||
_poll_tasks[internal_id] = asyncio.create_task(_poll_loop(internal_id))
|
||||
return internal_id
|
||||
|
||||
|
||||
def get_task(task_id: str) -> Optional[dict]:
|
||||
return _tasks.get(task_id)
|
||||
|
||||
|
||||
def list_tasks(
|
||||
status: Optional[TaskStatus] = None,
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
) -> tuple[list[dict], int]:
|
||||
tasks = list(_tasks.values())
|
||||
if status:
|
||||
tasks = [t for t in tasks if t['status'] == status]
|
||||
tasks.sort(key=lambda t: t['created_at'], reverse=True)
|
||||
total = len(tasks)
|
||||
return tasks[offset:offset + limit], total
|
||||
|
||||
|
||||
def delete_task(task_id: str) -> bool:
|
||||
if task_id not in _tasks:
|
||||
return False
|
||||
# Cancel polling if active
|
||||
poll = _poll_tasks.pop(task_id, None)
|
||||
if poll and not poll.done():
|
||||
poll.cancel()
|
||||
_tasks.pop(task_id, None)
|
||||
return True
|
||||
|
||||
|
||||
def active_count() -> int:
|
||||
return sum(
|
||||
1 for t in _tasks.values()
|
||||
if t['status'] in (TaskStatus.pending, TaskStatus.processing)
|
||||
)
|
||||
|
||||
|
||||
async def _poll_loop(internal_id: str):
|
||||
"""Poll Freepik API until the task completes or times out."""
|
||||
task = _tasks.get(internal_id)
|
||||
if not task:
|
||||
return
|
||||
|
||||
freepik_id = task['freepik_task_id']
|
||||
elapsed = 0
|
||||
|
||||
try:
|
||||
while elapsed < settings.task_poll_timeout_seconds:
|
||||
await asyncio.sleep(settings.task_poll_interval_seconds)
|
||||
elapsed += settings.task_poll_interval_seconds
|
||||
|
||||
try:
|
||||
result = await freepik_client.get_task_status(freepik_id)
|
||||
except Exception as exc:
|
||||
logger.warning(f'Poll error for {internal_id}: {exc}')
|
||||
continue
|
||||
|
||||
data = result.get('data', result)
|
||||
fp_status = data.get('status', '')
|
||||
|
||||
task['updated_at'] = datetime.now(timezone.utc)
|
||||
|
||||
if fp_status in ('IN_PROGRESS', 'PROCESSING', 'processing'):
|
||||
task['status'] = TaskStatus.processing
|
||||
task['progress'] = data.get('progress')
|
||||
continue
|
||||
|
||||
if fp_status in ('COMPLETED', 'completed', 'done'):
|
||||
task['status'] = TaskStatus.completed
|
||||
task['progress'] = 1.0
|
||||
# Extract result URL from various response shapes
|
||||
result_url = (
|
||||
data.get('result_url')
|
||||
or data.get('output', {}).get('url')
|
||||
or _extract_first_url(data)
|
||||
)
|
||||
task['result_url'] = result_url
|
||||
if result_url:
|
||||
try:
|
||||
task['local_path'] = await download_result(
|
||||
internal_id, result_url
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f'Download failed for {internal_id}: {exc}')
|
||||
logger.info(f'Task {internal_id} completed')
|
||||
return
|
||||
|
||||
if fp_status in ('FAILED', 'failed', 'error'):
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = data.get('error', data.get('message', 'Unknown error'))
|
||||
logger.warning(f'Task {internal_id} failed: {task["error"]}')
|
||||
return
|
||||
|
||||
# Timeout
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = f'Polling timed out after {settings.task_poll_timeout_seconds}s'
|
||||
logger.warning(f'Task {internal_id} timed out')
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f'Polling cancelled for {internal_id}')
|
||||
except Exception as exc:
|
||||
task['status'] = TaskStatus.failed
|
||||
task['error'] = str(exc)
|
||||
logger.error(f'Unexpected error polling {internal_id}: {exc}')
|
||||
finally:
|
||||
_poll_tasks.pop(internal_id, None)
|
||||
|
||||
|
||||
def _extract_first_url(data: dict) -> Optional[str]:
|
||||
"""Try to extract the first URL from common Freepik response shapes."""
|
||||
# Some endpoints return {"data": {"images": [{"url": "..."}]}}
|
||||
for key in ('images', 'videos', 'results', 'outputs'):
|
||||
items = data.get(key, [])
|
||||
if isinstance(items, list) and items:
|
||||
first = items[0]
|
||||
if isinstance(first, dict) and 'url' in first:
|
||||
return first['url']
|
||||
if isinstance(first, str) and first.startswith('http'):
|
||||
return first
|
||||
# Direct URL field
|
||||
if 'url' in data and isinstance(data['url'], str):
|
||||
return data['url']
|
||||
return None
|
||||
|
||||
|
||||
def handle_webhook_completion(freepik_task_id: str, result_data: dict):
|
||||
"""Called when a webhook notification arrives for a completed task."""
|
||||
for task in _tasks.values():
|
||||
if task['freepik_task_id'] == freepik_task_id:
|
||||
task['status'] = TaskStatus.completed
|
||||
task['progress'] = 1.0
|
||||
task['updated_at'] = datetime.now(timezone.utc)
|
||||
result_url = (
|
||||
result_data.get('result_url')
|
||||
or result_data.get('output', {}).get('url')
|
||||
or _extract_first_url(result_data)
|
||||
)
|
||||
task['result_url'] = result_url
|
||||
# Cancel polling since webhook already notified us
|
||||
poll = _poll_tasks.pop(task['task_id'], None)
|
||||
if poll and not poll.done():
|
||||
poll.cancel()
|
||||
logger.info(f'Task {task["task_id"]} completed via webhook')
|
||||
break
|
||||
Reference in New Issue
Block a user