Yash030/claude-code-proxy

1"""Model capability registry for intelligent routing."""2 3from __future__ import annotations4 5from dataclasses import dataclass6from typing import TYPE_CHECKING7 8if TYPE_CHECKING:9    from collections.abc import Sequence10 11 12@dataclass(frozen=True, slots=True)13class ModelCapabilities:14    """Capabilities of a specific model for routing decisions."""15 16    provider_id: str17    model_id: str18    model_ref: str  # provider/model format19    vision: bool = False  # Can process images20    supports_base64_images: bool = False  # Accepts data: URLs with base6421    supports_remote_images: bool = False  # Accepts http/https URLs22    supports_pdfs: bool = False  # Can process PDF documents23    max_images: int = 0  # Max images per request (0 = unlimited)24    max_image_size_mb: float = 10.0  # Max size per image in MB25    coding: bool = False  # Good at code generation/analysis26    reasoning: bool = False  # Strong reasoning/thinking27    general_text: bool = True  # General text generation28    multimodal_input: bool = False  # Can handle multiple input types29    multimodal_output: bool = False  # Can produce multiple output types30    max_tokens: int = 409631    speed: str = "medium"  # "fast", "medium", "slow"32    priority: int = 100  # Higher = preferred for its capabilities33 34 35# Registry of all available models and their capabilities36# This can be extended with actual model discovery later37MODEL_CAPABILITIES: dict[str, ModelCapabilities] = {38    # Zen/minimax models39    "zen/minimax-m2.5-free": ModelCapabilities(40        provider_id="zen",41        model_id="minimax-m2.5-free",42        model_ref="zen/minimax-m2.5-free",43        coding=True,44        reasoning=True,45        general_text=True,46        max_tokens=32000,47        speed="fast",48        priority=80,49    ),50    # NVIDIA NIM models51    "nvidia_nim/stepfun-ai/step-3.5-flash": ModelCapabilities(52        provider_id="nvidia_nim",53        model_id="step-3.5-flash",54        model_ref="nvidia_nim/stepfun-ai/step-3.5-flash",55        coding=True,56        reasoning=True,57        general_text=True,58        max_tokens=32000,59        speed="fast",60        priority=70,61    ),62    "nvidia_nim/qwen/qwen3-coder-480b-a35b-instruct": ModelCapabilities(63        provider_id="nvidia_nim",64        model_id="qwen3-coder-480b-a35b-instruct",65        model_ref="nvidia_nim/qwen/qwen3-coder-480b-a35b-instruct",66        coding=True,67        reasoning=True,68        general_text=True,69        max_tokens=32000,70        speed="slow",71        priority=90,72    ),73    "nvidia_nim/mistralai/mistral-large-3-675b-instruct-2512": ModelCapabilities(74        provider_id="nvidia_nim",75        model_id="mistral-large-3-675b-instruct-2512",76        model_ref="nvidia_nim/mistralai/mistral-large-3-675b-instruct-2512",77        vision=True,78        supports_base64_images=True,79        supports_remote_images=False,80        max_images=16,81        max_image_size_mb=10.0,82        multimodal_input=True,83        coding=True,84        reasoning=True,85        general_text=True,86        max_tokens=32000,87        speed="slow",88        priority=90,89    ),90    "nvidia_nim/abacusai/dracarys-llama-3.1-70b-instruct": ModelCapabilities(91        provider_id="nvidia_nim",92        model_id="dracarys-llama-3.1-70b-instruct",93        model_ref="nvidia_nim/abacusai/dracarys-llama-3.1-70b-instruct",94        coding=True,95        reasoning=True,96        general_text=True,97        max_tokens=32000,98        speed="medium",99        priority=75,100    ),101    "nvidia_nim/z-ai/glm4.7": ModelCapabilities(102        provider_id="nvidia_nim",103        model_id="glm4.7",104        model_ref="nvidia_nim/z-ai/glm4.7",105        coding=True,106        reasoning=True,107        general_text=True,108        max_tokens=32000,109        speed="medium",110        priority=70,111    ),112    "nvidia_nim/bytedance/seed-oss-36b-instruct": ModelCapabilities(113        provider_id="nvidia_nim",114        model_id="seed-oss-36b-instruct",115        model_ref="nvidia_nim/bytedance/seed-oss-36b-instruct",116        coding=True,117        reasoning=True,118        general_text=True,119        max_tokens=32000,120        speed="medium",121        priority=65,122    ),123    "nvidia_nim/mistralai/mistral-nemotron": ModelCapabilities(124        provider_id="nvidia_nim",125        model_id="mistral-nemotron",126        model_ref="nvidia_nim/mistralai/mistral-nemotron",127        coding=True,128        reasoning=True,129        general_text=True,130        max_tokens=32000,131        speed="medium",132        priority=60,133    ),134    # Cerebras models (key only has access to llama3.1-8b currently)135    # Note: qwen-3-235b-a22b-instruct-2507 exists but is rate-limited136    # Note: zai-glm-4.7 and gpt-oss-120b are not accessible with current key137    "cerebras/llama3.1-8b": ModelCapabilities(138        provider_id="cerebras",139        model_id="llama3.1-8b",140        model_ref="cerebras/llama3.1-8b",141        coding=True,142        reasoning=False,143        general_text=True,144        max_tokens=32000,145        speed="fast",146        priority=60,147    ),148    # Silicon Flow models149    "silicon/Qwen/Qwen3.6-35B-A3B": ModelCapabilities(150        provider_id="silicon",151        model_id="Qwen/Qwen3.6-35B-A3B",152        model_ref="silicon/Qwen/Qwen3.6-35B-A3B",153        vision=True,154        supports_base64_images=True,155        max_images=1,156        multimodal_input=True,157        coding=True,158        reasoning=True,159        general_text=True,160        max_tokens=262144,161        speed="medium",162        priority=85,163    ),164    "silicon/Qwen/Qwen3.6-27B": ModelCapabilities(165        provider_id="silicon",166        model_id="Qwen/Qwen3.6-27B",167        model_ref="silicon/Qwen/Qwen3.6-27B",168        vision=True,169        supports_base64_images=True,170        max_images=1,171        multimodal_input=True,172        coding=True,173        reasoning=True,174        general_text=True,175        max_tokens=262144,176        speed="medium",177        priority=82,178    ),179    "silicon/Qwen/Qwen3.5-35B-A3B": ModelCapabilities(180        provider_id="silicon",181        model_id="Qwen/Qwen3.5-35B-A3B",182        model_ref="silicon/Qwen/Qwen3.5-35B-A3B",183        vision=True,184        supports_base64_images=True,185        max_images=1,186        multimodal_input=True,187        coding=True,188        reasoning=True,189        general_text=True,190        max_tokens=262144,191        speed="medium",192        priority=80,193    ),194    "silicon/Qwen/Qwen3.5-27B": ModelCapabilities(195        provider_id="silicon",196        model_id="Qwen/Qwen3.5-27B",197        model_ref="silicon/Qwen/Qwen3.5-27B",198        vision=True,199        supports_base64_images=True,200        max_images=1,201        multimodal_input=True,202        coding=True,203        reasoning=True,204        general_text=True,205        max_tokens=262144,206        speed="medium",207        priority=78,208    ),209    "silicon/google/gemma-4-26B-A4B-it": ModelCapabilities(210        provider_id="silicon",211        model_id="google/gemma-4-26B-A4B-it",212        model_ref="silicon/google/gemma-4-26B-A4B-it",213        coding=True,214        reasoning=True,215        general_text=True,216        max_tokens=262144,217        speed="fast",218        priority=75,219    ),220    "silicon/google/gemma-4-31B-it": ModelCapabilities(221        provider_id="silicon",222        model_id="google/gemma-4-31B-it",223        model_ref="silicon/google/gemma-4-31B-it",224        coding=True,225        reasoning=True,226        general_text=True,227        max_tokens=262144,228        speed="fast",229        priority=76,230    ),231    # Groq models232    "groq/llama-3.3-70b-versatile": ModelCapabilities(233        provider_id="groq",234        model_id="llama-3.3-70b-versatile",235        model_ref="groq/llama-3.3-70b-versatile",236        coding=True,237        reasoning=True,238        general_text=True,239        max_tokens=32768,240        speed="fast",241        priority=85,242    ),243    "groq/llama-3.1-8b-instant": ModelCapabilities(244        provider_id="groq",245        model_id="llama-3.1-8b-instant",246        model_ref="groq/llama-3.1-8b-instant",247        coding=True,248        general_text=True,249        max_tokens=131072,250        speed="fast",251        priority=90,252    ),253    "groq/qwen3-32b": ModelCapabilities(254        provider_id="groq",255        model_id="qwen3-32b",256        model_ref="groq/qwen3-32b",257        coding=True,258        reasoning=True,259        general_text=True,260        max_tokens=40960,261        speed="medium",262        priority=88,263    ),264}265 266 267def get_model_capabilities(model_ref: str) -> ModelCapabilities | None:268    """Get capabilities for a specific model reference."""269    return MODEL_CAPABILITIES.get(model_ref)270 271 272def find_models_with_capability(capability: str) -> list[ModelCapabilities]:273    """Find all models that have a specific capability."""274    results = []275    for caps in MODEL_CAPABILITIES.values():276        if getattr(caps, capability, False):277            results.append(caps)278    # Sort by priority (higher = better)279    results.sort(key=lambda x: x.priority, reverse=True)280    return results281 282 283def find_best_model_for_task(284    required_capabilities: set[str],285    available_models: Sequence[str] | None = None,286) -> ModelCapabilities | None:287    """Find the best model matching required capabilities.288 289    Args:290        required_capabilities: Set of capability names needed (e.g., {"coding", "vision"})291        available_models: Optional list of model refs to filter by292 293    Returns:294        Best matching ModelCapabilities or None295    """296    candidates = []297 298    models_to_check = (299        [MODEL_CAPABILITIES[m] for m in available_models if m in MODEL_CAPABILITIES]300        if available_models301        else list(MODEL_CAPABILITIES.values())302    )303 304    for caps in models_to_check:305        # Check if model has all required capabilities306        if all(getattr(caps, cap, False) for cap in required_capabilities):307            candidates.append(caps)308 309    if not candidates:310        return None311 312    # Sort by priority and return best313    candidates.sort(key=lambda x: x.priority, reverse=True)314    return candidates[0]315 316 317def get_capability_match_score(318    model_caps: ModelCapabilities,319    required: set[str],320) -> tuple[int, int]:321    """Calculate match score for routing.322 323    Returns (matched_count, priority) for sorting.324    """325    matched = sum(1 for cap in required if getattr(model_caps, cap, False))326    return (matched, model_caps.priority)327