gemini_backend.py

import os
import time
from typing import TYPE_CHECKING
from PIL import Image
from logging_helper import log as _log, log_debug as _log_debug, _log_model_response
from common import MODELS_MAP

try:
    # New google-genai client library
    from google import genai
except ImportError:  # pragma: no cover
    genai = None  # type: ignore

if TYPE_CHECKING:  # pragma: no cover
    # Kept for type-checkers; Image is also imported at runtime above
    from PIL import Image as _ImageType


def _run_gemini_vision(image: Image.Image, prompt: str, model_choice: str) -> str:
    if genai is None:
        raise RuntimeError("google-genai package is not installed. Please install it to use Gemini backend.")

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable is not set.")

    model_name = model_choice

    # Instantiate the google-genai client and call the model
    client = genai.Client(api_key=api_key)

    _log_debug(f"Using Gemini model: {model_name}")
    _log_debug(f"Input image size: {image.size}")

    start_time = time.perf_counter()

    # google-genai accepts mixed text and image content
    response = client.models.generate_content(
        model=model_name,
        contents=[prompt, image],
    )

    duration = time.perf_counter() - start_time
    print(f"Response: {response}")
    print(f"Response text: {getattr(response, 'text', 'No text attribute')}")
    
    content = response.text or ""

    usage = getattr(response, "usage_metadata", None)
    if usage is None:
        usage = getattr(response, "usage", None)

    _log_model_response(
        model_name=model_name,
        content=content,
        duration=duration,
        usage=usage,
        pricing=MODELS_MAP,
    )

    return content