"""Model configuration helpers and preset models.""" from typing import Any # Preset provider models for quick selection PROVIDER_PRESETS = { "GPT-5.2 Pro": { "model": "gpt-5.2-pro-2025-12-11", "tokenizer": "openai/gpt-5.2-pro-2025-12-11", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "GPT-5.2": { "model": "gpt-5.2-2025-12-11", "tokenizer": "openai/gpt-5.2-2025-12-11", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "GPT-5.1": { "model": "gpt-5.1-2025-11-13", "tokenizer": "openai/gpt-5.1-2025-11-13", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "GPT-5": { "model": "gpt-5-2025-08-07", "tokenizer": "openai/gpt-5-2025-08-07", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "GPT-5 Mini": { "model": "gpt-5-mini-2025-08-07", "tokenizer": "openai/gpt-5-mini-2025-08-07", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "GPT-5 Nano": { "model": "gpt-5-nano-2025-08-07", "tokenizer": "openai/gpt-5-nano-2025-08-07", "max_context": 272000, "api_key_env": "OPENAI_API_KEY", }, "Claude Opus 4.5": { "model": "claude-opus-4-5-202511019", "tokenizer": "anthropic/claude-opus-4-5-20251101", "max_context": 200000, "api_key_env": "ANTHROPIC_API_KEY", }, "Claude Sonnet 4.5": { "model": "claude-sonnet-4-5-20250929", "tokenizer": "anthropic/claude-sonnet-4-5-20250929", "max_context": 200000, "api_key_env": "ANTHROPIC_API_KEY", }, "Claude Haiku 4.5": { "model": "claude-haiku-4-5-20251001", "tokenizer": "anthropic/claude-haiku-4-5-20251001", "max_context": 200000, "api_key_env": "ANTHROPIC_API_KEY", }, "Gemini 3.0 Pro": { "model": "gemini/gemini-3-pro-preview", "tokenizer": "gemini/gemini-3-pro-preview", "max_context": 2000000, "api_key_env": "GEMINI_API_KEY", }, "Gemini 3.0 Flash": { "model": "gemini/gemini-3-flash-preview", "tokenizer": "gemini/gemini-3-flash-preview", "max_context": 1000000, "api_key_env": "GEMINI_API_KEY", }, "Gemini 2.5 Pro": { "model": "gemini/gemini-2.5-pro", "tokenizer": "gemini/gemini-2.5-pro", "max_context": 2000000, "api_key_env": "GEMINI_API_KEY", }, "Gemini 2.5 Flash": { "model": "gemini/gemini-2.5-flash", "tokenizer": "gemini/gemini-2.5-flash", "max_context": 1000000, "api_key_env": "GEMINI_API_KEY", }, } def create_local_model_config( model: str, api_base: str | None = None, max_context: int = 131072, tokenizer: str | None = None, ) -> dict[str, Any]: """ Create a local model configuration. Args: model: Model name (e.g., "ollama/llama2" or "gpt-3.5-turbo" for vLLM) api_base: API base URL (defaults based on model type) max_context: Maximum context window size tokenizer: Tokenizer name for token counting Returns: Model configuration dictionary """ # Set default API base based on model type if api_base is None: if model.startswith("ollama/"): api_base = "http://localhost:11434" elif model.startswith("vllm/") or not model.startswith(("ollama/", "openrouter/")): # Assume OpenAI-compatible (vLLM) api_base = "http://localhost:8000/v1" # Infer tokenizer if not provided if tokenizer is None: if model.startswith("ollama/"): # Try to infer from model name model_name = model.replace("ollama/", "") tokenizer = f"hf/{model_name}" else: # For vLLM/OpenAI-compatible, try to infer tokenizer = model.replace("vllm/", "") return { "type": "local", "model": model, "api_base": api_base, "max_context": max_context, "tokenizer": tokenizer, } def create_provider_model_config( model: str, api_key: str, max_context: int = 128000, tokenizer: str | None = None, ) -> dict[str, Any]: """ Create a provider model configuration. Args: model: Model name in litellm format api_key: API key for the provider max_context: Maximum context window size tokenizer: Tokenizer name for token counting Returns: Model configuration dictionary """ # Infer tokenizer if not provided if tokenizer is None: # Try to infer from model name if model.startswith("openai/") or not "/" in model: # OpenAI models model_name = model.replace("openai/", "") tokenizer = f"openai/{model_name}" elif model.startswith("anthropic/") or model.startswith("claude-"): # Anthropic models model_name = model.replace("anthropic/", "") tokenizer = f"anthropic/{model_name}" elif model.startswith("gemini/"): # Gemini models model_name = model.replace("gemini/", "") tokenizer = f"gemini/{model_name}" else: # Generic fallback tokenizer = "gpt2" return { "type": "provider", "model": model, "api_key": api_key, "max_context": max_context, "tokenizer": tokenizer, } def get_provider_from_model(model: str) -> str: """ Infer provider name from model identifier. Args: model: Model name in litellm format Returns: Provider name hint (e.g., "openai", "anthropic", "gemini") """ model_lower = model.lower() if model_lower.startswith("openai/") or not "/" in model: return "openai" elif model_lower.startswith("anthropic/") or model_lower.startswith("claude-"): return "anthropic" elif model_lower.startswith("gemini/"): return "gemini" elif model_lower.startswith("openrouter/"): return "openrouter" elif model_lower.startswith("cohere/"): return "cohere" elif model_lower.startswith("mistral/"): return "mistral" else: return "other" def get_api_key_env_name(provider: str) -> str: """ Get the environment variable name for API key based on provider. Args: provider: Provider name Returns: Environment variable name for API key """ provider_to_key = { "openai": "OPENAI_API_KEY", "anthropic": "ANTHROPIC_API_KEY", "gemini": "GEMINI_API_KEY", "openrouter": "OPENROUTER_API_KEY", "cohere": "COHERE_API_KEY", "mistral": "MISTRAL_API_KEY", "other": "API_KEY", } return provider_to_key.get(provider.lower(), "API_KEY")