ai-flow-python/app/service/task/impl/nano_banana_task.py

import logging
import time

from app.config import settings
from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util.http_util import logged_client
from app.util.retry_util import retry_on_failure

logger = logging.getLogger(__name__)


class NanoBananaTask(BaseTask):
    def get_type(self) -> str:
        return "nano_banana"

    def _submit(self, endpoint: str, headers: dict, inputs: dict) -> dict:
        with logged_client(timeout=120) as client:
            resp = client.post(endpoint, headers=headers, json=inputs)
            logger.info("[nano_banana] Submit response: status=%d, body_size=%d",
                        resp.status_code, len(resp.content))
            resp.raise_for_status()
            data = resp.json()
        if not data.get("candidates"):
            raise RuntimeError("No candidates in response")
        return data

    def execute(self, param: TaskParam) -> TaskResult:
        endpoint = f"{settings.gemini_base_url}/models/{settings.gemini_nano_banana_model}:generateContent"
        headers = {
            "x-goog-api-key": settings.gemini_api_key,
            "Content-Type": "application/json",
        }

        logger.info("[nano_banana] Calling Gemini API: POST %s", endpoint)
        t0 = time.time()
        try:
            data = retry_on_failure(
                lambda: self._submit(endpoint, headers, param.inputs),
                task_name="nano_banana",
            )
        except Exception as e:
            elapsed = time.time() - t0
            logger.error("[nano_banana] Request failed in %.2fs: %s", elapsed, e)
            return TaskResult.failure(f"submit execute fail: {e}")

        encoded_files: list[str] = []
        candidates = data.get("candidates", [])
        parts = candidates[0].get("content", {}).get("parts", [])
        for part in parts:
            text = part.get("text")
            if text:
                logger.info("[nano_banana] Text part: %s", text[:200])
            inline_data = part.get("inlineData")
            if inline_data and inline_data.get("data"):
                encoded_files.append(inline_data["data"])

        logger.info("[nano_banana] Completed: %d images generated", len(encoded_files))
        return TaskResult.success({"images": encoded_files})