open-swe/agent/dashboard/options.py

"""Supported models and reasoning efforts surfaced in the profile editor."""

from __future__ import annotations

from typing import TypedDict


class ModelOption(TypedDict):
    id: str
    label: str
    efforts: list[str]
    default_effort: str
    supports_images: bool


SUPPORTED_MODELS: list[ModelOption] = [
    {
        "id": "bedrock_converse:us.anthropic.claude-opus-4-8",
        "label": "Opus 4.8 (Bedrock)",
        "efforts": ["low", "medium", "high", "xhigh", "max"],
        "default_effort": "high",
        "supports_images": True,
    },
    {
        "id": "bedrock_converse:us.anthropic.claude-sonnet-5",
        "label": "Sonnet 5 (Bedrock)",
        "efforts": ["low", "medium", "high", "xhigh", "max"],
        "default_effort": "high",
        "supports_images": True,
    },
    {
        "id": "fireworks:accounts/fireworks/models/kimi-k2p7-code",
        "label": "Kimi K2.7",
        "efforts": ["low", "medium", "high"],
        "default_effort": "high",
        "supports_images": False,
    },
    {
        "id": "fireworks:accounts/fireworks/models/deepseek-v4-pro",
        "label": "DeepSeek V4 Pro",
        "efforts": ["none", "low", "medium", "high", "xhigh", "max"],
        "default_effort": "high",
        "supports_images": False,
    },
    {
        "id": "fireworks:accounts/fireworks/models/glm-5p2",
        "label": "GLM 5.2",
        "efforts": ["none", "high", "max"],
        "default_effort": "high",
        "supports_images": False,
    },
    {
        "id": "fireworks:accounts/fireworks/models/minimax-m3",
        "label": "MiniMax M3",
        "efforts": ["medium", "high"],
        "default_effort": "high",
        "supports_images": True,
    },
    {
        "id": "fireworks:accounts/fireworks/models/gpt-oss-120b",
        "label": "gpt-oss-120b",
        "efforts": ["low", "medium", "high"],
        "default_effort": "medium",
        "supports_images": False,
    },
    {
        "id": "fireworks:accounts/fireworks/models/deepseek-v4-flash",
        "label": "DeepSeek V4 Flash",
        "efforts": ["none", "medium", "high"],
        "default_effort": "high",
        "supports_images": False,
    },
]

SUPPORTED_MODEL_IDS: frozenset[str] = frozenset(m["id"] for m in SUPPORTED_MODELS)

DEFAULT_MODEL_ID: str = "bedrock_converse:us.anthropic.claude-opus-4-8"
DEFAULT_MODEL_EFFORT: str = "medium"


def model_supports_effort(model_id: str, effort: str) -> bool:
    for m in SUPPORTED_MODELS:
        if m["id"] == model_id:
            return effort in m["efforts"]
    return False


def model_supports_images(model_id: str) -> bool:
    for m in SUPPORTED_MODELS:
        if m["id"] == model_id:
            return m["supports_images"]
    return False


def _provider_of(model_id: str) -> str | None:
    provider, _, rest = model_id.partition(":")
    return provider if rest else None


def _claude_family_of(model_id: str) -> str | None:
    provider, _, name = model_id.partition(":")
    if provider == "anthropic":
        claude = name
    elif provider == "bedrock_converse":
        # Bedrock ids embed the model as a region-prefixed path, e.g.
        # "us.anthropic.claude-sonnet-5" — take the trailing "claude-*" segment so
        # our Bedrock Claude models get the same family-aware fallback.
        claude = name.rpartition(".")[2]
    else:
        return None
    if not claude.startswith("claude-"):
        return None
    parts = claude.split("-")
    if len(parts) < 2:
        return None
    return "-".join(parts[:2])


def _fallback_effort_for(model: ModelOption, effort: object) -> str | None:
    if not isinstance(effort, str):
        return None
    if effort in model["efforts"]:
        return effort
    if (
        model["id"].startswith("google_genai:")
        and effort == "none"
        and "minimal" in model["efforts"]
    ):
        return "minimal"
    return None


def provider_fallback_pair(model_id: object, effort: object = None) -> tuple[str, str] | None:
    """Newest supported ``(model_id, effort)`` for the same provider/family.

    Keeps a stored selection on its original provider when its exact id has
    dropped out of the supported set (e.g. an Opus minor-version bump), preferring
    the same Claude family when available instead of falling through to the
    cross-provider global default. Preserves ``effort`` when the fallback model
    supports it, otherwise uses that model's default effort. Returns ``None`` when
    no supported model shares the provider.
    """
    if not isinstance(model_id, str):
        return None
    provider = _provider_of(model_id)
    if provider is None:
        return None
    family = _claude_family_of(model_id)
    if family is not None:
        for m in SUPPORTED_MODELS:
            if _provider_of(m["id"]) == provider and _claude_family_of(m["id"]) == family:
                return m["id"], _fallback_effort_for(m, effort) or m["default_effort"]
    for m in SUPPORTED_MODELS:
        if _provider_of(m["id"]) == provider:
            return m["id"], _fallback_effort_for(m, effort) or m["default_effort"]
    return None


def default_model_pair() -> tuple[str, str]:
    """Hardcoded fallback (model_id, reasoning_effort) used when no team default is set."""
    if DEFAULT_MODEL_ID in SUPPORTED_MODEL_IDS and model_supports_effort(
        DEFAULT_MODEL_ID, DEFAULT_MODEL_EFFORT
    ):
        return DEFAULT_MODEL_ID, DEFAULT_MODEL_EFFORT
    first = SUPPORTED_MODELS[0]
    return first["id"], first["default_effort"]


def default_vision_model_pair() -> tuple[str, str]:
    """Default (model_id, reasoning_effort) to use when image input is required.

    Prefers the configured default model when it is vision-capable (Opus 4.8),
    otherwise the first vision-capable supported model. The fork ships only
    Bedrock/Fireworks models, so this selects on ``supports_images`` rather than
    upstream's ``openai:``/``anthropic:`` provider filter.
    """
    if (
        DEFAULT_MODEL_ID in SUPPORTED_MODEL_IDS
        and model_supports_images(DEFAULT_MODEL_ID)
        and model_supports_effort(DEFAULT_MODEL_ID, DEFAULT_MODEL_EFFORT)
    ):
        return DEFAULT_MODEL_ID, DEFAULT_MODEL_EFFORT
    for model in SUPPORTED_MODELS:
        if model["supports_images"]:
            return model["id"], model["default_effort"]
    return default_model_pair()