publisher/ai/gemini_client.py

import os
import time
import random
import logging

logger = logging.getLogger("gemini-client")

# =========================
# MODEL POOL (Gemini 3 Era)
# =========================

PRIMARY_MODELS = [
    "gemini-3.1-pro",
    "gemini-3.1-flash",
]

FALLBACK_MODELS = [
    "gemini-2.5-flash",
    "gemini-3.1-flash-lite",
]

# =========================
# SIMPLE QUOTA TRACKER
# =========================
_model_fail_count = {
    "gemini-3.1-pro": 0,
    "gemini-3.1-flash": 0,
}


MAX_FAILS = 3


# =========================
# CORE MODEL RESOLVER
# =========================

def _pick_model():
    """
    Select best available model with fallback logic.
    """
    for m in PRIMARY_MODELS:
        if _model_fail_count.get(m, 0) < MAX_FAILS:
            return m

    return random.choice(FALLBACK_MODELS)


# =========================
# MAIN CLIENT INTERFACE
# =========================

def get_model():
    """
    Public entrypoint used by publisher_ai.
    Returns an initialized Gemini model.
    """
    try:
        import google.generativeai as genai
    except ImportError as exc:
        raise RuntimeError("google-generativeai is required for Gemini features") from exc

    api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
    if api_key:
        genai.configure(api_key=api_key)

    model_name = _pick_model()
    logger.info(f"[Gemini] Selected model: {model_name}")
    return genai.GenerativeModel(model_name)


def safe_generate(prompt: str, client_callable):
    """
    Wrapper for Gemini calls with automatic fallback.
    """

    last_error = None

    for _ in range(3):
        model = _pick_model()

        try:
            result = client_callable(model, prompt)
            return result

        except Exception as e:
            last_error = e
            _model_fail_count[model] = _model_fail_count.get(model, 0) + 1
            logger.warning(f"[Gemini FAIL] {model}: {str(e)}")
            time.sleep(0.5)

    raise RuntimeError(f"All Gemini models failed: {last_error}")