gemini_proxy/models.py

"""Model definitions and mapping from Gemini frontend JS source.

Model hex IDs and capacity_tail values from HanaokaYuzu/Gemini-API:
  - x-goog-ext-525001261-jspb header format:
    [1,null,null,null,"<hex_model_id>",null,null,0,[4],null,null,<capacity_tail>]
  - capacity_tail: 1=Basic/Free, 2=Advanced/Pro, 4=Plus
"""

# MODE_CATEGORY enum from 028-6eb337387583.js:
#   1=FAST, 2=THINKING, 3=PRO, 4=AUTO, 5=FAST_DYNAMIC_THINKING, 6=FLASH_LITE

# Gemini model hex IDs (from Gemini web frontend JS)
MODEL_HEX_IDS = {
    # Basic/Free tier models
    "gemini-3-pro": "9d8ca3786ebdfbea",
    "gemini-3-flash": "fbb127bbb056c959",
    "gemini-3-flash-thinking": "5bf011840784117a",
    # Plus/Advanced tier models (same hex IDs, different capacity_tail)
    "gemini-3-pro-plus": "e6fa609c3fa255c0",
    "gemini-3-flash-plus": "56fdd199312815e2",
    "gemini-3-flash-thinking-plus": "e051ce1aa80aa576",
}

# capacity_tail per account tier
CAPACITY_TAILS = {
    "free": 1,
    "pro": 2,     # Advanced tier in Gemini
    "plus": 4,
}

MODELS = {
    "gemini-3.5-flash": {
        "mode": 1, "think": 4,
        "hex_id": "fbb127bbb056c959",
        "desc": "Fast general-purpose model",
    },
    "gemini-3.5-flash-thinking": {
        "mode": 2, "think": 0,
        "hex_id": "5bf011840784117a",
        "desc": "Deep thinking mode, longest output (~20k chars)",
    },
    "gemini-3.1-pro": {
        "mode": 3, "think": 4,
        "hex_id": "9d8ca3786ebdfbea",
        "desc": "Pro model (requires cookie for real routing)",
    },
    "gemini-3.1-pro-enhanced": {
        "mode": 3, "think": 4, "extra": {31: 2, 80: 3},
        "hex_id": "9d8ca3786ebdfbea",
        "desc": "Pro with enhanced output (experimental)",
    },
    "gemini-auto": {
        "mode": 4, "think": 4,
        "hex_id": None,
        "desc": "Auto model selection",
    },
    "gemini-3.5-flash-thinking-lite": {
        "mode": 5, "think": 0,
        "hex_id": None,
        "desc": "Dynamic thinking with adaptive depth",
    },
    "gemini-flash-lite": {
        "mode": 6, "think": 4,
        "hex_id": None,
        "desc": "Lightweight fast model",
    },
    # Image generation models — use same mode/think/hex_id as base model,
    # Gemini detects image intent from prompt automatically.
    # The "-image" suffix is our routing flag, not a Gemini model name.
    "gemini-3.5-flash-image": {
        "mode": 1, "think": 4,
        "hex_id": "fbb127bbb056c959",
        "desc": "Flash model with image generation",
        "image_gen": True,
    },
    "gemini-3.1-pro-image": {
        "mode": 3, "think": 4,
        "hex_id": "9d8ca3786ebdfbea",
        "desc": "Pro model with image generation",
        "image_gen": True,
    },
}


def resolve_model(model_name: str, default: str = "gemini-3.5-flash"):
    """Resolve model name to (name, mode_id, think_mode, error, extra_fields, image_gen, hex_id).

    Returns: (model_name, mode_id, think_mode, error, extra_fields, image_gen, hex_id)
    hex_id is the model hex string for the x-goog-ext-525001261-jspb header.
    """
    think_override = None
    image_gen = False

    # Strip -image suffix for resolution, but keep the flag
    if model_name.endswith("-image"):
        image_gen = True
        model_name = model_name[:-6]  # Remove "-image" suffix

    if "@think=" in model_name:
        model_name, think_str = model_name.rsplit("@think=", 1)
        try:
            think_override = int(think_str)
        except ValueError:
            return None, None, None, f"Invalid think level: {think_str}", None, False, None

    # Check if the full name (including -image variant) is in MODELS
    full_name = model_name + ("-image" if image_gen else "")
    cfg = MODELS.get(full_name)
    if cfg:
        # Use the explicit model config (may have image_gen=True set)
        image_gen = cfg.get("image_gen", image_gen)
        model_name = full_name
    else:
        cfg = MODELS.get(model_name)
        if not cfg:
            from .gemini import log
            log(f"Unknown model '{model_name}', falling back to '{default}'")
            model_name = default
            cfg = MODELS[default]

    mode_id = cfg["mode"]
    think_mode = think_override if think_override is not None else cfg["think"]
    extra = cfg.get("extra")
    hex_id = cfg.get("hex_id")
    return model_name, mode_id, think_mode, None, extra, image_gen, hex_id