diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 4e55978793e..e5c17ee6cc6 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5855 }, "reportMissingTypeArgument": { - "limit": 15849 + "limit": 15727 }, "reportMissingTypeStubs": { "limit": 41 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45269 + "limit": 45131 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40452 + "limit": 40040 }, "reportUnknownParameterType": { - "limit": 20309 + "limit": 20308 }, "reportUnknownVariableType": { - "limit": 31978 + "limit": 31834 }, "reportUnnecessaryCast": { "limit": 173 diff --git a/litellm/__init__.py b/litellm/__init__.py index 62c41c2959b..d4045985503 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -50,7 +50,8 @@ from litellm._logging import ( _turn_on_json, log_level, ) -import re +from collections.abc import Mapping, Sequence +from functools import lru_cache from litellm.constants import ( DEFAULT_BATCH_SIZE, DEFAULT_FLUSH_INTERVAL_SECONDS, @@ -71,8 +72,6 @@ from litellm.constants import ( replicate_models, clarifai_models, huggingface_models, - modelscope_models, - empower_models, together_ai_models, baseten_models, WANDB_MODELS, @@ -556,402 +555,13 @@ config_path = None vertex_ai_safety_settings: Optional[dict] = None ####### COMPLETION MODELS ################### -from typing import Set - -open_ai_chat_completion_models: Set = set() -open_ai_text_completion_models: Set = set() -cohere_models: Set = set() -cohere_chat_models: Set = set() -mistral_chat_models: Set = set() -text_completion_codestral_models: Set = set() -text_completion_inception_models: Set = set() -anthropic_models: Set = set() -openrouter_models: Set = set() -datarobot_models: Set = set() -vertex_language_models: Set = set() -vertex_vision_models: Set = set() -vertex_chat_models: Set = set() -vertex_code_chat_models: Set = set() -vertex_ai_image_models: Set = set() -vertex_ai_video_models: Set = set() -vertex_text_models: Set = set() -vertex_code_text_models: Set = set() -vertex_embedding_models: Set = set() -vertex_anthropic_models: Set = set() -vertex_llama3_models: Set = set() -vertex_deepseek_models: Set = set() -vertex_ai_ai21_models: Set = set() -vertex_mistral_models: Set = set() -vertex_openai_models: Set = set() -vertex_minimax_models: Set = set() -vertex_moonshot_models: Set = set() -vertex_zai_models: Set = set() -ai21_models: Set = set() -ai21_chat_models: Set = set() -nlp_cloud_models: Set = set() -aleph_alpha_models: Set = set() -bedrock_models: Set = set() -bedrock_converse_models: Set = set(BEDROCK_CONVERSE_MODELS) -fal_ai_models: Set = set() -fireworks_ai_models: Set = set() -fireworks_ai_embedding_models: Set = set() -deepinfra_models: Set = set() -perplexity_models: Set = set() -watsonx_models: Set = set() -gemini_models: Set = set() -xai_models: Set = set() -zai_models: Set = set() -deepseek_models: Set = set() -tencent_models: Set = set() -runwayml_models: Set = set() -azure_ai_models: Set = set() -jina_ai_models: Set = set() -voyage_models: Set = set() -infinity_models: Set = set() -heroku_models: Set = set() -databricks_models: Set = set() -cloudflare_models: Set = set() -codestral_models: Set = set() -friendliai_models: Set = set() -featherless_ai_models: Set = set() -palm_models: Set = set() -groq_models: Set = set() -azure_models: Set = set() -azure_anthropic_models: Set = set() -azure_text_models: Set = set() -anyscale_models: Set = set() -cerebras_models: Set = set() -galadriel_models: Set = set() -nvidia_nim_models: Set = set() -nvidia_riva_models: Set = set() -soniox_models: Set = set() -sambanova_models: Set = set() -sambanova_embedding_models: Set = set() -novita_models: Set = set() -assemblyai_models: Set = set() -snowflake_models: Set = set() -gradient_ai_models: Set = set() -llama_models: Set = set() -nscale_models: Set = set() -nebius_models: Set = set() -nebius_embedding_models: Set = set() -aiml_models: Set = set() -deepgram_models: Set = set() -elevenlabs_models: Set = set() -dashscope_models: Set = set() -moonshot_models: Set = set() -publicai_models: Set = set() -darkbloom_models: Set = set() -v0_models: Set = set() -morph_models: Set = set() -lambda_ai_models: Set = set() -inception_models: Set = set() -hyperbolic_models: Set = set() -black_forest_labs_models: Set = set() -recraft_models: Set = set() -cometapi_models: Set = set() -oci_models: Set = set() -vercel_ai_gateway_models: Set = set() -volcengine_models: Set = set() -wandb_models: Set = set(WANDB_MODELS) -ovhcloud_models: Set = set() -ovhcloud_embedding_models: Set = set() -lemonade_models: Set = set() -docker_model_runner_models: Set = set() -amazon_nova_models: Set = set() -stability_models: Set = set() -github_copilot_models: Set = set() -chatgpt_models: Set = set() -minimax_models: Set = set() -aws_polly_models: Set = set() -gigachat_models: Set = set() -llamagate_models: Set = set() -reducto_models: Set = set() -bedrock_mantle_models: Set = set() - - -def is_bedrock_pricing_only_model(key: str) -> bool: - """ - Excludes keys with the pattern 'bedrock//'. These are in the model_prices_and_context_window.json file for pricing purposes only. - - Args: - key (str): A key to filter. - - Returns: - bool: True if the key matches the Bedrock pattern, False otherwise. - """ - # Regex to match 'bedrock//' - bedrock_pattern = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$") - - if "month-commitment" in key: - return True - - is_match = bedrock_pattern.match(key) - return is_match is not None - - -def is_openai_finetune_model(key: str) -> bool: - """ - Excludes model cost keys with the pattern 'ft:'. These are in the model_prices_and_context_window.json file for pricing purposes only. - - Args: - key (str): A key to filter. - - Returns: - bool: True if the key matches the OpenAI finetune pattern, False otherwise. - """ - return key.startswith("ft:") and not key.count(":") > 1 - - -def add_known_models(model_cost_map: Optional[Dict] = None): - _map = model_cost_map if model_cost_map is not None else model_cost - for key, value in _map.items(): - if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): - open_ai_chat_completion_models.add(key) - elif value.get("litellm_provider") == "text-completion-openai": - open_ai_text_completion_models.add(key) - elif value.get("litellm_provider") == "azure_text": - azure_text_models.add(key) - elif value.get("litellm_provider") == "cohere": - cohere_models.add(key) - elif value.get("litellm_provider") == "cohere_chat": - cohere_chat_models.add(key) - elif value.get("litellm_provider") == "mistral": - mistral_chat_models.add(key) - elif value.get("litellm_provider") == "anthropic": - anthropic_models.add(key) - elif value.get("litellm_provider") == "empower": - empower_models.add(key) - elif value.get("litellm_provider") == "openrouter": - openrouter_models.add(key) - elif value.get("litellm_provider") == "vercel_ai_gateway": - vercel_ai_gateway_models.add(key) - elif value.get("litellm_provider") == "datarobot": - datarobot_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-text-models": - vertex_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-text-models": - vertex_code_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-language-models": - vertex_language_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-vision-models": - vertex_vision_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-chat-models": - vertex_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-chat-models": - vertex_code_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-embedding-models": - vertex_embedding_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-anthropic_models": - key = key.replace("vertex_ai/", "") - vertex_anthropic_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-llama_models": - key = key.replace("vertex_ai/", "") - vertex_llama3_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-deepseek_models": - key = key.replace("vertex_ai/", "") - vertex_deepseek_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-mistral_models": - key = key.replace("vertex_ai/", "") - vertex_mistral_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-ai21_models": - key = key.replace("vertex_ai/", "") - vertex_ai_ai21_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-image-models": - key = key.replace("vertex_ai/", "") - vertex_ai_image_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-video-models": - key = key.replace("vertex_ai/", "") - vertex_ai_video_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-openai_models": - key = key.replace("vertex_ai/", "") - vertex_openai_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-minimax_models": - key = key.replace("vertex_ai/", "") - vertex_minimax_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-moonshot_models": - key = key.replace("vertex_ai/", "") - vertex_moonshot_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-zai_models": - key = key.replace("vertex_ai/", "") - vertex_zai_models.add(key) - elif value.get("litellm_provider") == "ai21": - if value.get("mode") == "chat": - ai21_chat_models.add(key) - else: - ai21_models.add(key) - elif value.get("litellm_provider") == "nlp_cloud": - nlp_cloud_models.add(key) - elif value.get("litellm_provider") == "aleph_alpha": - aleph_alpha_models.add(key) - elif value.get("litellm_provider") == "bedrock" and not is_bedrock_pricing_only_model(key): - bedrock_models.add(key) - elif value.get("litellm_provider") == "bedrock_converse": - bedrock_converse_models.add(key) - elif value.get("litellm_provider") == "deepinfra": - deepinfra_models.add(key) - elif value.get("litellm_provider") == "perplexity": - perplexity_models.add(key) - elif value.get("litellm_provider") == "watsonx": - watsonx_models.add(key) - elif value.get("litellm_provider") == "gemini": - gemini_models.add(key) - elif value.get("litellm_provider") == "fireworks_ai": - # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. - if "-to-" not in key and "fireworks-ai-default" not in key: - fireworks_ai_models.add(key) - elif value.get("litellm_provider") == "fireworks_ai-embedding-models": - # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. - if "-to-" not in key: - fireworks_ai_embedding_models.add(key) - elif value.get("litellm_provider") == "text-completion-codestral": - text_completion_codestral_models.add(key) - elif value.get("litellm_provider") == "text-completion-inception": - text_completion_inception_models.add(key) - elif value.get("litellm_provider") == "xai": - xai_models.add(key) - elif value.get("litellm_provider") == "zai": - zai_models.add(key) - elif value.get("litellm_provider") == "fal_ai": - fal_ai_models.add(key) - elif value.get("litellm_provider") == "deepseek": - deepseek_models.add(key) - elif value.get("litellm_provider") == "tencent": - tencent_models.add(key) - elif value.get("litellm_provider") == "runwayml": - runwayml_models.add(key) - elif value.get("litellm_provider") == "meta_llama": - llama_models.add(key) - elif value.get("litellm_provider") == "nscale": - nscale_models.add(key) - elif value.get("litellm_provider") == "azure_ai": - azure_ai_models.add(key) - elif value.get("litellm_provider") == "voyage": - voyage_models.add(key) - elif value.get("litellm_provider") == "infinity": - infinity_models.add(key) - elif value.get("litellm_provider") == "databricks": - databricks_models.add(key) - elif value.get("litellm_provider") == "cloudflare": - cloudflare_models.add(key) - elif value.get("litellm_provider") == "codestral": - codestral_models.add(key) - elif value.get("litellm_provider") == "friendliai": - friendliai_models.add(key) - elif value.get("litellm_provider") == "palm": - palm_models.add(key) - elif value.get("litellm_provider") == "groq": - groq_models.add(key) - elif value.get("litellm_provider") == "azure": - azure_models.add(key) - elif value.get("litellm_provider") == "azure_anthropic": - azure_anthropic_models.add(key) - elif value.get("litellm_provider") == "anyscale": - anyscale_models.add(key) - elif value.get("litellm_provider") == "cerebras": - cerebras_models.add(key) - elif value.get("litellm_provider") == "galadriel": - galadriel_models.add(key) - elif value.get("litellm_provider") == "nvidia_nim": - nvidia_nim_models.add(key) - elif value.get("litellm_provider") == "nvidia_riva": - nvidia_riva_models.add(key) - elif value.get("litellm_provider") == "soniox": - soniox_models.add(key) - elif value.get("litellm_provider") == "sambanova": - sambanova_models.add(key) - elif value.get("litellm_provider") == "sambanova-embedding-models": - sambanova_embedding_models.add(key) - elif value.get("litellm_provider") == "novita": - novita_models.add(key) - elif value.get("litellm_provider") == "nebius-chat-models": - nebius_models.add(key) - elif value.get("litellm_provider") == "nebius-embedding-models": - nebius_embedding_models.add(key) - elif value.get("litellm_provider") == "aiml": - aiml_models.add(key) - elif value.get("litellm_provider") == "assemblyai": - assemblyai_models.add(key) - elif value.get("litellm_provider") == "jina_ai": - jina_ai_models.add(key) - elif value.get("litellm_provider") == "snowflake": - snowflake_models.add(key) - elif value.get("litellm_provider") == "gradient_ai": - gradient_ai_models.add(key) - elif value.get("litellm_provider") == "featherless_ai": - featherless_ai_models.add(key) - elif value.get("litellm_provider") == "deepgram": - deepgram_models.add(key) - elif value.get("litellm_provider") == "elevenlabs": - elevenlabs_models.add(key) - elif value.get("litellm_provider") == "heroku": - heroku_models.add(key) - elif value.get("litellm_provider") == "dashscope": - dashscope_models.add(key) - elif value.get("litellm_provider") == "modelscope": - modelscope_models.add(key) - elif value.get("litellm_provider") == "moonshot": - moonshot_models.add(key) - elif value.get("litellm_provider") == "publicai": - publicai_models.add(key) - elif value.get("litellm_provider") == "darkbloom": - darkbloom_models.add(key) - elif value.get("litellm_provider") == "v0": - v0_models.add(key) - elif value.get("litellm_provider") == "morph": - morph_models.add(key) - elif value.get("litellm_provider") == "lambda_ai": - lambda_ai_models.add(key) - elif value.get("litellm_provider") == "inception": - inception_models.add(key) - elif value.get("litellm_provider") == "hyperbolic": - hyperbolic_models.add(key) - elif value.get("litellm_provider") == "black_forest_labs": - black_forest_labs_models.add(key) - elif value.get("litellm_provider") == "recraft": - recraft_models.add(key) - elif value.get("litellm_provider") == "cometapi": - cometapi_models.add(key) - elif value.get("litellm_provider") == "oci": - oci_models.add(key) - elif value.get("litellm_provider") == "volcengine": - volcengine_models.add(key) - elif value.get("litellm_provider") == "wandb": - wandb_models.add(key) - elif value.get("litellm_provider") == "ovhcloud": - ovhcloud_models.add(key) - elif value.get("litellm_provider") == "ovhcloud-embedding-models": - ovhcloud_embedding_models.add(key) - elif value.get("litellm_provider") == "lemonade": - lemonade_models.add(key) - elif value.get("litellm_provider") == "docker_model_runner": - docker_model_runner_models.add(key) - elif value.get("litellm_provider") == "amazon_nova": - amazon_nova_models.add(key) - elif value.get("litellm_provider") == "stability": - stability_models.add(key) - elif value.get("litellm_provider") == "github_copilot": - github_copilot_models.add(key) - elif value.get("litellm_provider") == "chatgpt": - chatgpt_models.add(key) - elif value.get("litellm_provider") == "minimax": - minimax_models.add(key) - elif value.get("litellm_provider") == "aws_polly": - aws_polly_models.add(key) - elif value.get("litellm_provider") == "gigachat": - gigachat_models.add(key) - elif value.get("litellm_provider") == "llamagate": - llamagate_models.add(key) - elif value.get("litellm_provider") == "reducto": - reducto_models.add(key) - elif value.get("litellm_provider") == "bedrock_mantle": - bedrock_mantle_models.add(key) - - -add_known_models() -# known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary - -# this is maintained for Exception Mapping +from litellm.litellm_core_utils.model_registry import ( + ModelRegistrySnapshot, + build_snapshot, + extend_snapshot, + is_bedrock_pricing_only_model as is_bedrock_pricing_only_model, + is_openai_finetune_model as is_openai_finetune_model, +) # used for Cost Tracking & Token counting @@ -978,205 +588,188 @@ ollama_models = ["llama2"] maritalk_models = ["maritalk"] -model_list = list( - open_ai_chat_completion_models - | open_ai_text_completion_models - | cohere_models - | cohere_chat_models - | anthropic_models - | set(replicate_models) - | openrouter_models - | datarobot_models - | set(huggingface_models) - | vertex_chat_models - | vertex_text_models - | ai21_models - | ai21_chat_models - | set(together_ai_models) - | set(baseten_models) - | aleph_alpha_models - | nlp_cloud_models - | set(ollama_models) - | bedrock_models - | deepinfra_models - | perplexity_models - | set(maritalk_models) - | runwayml_models - | vertex_language_models - | watsonx_models - | gemini_models - | text_completion_codestral_models - | text_completion_inception_models - | xai_models - | zai_models - | fal_ai_models - | deepseek_models - | modelscope_models - | azure_ai_models - | voyage_models - | infinity_models - | databricks_models - | cloudflare_models - | codestral_models - | friendliai_models - | palm_models - | groq_models - | azure_models - | azure_anthropic_models - | anyscale_models - | cerebras_models - | galadriel_models - | nvidia_nim_models - | nvidia_riva_models - | soniox_models - | sambanova_models - | azure_text_models - | novita_models - | assemblyai_models - | jina_ai_models - | snowflake_models - | gradient_ai_models - | llama_models - | featherless_ai_models - | nscale_models - | deepgram_models - | elevenlabs_models - | dashscope_models - | moonshot_models - | publicai_models - | darkbloom_models - | v0_models - | morph_models - | lambda_ai_models - | inception_models - | black_forest_labs_models - | recraft_models - | cometapi_models - | oci_models - | heroku_models - | vercel_ai_gateway_models - | volcengine_models - | wandb_models - | ovhcloud_models - | lemonade_models - | docker_model_runner_models - | reducto_models - | bedrock_mantle_models - | set(clarifai_models) -) +_registry_static_model_names: Mapping[str, frozenset[str]] = { + "petals_models": frozenset(petals_models), + "ollama_models": frozenset(ollama_models), + "maritalk_models": frozenset(maritalk_models), +} -model_list_set = set(model_list) +_model_registry_snapshot: ModelRegistrySnapshot + + +@lru_cache(maxsize=1) +def _known_llm_providers() -> frozenset[str]: + from litellm.types.utils import LlmProviders + + return frozenset(provider.value for provider in LlmProviders) + + +def add_known_models(model_cost_map: Mapping[str, Mapping[str, object]] | None = None) -> None: + """ + Re-derive every provider model collection (``litellm.anthropic_models``, + ``litellm.models_by_provider``, ``litellm.model_list``, ...) from the cost map. + + ``model_cost_map`` is merged over ``litellm.model_cost``; models absent from the merged + map are dropped, so callers that replace the cost map must assign it before rebuilding. + """ + global _model_registry_snapshot + _cost_map: Mapping[str, Mapping[str, object]] = ( + {**model_cost, **model_cost_map} if model_cost_map is not None else model_cost + ) + _model_registry_snapshot = build_snapshot( + model_cost=_cost_map, + known_providers=_known_llm_providers(), + static_model_names=_registry_static_model_names, + ) + + +def extend_known_models(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: + """ + Add newly registered cost-map entries to the registry without re-deriving the whole map. + + ``register_model`` runs twice per deployment on proxy startup, which a full rebuild + cannot absorb; additions go through the same derivation rules as ``add_known_models``. + """ + global _model_registry_snapshot + _model_registry_snapshot = extend_snapshot( + snapshot=_model_registry_snapshot, + model_cost_additions=model_cost_map, + known_providers=_known_llm_providers(), + ) + + +add_known_models() + +_REGISTRY_VIEWS: Mapping[str, Callable[[ModelRegistrySnapshot], object]] = { + "models_by_provider": lambda snapshot: snapshot.models_by_provider, + "model_list": lambda snapshot: list(snapshot.model_list), + "model_list_set": lambda snapshot: snapshot.model_list_set, + "all_embedding_models": lambda snapshot: snapshot.all_embedding_models, +} + +if TYPE_CHECKING: + models_by_provider: Mapping[str, frozenset[str]] + model_list: Sequence[str] + model_list_set: frozenset[str] + all_embedding_models: frozenset[str] + ai21_chat_models: frozenset[str] + ai21_models: frozenset[str] + aiml_models: frozenset[str] + aleph_alpha_models: frozenset[str] + amazon_nova_models: frozenset[str] + anthropic_models: frozenset[str] + anyscale_models: frozenset[str] + assemblyai_models: frozenset[str] + aws_polly_models: frozenset[str] + azure_ai_models: frozenset[str] + azure_anthropic_models: frozenset[str] + azure_models: frozenset[str] + azure_text_models: frozenset[str] + bedrock_converse_models: frozenset[str] + bedrock_mantle_models: frozenset[str] + bedrock_models: frozenset[str] + black_forest_labs_models: frozenset[str] + cerebras_models: frozenset[str] + chatgpt_models: frozenset[str] + cloudflare_models: frozenset[str] + codestral_models: frozenset[str] + cohere_chat_models: frozenset[str] + cohere_models: frozenset[str] + cometapi_models: frozenset[str] + darkbloom_models: frozenset[str] + dashscope_models: frozenset[str] + databricks_models: frozenset[str] + datarobot_models: frozenset[str] + deepgram_models: frozenset[str] + deepinfra_models: frozenset[str] + deepseek_models: frozenset[str] + docker_model_runner_models: frozenset[str] + elevenlabs_models: frozenset[str] + empower_models: frozenset[str] + fal_ai_models: frozenset[str] + featherless_ai_models: frozenset[str] + fireworks_ai_embedding_models: frozenset[str] + fireworks_ai_models: frozenset[str] + friendliai_models: frozenset[str] + galadriel_models: frozenset[str] + gemini_models: frozenset[str] + gigachat_models: frozenset[str] + github_copilot_models: frozenset[str] + gradient_ai_models: frozenset[str] + groq_models: frozenset[str] + heroku_models: frozenset[str] + hyperbolic_models: frozenset[str] + inception_models: frozenset[str] + infinity_models: frozenset[str] + jina_ai_models: frozenset[str] + lambda_ai_models: frozenset[str] + lemonade_models: frozenset[str] + llama_models: frozenset[str] + llamagate_models: frozenset[str] + minimax_models: frozenset[str] + mistral_chat_models: frozenset[str] + modelscope_models: frozenset[str] + moonshot_models: frozenset[str] + morph_models: frozenset[str] + nebius_embedding_models: frozenset[str] + nebius_models: frozenset[str] + nlp_cloud_models: frozenset[str] + novita_models: frozenset[str] + nscale_models: frozenset[str] + nvidia_nim_models: frozenset[str] + nvidia_riva_models: frozenset[str] + oci_models: frozenset[str] + open_ai_chat_completion_models: frozenset[str] + open_ai_text_completion_models: frozenset[str] + openrouter_models: frozenset[str] + ovhcloud_embedding_models: frozenset[str] + ovhcloud_models: frozenset[str] + palm_models: frozenset[str] + perplexity_models: frozenset[str] + publicai_models: frozenset[str] + recraft_models: frozenset[str] + reducto_models: frozenset[str] + runwayml_models: frozenset[str] + sambanova_embedding_models: frozenset[str] + sambanova_models: frozenset[str] + snowflake_models: frozenset[str] + soniox_models: frozenset[str] + stability_models: frozenset[str] + tencent_models: frozenset[str] + text_completion_codestral_models: frozenset[str] + text_completion_inception_models: frozenset[str] + v0_models: frozenset[str] + vercel_ai_gateway_models: frozenset[str] + vertex_ai_ai21_models: frozenset[str] + vertex_ai_image_models: frozenset[str] + vertex_ai_video_models: frozenset[str] + vertex_anthropic_models: frozenset[str] + vertex_chat_models: frozenset[str] + vertex_code_chat_models: frozenset[str] + vertex_code_text_models: frozenset[str] + vertex_deepseek_models: frozenset[str] + vertex_embedding_models: frozenset[str] + vertex_language_models: frozenset[str] + vertex_llama3_models: frozenset[str] + vertex_minimax_models: frozenset[str] + vertex_mistral_models: frozenset[str] + vertex_moonshot_models: frozenset[str] + vertex_openai_models: frozenset[str] + vertex_text_models: frozenset[str] + vertex_vision_models: frozenset[str] + vertex_zai_models: frozenset[str] + volcengine_models: frozenset[str] + voyage_models: frozenset[str] + wandb_models: frozenset[str] + watsonx_models: frozenset[str] + xai_models: frozenset[str] + zai_models: frozenset[str] +# known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary + +# this is maintained for Exception Mapping # provider_list is lazy-loaded via __getattr__ to avoid importing LlmProviders at import time - -models_by_provider: dict = { - "openai": open_ai_chat_completion_models | open_ai_text_completion_models, - "text-completion-openai": open_ai_text_completion_models, - "cohere": cohere_models | cohere_chat_models, - "cohere_chat": cohere_chat_models, - "anthropic": anthropic_models, - "replicate": replicate_models, - "huggingface": huggingface_models, - "together_ai": together_ai_models, - "baseten": baseten_models, - "openrouter": openrouter_models, - "vercel_ai_gateway": vercel_ai_gateway_models, - "datarobot": datarobot_models, - "vertex_ai": vertex_chat_models - | vertex_text_models - | vertex_anthropic_models - | vertex_vision_models - | vertex_language_models - | vertex_deepseek_models - | vertex_minimax_models - | vertex_moonshot_models - | vertex_zai_models, - "ai21": ai21_models, - "bedrock": bedrock_models | bedrock_converse_models, - "petals": petals_models, - "ollama": ollama_models, - "ollama_chat": ollama_models, - "deepinfra": deepinfra_models, - "perplexity": perplexity_models, - "maritalk": maritalk_models, - "watsonx": watsonx_models, - "gemini": gemini_models, - "fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models, - "aleph_alpha": aleph_alpha_models, - "text-completion-codestral": text_completion_codestral_models, - "text-completion-inception": text_completion_inception_models, - "xai": xai_models, - "zai": zai_models, - "fal_ai": fal_ai_models, - "deepseek": deepseek_models, - "tencent": tencent_models, - "runwayml": runwayml_models, - "mistral": mistral_chat_models, - "azure_ai": azure_ai_models, - "voyage": voyage_models, - "infinity": infinity_models, - "databricks": databricks_models, - "cloudflare": cloudflare_models, - "codestral": codestral_models, - "nlp_cloud": nlp_cloud_models, - "friendliai": friendliai_models, - "palm": palm_models, - "groq": groq_models, - "azure": azure_models | azure_text_models, - "azure_anthropic": azure_anthropic_models, - "azure_text": azure_text_models, - "anyscale": anyscale_models, - "cerebras": cerebras_models, - "galadriel": galadriel_models, - "nvidia_nim": nvidia_nim_models, - "nvidia_riva": nvidia_riva_models, - "soniox": soniox_models, - "sambanova": sambanova_models | sambanova_embedding_models, - "novita": novita_models, - "nebius": nebius_models | nebius_embedding_models, - "aiml": aiml_models, - "assemblyai": assemblyai_models, - "jina_ai": jina_ai_models, - "snowflake": snowflake_models, - "gradient_ai": gradient_ai_models, - "meta_llama": llama_models, - "nscale": nscale_models, - "featherless_ai": featherless_ai_models, - "deepgram": deepgram_models, - "elevenlabs": elevenlabs_models, - "heroku": heroku_models, - "dashscope": dashscope_models, - "modelscope": modelscope_models, - "moonshot": moonshot_models, - "publicai": publicai_models, - "darkbloom": darkbloom_models, - "v0": v0_models, - "morph": morph_models, - "lambda_ai": lambda_ai_models, - "inception": inception_models, - "hyperbolic": hyperbolic_models, - "black_forest_labs": black_forest_labs_models, - "recraft": recraft_models, - "cometapi": cometapi_models, - "oci": oci_models, - "volcengine": volcengine_models, - "wandb": wandb_models, - "ovhcloud": ovhcloud_models | ovhcloud_embedding_models, - "lemonade": lemonade_models, - "clarifai": clarifai_models, - "amazon_nova": amazon_nova_models, - "stability": stability_models, - "github_copilot": github_copilot_models, - "chatgpt": chatgpt_models, - "minimax": minimax_models, - "aws_polly": aws_polly_models, - "gigachat": gigachat_models, - "llamagate": llamagate_models, - "reducto": reducto_models, - "bedrock_mantle": bedrock_mantle_models, -} - # mapping for those models which have larger equivalents longer_context_model_fallback_dict: dict = { # openai chat completion models @@ -1199,19 +792,6 @@ longer_context_model_fallback_dict: dict = { "openrouter/anthropic/claude-instant-v1": "openrouter/anthropic/claude-2", } -####### EMBEDDING MODELS ################### - -all_embedding_models = ( - open_ai_embedding_models - | set(cohere_embedding_models) - | set(bedrock_embedding_models) - | vertex_embedding_models - | fireworks_ai_embedding_models - | nebius_embedding_models - | sambanova_embedding_models - | ovhcloud_embedding_models -) - ####### IMAGE GENERATION MODELS ################### openai_image_generation_models = ["dall-e-2", "dall-e-3"] @@ -2128,6 +1708,15 @@ if os.getenv("LITELLM_DISABLE_LAZY_LOADING", "").lower() in ("1", "true", "yes", def __getattr__(name: str) -> Any: """Lazy import handler with cached registry for improved performance.""" global _async_client_cleanup_registered + + _snapshot = _model_registry_snapshot + _legacy_set = _snapshot.legacy_sets.get(name) + if _legacy_set is not None: + return _legacy_set + _registry_view = _REGISTRY_VIEWS.get(name) + if _registry_view is not None: + return _registry_view(_snapshot) + # Register async client cleanup on first access (only once) if not _async_client_cleanup_registered: from litellm.llms.custom_httpx.async_client_cleanup import ( @@ -2323,4 +1912,9 @@ def __getattr__(name: str) -> Any: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") +def __dir__() -> Sequence[str]: + """Keep registry-served names discoverable even though they never enter the namespace.""" + return sorted({*globals(), *_model_registry_snapshot.legacy_sets, *_REGISTRY_VIEWS}) + + # ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time diff --git a/litellm/constants.py b/litellm/constants.py index 164f5a77a76..6ae9b3ba653 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -846,7 +846,7 @@ _openai_like_providers: list = [ "watsonx", ] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk # well supported replicate llms -replicate_models: set = set( +replicate_models: frozenset[str] = frozenset( [ # llama replicate supported LLMs "replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf", @@ -863,7 +863,7 @@ replicate_models: set = set( ] ) -clarifai_models: set = set( +clarifai_models: frozenset[str] = frozenset( [ "clarifai/openai.chat-completion.gpt-oss-20b", "clarifai/qwen.qwenLM.Qwen3-30B-A3B-Instruct-2507", @@ -899,7 +899,7 @@ clarifai_models: set = set( ) -huggingface_models: set = set( +huggingface_models: frozenset[str] = frozenset( [ "meta-llama/Llama-2-7b-hf", "meta-llama/Llama-2-7b-chat-hf", @@ -915,14 +915,14 @@ huggingface_models: set = set( "meta-llama/Llama-2-70b-chat", ] ) # these have been tested on extensively. But by default all text2text-generation and text-generation models are supported by liteLLM. - https://docs.litellm.ai/docs/providers -empower_models = set( +empower_models: frozenset[str] = frozenset( [ "empower/empower-functions", "empower/empower-functions-small", ] ) -together_ai_models: set = set( +together_ai_models: frozenset[str] = frozenset( [ # llama llms - chat "togethercomputer/llama-2-70b-chat", @@ -956,7 +956,7 @@ together_ai_models: set = set( # supports all together ai models, just pass in the model id e.g. completion(model="together_computer/replit_code_3b",...) -baseten_models: set = set( +baseten_models: frozenset[str] = frozenset( [ "qvv0xeq", "q841o8w", @@ -1055,7 +1055,7 @@ nebius_embedding_models: set = set( ] ) -WANDB_MODELS: set = set( +WANDB_MODELS: frozenset[str] = frozenset( [ # openai models "openai/gpt-oss-120b", @@ -1084,7 +1084,7 @@ WANDB_MODELS: set = set( ] ) -modelscope_models: set = set( +modelscope_models: frozenset[str] = frozenset( [ # Qwen series models "Qwen/Qwen3-0.6B", @@ -1151,7 +1151,7 @@ BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[ "nova", ] -BEDROCK_CONVERSE_MODELS = [ +BEDROCK_CONVERSE_MODELS: tuple[str, ...] = ( "qwen.qwen3-coder-480b-a35b-v1:0", "qwen.qwen3-coder-next", "qwen.qwen3-235b-a22b-2507-v1:0", @@ -1209,11 +1209,11 @@ BEDROCK_CONVERSE_MODELS = [ "writer.palmyra-x5-v1:0", "minimax.minimax-m2.1", "moonshotai.kimi-k2.5", -] +) -open_ai_embedding_models: set = set(["text-embedding-ada-002"]) -cohere_embedding_models: set = set( +open_ai_embedding_models: frozenset[str] = frozenset(["text-embedding-ada-002"]) +cohere_embedding_models: frozenset[str] = frozenset( [ "embed-v4.0", "embed-english-v3.0", @@ -1224,7 +1224,7 @@ cohere_embedding_models: set = set( "embed-multilingual-v2.0", ] ) -bedrock_embedding_models: set = set( +bedrock_embedding_models: frozenset[str] = frozenset( [ "amazon.titan-embed-text-v1", "amazon.nova-2-multimodal-embeddings-v1:0", diff --git a/litellm/litellm_core_utils/model_registry.py b/litellm/litellm_core_utils/model_registry.py new file mode 100644 index 00000000000..84da79355bf --- /dev/null +++ b/litellm/litellm_core_utils/model_registry.py @@ -0,0 +1,635 @@ +"""Immutable provider/model registry derived from the litellm cost map. + +Every legacy ``litellm._models`` collection, ``models_by_provider``, +``model_list`` and ``all_embedding_models`` is a view over one snapshot built +here, so a cost-map reload can never leave them stale relative to each other. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from itertools import groupby +from typing import assert_never + +from litellm.constants import ( + BEDROCK_CONVERSE_MODELS, + WANDB_MODELS, + baseten_models, + bedrock_embedding_models, + clarifai_models, + cohere_embedding_models, + empower_models, + huggingface_models, + modelscope_models, + open_ai_embedding_models, + replicate_models, + together_ai_models, +) + +_BEDROCK_PRICING_ONLY_PATTERN = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$") +_VERTEX_KEY_PREFIX = "vertex_ai/" + + +def is_bedrock_pricing_only_model(key: str) -> bool: + """Keys shaped ``bedrock//`` exist for pricing only, not for routing.""" + if "month-commitment" in key: + return True + return _BEDROCK_PRICING_ONLY_PATTERN.match(key) is not None + + +def is_openai_finetune_model(key: str) -> bool: + """Keys shaped ``ft:`` exist for pricing only, not for routing.""" + return key.startswith("ft:") and not key.count(":") > 1 + + +def _is_enumerable_openai_model(key: str) -> bool: + return not is_openai_finetune_model(key) + + +def _is_enumerable_bedrock_model(key: str) -> bool: + return not is_bedrock_pricing_only_model(key) + + +def _is_enumerable_fireworks_model(key: str) -> bool: + return "-to-" not in key and "fireworks-ai-default" not in key + + +def _is_enumerable_fireworks_embedding_model(key: str) -> bool: + return "-to-" not in key + + +@dataclass(frozen=True, slots=True) +class _LegacyTarget: + set_name: str + strip_vertex_prefix: bool = False + key_filter: Callable[[str], bool] | None = None + + +@dataclass(frozen=True, slots=True) +class _ModeSplitTarget: + chat_set_name: str + default_set_name: str + + +_ProviderRule = _LegacyTarget | _ModeSplitTarget + + +@dataclass(frozen=True, slots=True) +class _LegacyMember: + set_name: str + model: str + + +@dataclass(frozen=True, slots=True) +class _ProviderMember: + provider: str + model: str + + +_DerivedMember = _LegacyMember | _ProviderMember + + +@dataclass(frozen=True, slots=True) +class ModelRegistrySnapshot: + """One atomically-swappable view of everything derived from the cost map.""" + + legacy_sets: Mapping[str, frozenset[str]] + models_by_provider: Mapping[str, frozenset[str]] + model_list: tuple[str, ...] + model_list_set: frozenset[str] + all_embedding_models: frozenset[str] + + +_PROVIDER_RULES: Mapping[str, _ProviderRule] = { + "openai": _LegacyTarget(set_name="open_ai_chat_completion_models", key_filter=_is_enumerable_openai_model), + "text-completion-openai": _LegacyTarget(set_name="open_ai_text_completion_models"), + "azure_text": _LegacyTarget(set_name="azure_text_models"), + "cohere": _LegacyTarget(set_name="cohere_models"), + "cohere_chat": _LegacyTarget(set_name="cohere_chat_models"), + "mistral": _LegacyTarget(set_name="mistral_chat_models"), + "anthropic": _LegacyTarget(set_name="anthropic_models"), + "empower": _LegacyTarget(set_name="empower_models"), + "openrouter": _LegacyTarget(set_name="openrouter_models"), + "vercel_ai_gateway": _LegacyTarget(set_name="vercel_ai_gateway_models"), + "datarobot": _LegacyTarget(set_name="datarobot_models"), + "vertex_ai-text-models": _LegacyTarget(set_name="vertex_text_models"), + "vertex_ai-code-text-models": _LegacyTarget(set_name="vertex_code_text_models"), + "vertex_ai-language-models": _LegacyTarget(set_name="vertex_language_models"), + "vertex_ai-vision-models": _LegacyTarget(set_name="vertex_vision_models"), + "vertex_ai-chat-models": _LegacyTarget(set_name="vertex_chat_models"), + "vertex_ai-code-chat-models": _LegacyTarget(set_name="vertex_code_chat_models"), + "vertex_ai-embedding-models": _LegacyTarget(set_name="vertex_embedding_models"), + "vertex_ai-anthropic_models": _LegacyTarget(set_name="vertex_anthropic_models", strip_vertex_prefix=True), + "vertex_ai-llama_models": _LegacyTarget(set_name="vertex_llama3_models", strip_vertex_prefix=True), + "vertex_ai-deepseek_models": _LegacyTarget(set_name="vertex_deepseek_models", strip_vertex_prefix=True), + "vertex_ai-mistral_models": _LegacyTarget(set_name="vertex_mistral_models", strip_vertex_prefix=True), + "vertex_ai-ai21_models": _LegacyTarget(set_name="vertex_ai_ai21_models", strip_vertex_prefix=True), + "vertex_ai-image-models": _LegacyTarget(set_name="vertex_ai_image_models", strip_vertex_prefix=True), + "vertex_ai-video-models": _LegacyTarget(set_name="vertex_ai_video_models", strip_vertex_prefix=True), + "vertex_ai-openai_models": _LegacyTarget(set_name="vertex_openai_models", strip_vertex_prefix=True), + "vertex_ai-minimax_models": _LegacyTarget(set_name="vertex_minimax_models", strip_vertex_prefix=True), + "vertex_ai-moonshot_models": _LegacyTarget(set_name="vertex_moonshot_models", strip_vertex_prefix=True), + "vertex_ai-zai_models": _LegacyTarget(set_name="vertex_zai_models", strip_vertex_prefix=True), + "ai21": _ModeSplitTarget(chat_set_name="ai21_chat_models", default_set_name="ai21_models"), + "nlp_cloud": _LegacyTarget(set_name="nlp_cloud_models"), + "aleph_alpha": _LegacyTarget(set_name="aleph_alpha_models"), + "bedrock": _LegacyTarget(set_name="bedrock_models", key_filter=_is_enumerable_bedrock_model), + "bedrock_converse": _LegacyTarget(set_name="bedrock_converse_models"), + "deepinfra": _LegacyTarget(set_name="deepinfra_models"), + "perplexity": _LegacyTarget(set_name="perplexity_models"), + "watsonx": _LegacyTarget(set_name="watsonx_models"), + "gemini": _LegacyTarget(set_name="gemini_models"), + "fireworks_ai": _LegacyTarget(set_name="fireworks_ai_models", key_filter=_is_enumerable_fireworks_model), + "fireworks_ai-embedding-models": _LegacyTarget( + set_name="fireworks_ai_embedding_models", key_filter=_is_enumerable_fireworks_embedding_model + ), + "text-completion-codestral": _LegacyTarget(set_name="text_completion_codestral_models"), + "text-completion-inception": _LegacyTarget(set_name="text_completion_inception_models"), + "xai": _LegacyTarget(set_name="xai_models"), + "zai": _LegacyTarget(set_name="zai_models"), + "fal_ai": _LegacyTarget(set_name="fal_ai_models"), + "deepseek": _LegacyTarget(set_name="deepseek_models"), + "tencent": _LegacyTarget(set_name="tencent_models"), + "runwayml": _LegacyTarget(set_name="runwayml_models"), + "meta_llama": _LegacyTarget(set_name="llama_models"), + "nscale": _LegacyTarget(set_name="nscale_models"), + "azure_ai": _LegacyTarget(set_name="azure_ai_models"), + "voyage": _LegacyTarget(set_name="voyage_models"), + "infinity": _LegacyTarget(set_name="infinity_models"), + "databricks": _LegacyTarget(set_name="databricks_models"), + "cloudflare": _LegacyTarget(set_name="cloudflare_models"), + "codestral": _LegacyTarget(set_name="codestral_models"), + "friendliai": _LegacyTarget(set_name="friendliai_models"), + "palm": _LegacyTarget(set_name="palm_models"), + "groq": _LegacyTarget(set_name="groq_models"), + "azure": _LegacyTarget(set_name="azure_models"), + "azure_anthropic": _LegacyTarget(set_name="azure_anthropic_models"), + "anyscale": _LegacyTarget(set_name="anyscale_models"), + "cerebras": _LegacyTarget(set_name="cerebras_models"), + "galadriel": _LegacyTarget(set_name="galadriel_models"), + "nvidia_nim": _LegacyTarget(set_name="nvidia_nim_models"), + "nvidia_riva": _LegacyTarget(set_name="nvidia_riva_models"), + "soniox": _LegacyTarget(set_name="soniox_models"), + "sambanova": _LegacyTarget(set_name="sambanova_models"), + "sambanova-embedding-models": _LegacyTarget(set_name="sambanova_embedding_models"), + "novita": _LegacyTarget(set_name="novita_models"), + "nebius-chat-models": _LegacyTarget(set_name="nebius_models"), + "nebius-embedding-models": _LegacyTarget(set_name="nebius_embedding_models"), + "aiml": _LegacyTarget(set_name="aiml_models"), + "assemblyai": _LegacyTarget(set_name="assemblyai_models"), + "jina_ai": _LegacyTarget(set_name="jina_ai_models"), + "snowflake": _LegacyTarget(set_name="snowflake_models"), + "gradient_ai": _LegacyTarget(set_name="gradient_ai_models"), + "featherless_ai": _LegacyTarget(set_name="featherless_ai_models"), + "deepgram": _LegacyTarget(set_name="deepgram_models"), + "elevenlabs": _LegacyTarget(set_name="elevenlabs_models"), + "heroku": _LegacyTarget(set_name="heroku_models"), + "dashscope": _LegacyTarget(set_name="dashscope_models"), + "modelscope": _LegacyTarget(set_name="modelscope_models"), + "moonshot": _LegacyTarget(set_name="moonshot_models"), + "publicai": _LegacyTarget(set_name="publicai_models"), + "darkbloom": _LegacyTarget(set_name="darkbloom_models"), + "v0": _LegacyTarget(set_name="v0_models"), + "morph": _LegacyTarget(set_name="morph_models"), + "lambda_ai": _LegacyTarget(set_name="lambda_ai_models"), + "inception": _LegacyTarget(set_name="inception_models"), + "hyperbolic": _LegacyTarget(set_name="hyperbolic_models"), + "black_forest_labs": _LegacyTarget(set_name="black_forest_labs_models"), + "recraft": _LegacyTarget(set_name="recraft_models"), + "cometapi": _LegacyTarget(set_name="cometapi_models"), + "oci": _LegacyTarget(set_name="oci_models"), + "volcengine": _LegacyTarget(set_name="volcengine_models"), + "wandb": _LegacyTarget(set_name="wandb_models"), + "ovhcloud": _LegacyTarget(set_name="ovhcloud_models"), + "ovhcloud-embedding-models": _LegacyTarget(set_name="ovhcloud_embedding_models"), + "lemonade": _LegacyTarget(set_name="lemonade_models"), + "docker_model_runner": _LegacyTarget(set_name="docker_model_runner_models"), + "amazon_nova": _LegacyTarget(set_name="amazon_nova_models"), + "stability": _LegacyTarget(set_name="stability_models"), + "github_copilot": _LegacyTarget(set_name="github_copilot_models"), + "chatgpt": _LegacyTarget(set_name="chatgpt_models"), + "minimax": _LegacyTarget(set_name="minimax_models"), + "aws_polly": _LegacyTarget(set_name="aws_polly_models"), + "gigachat": _LegacyTarget(set_name="gigachat_models"), + "llamagate": _LegacyTarget(set_name="llamagate_models"), + "reducto": _LegacyTarget(set_name="reducto_models"), + "bedrock_mantle": _LegacyTarget(set_name="bedrock_mantle_models"), +} + +_PROVIDER_COMPOSITION: Mapping[str, tuple[str, ...]] = { + "openai": ("open_ai_chat_completion_models", "open_ai_text_completion_models"), + "text-completion-openai": ("open_ai_text_completion_models",), + "cohere": ("cohere_models", "cohere_chat_models"), + "cohere_chat": ("cohere_chat_models",), + "anthropic": ("anthropic_models",), + "replicate": ("replicate_models",), + "huggingface": ("huggingface_models",), + "together_ai": ("together_ai_models",), + "baseten": ("baseten_models",), + "openrouter": ("openrouter_models",), + "vercel_ai_gateway": ("vercel_ai_gateway_models",), + "datarobot": ("datarobot_models",), + "vertex_ai": ( + "vertex_chat_models", + "vertex_text_models", + "vertex_anthropic_models", + "vertex_vision_models", + "vertex_language_models", + "vertex_deepseek_models", + "vertex_minimax_models", + "vertex_moonshot_models", + "vertex_zai_models", + ), + "ai21": ("ai21_models",), + "bedrock": ("bedrock_models", "bedrock_converse_models"), + "petals": ("petals_models",), + "ollama": ("ollama_models",), + "ollama_chat": ("ollama_models",), + "deepinfra": ("deepinfra_models",), + "perplexity": ("perplexity_models",), + "maritalk": ("maritalk_models",), + "watsonx": ("watsonx_models",), + "gemini": ("gemini_models",), + "fireworks_ai": ("fireworks_ai_models", "fireworks_ai_embedding_models"), + "aleph_alpha": ("aleph_alpha_models",), + "text-completion-codestral": ("text_completion_codestral_models",), + "text-completion-inception": ("text_completion_inception_models",), + "xai": ("xai_models",), + "zai": ("zai_models",), + "fal_ai": ("fal_ai_models",), + "deepseek": ("deepseek_models",), + "tencent": ("tencent_models",), + "runwayml": ("runwayml_models",), + "mistral": ("mistral_chat_models",), + "azure_ai": ("azure_ai_models",), + "voyage": ("voyage_models",), + "infinity": ("infinity_models",), + "databricks": ("databricks_models",), + "cloudflare": ("cloudflare_models",), + "codestral": ("codestral_models",), + "nlp_cloud": ("nlp_cloud_models",), + "friendliai": ("friendliai_models",), + "palm": ("palm_models",), + "groq": ("groq_models",), + "azure": ("azure_models", "azure_text_models"), + "azure_anthropic": ("azure_anthropic_models",), + "azure_text": ("azure_text_models",), + "anyscale": ("anyscale_models",), + "cerebras": ("cerebras_models",), + "galadriel": ("galadriel_models",), + "nvidia_nim": ("nvidia_nim_models",), + "nvidia_riva": ("nvidia_riva_models",), + "soniox": ("soniox_models",), + "sambanova": ("sambanova_models", "sambanova_embedding_models"), + "novita": ("novita_models",), + "nebius": ("nebius_models", "nebius_embedding_models"), + "aiml": ("aiml_models",), + "assemblyai": ("assemblyai_models",), + "jina_ai": ("jina_ai_models",), + "snowflake": ("snowflake_models",), + "gradient_ai": ("gradient_ai_models",), + "meta_llama": ("llama_models",), + "nscale": ("nscale_models",), + "featherless_ai": ("featherless_ai_models",), + "deepgram": ("deepgram_models",), + "elevenlabs": ("elevenlabs_models",), + "heroku": ("heroku_models",), + "dashscope": ("dashscope_models",), + "modelscope": ("modelscope_models",), + "moonshot": ("moonshot_models",), + "publicai": ("publicai_models",), + "darkbloom": ("darkbloom_models",), + "v0": ("v0_models",), + "morph": ("morph_models",), + "lambda_ai": ("lambda_ai_models",), + "inception": ("inception_models",), + "hyperbolic": ("hyperbolic_models",), + "black_forest_labs": ("black_forest_labs_models",), + "recraft": ("recraft_models",), + "cometapi": ("cometapi_models",), + "oci": ("oci_models",), + "volcengine": ("volcengine_models",), + "wandb": ("wandb_models",), + "ovhcloud": ("ovhcloud_models", "ovhcloud_embedding_models"), + "lemonade": ("lemonade_models",), + "clarifai": ("clarifai_models",), + "amazon_nova": ("amazon_nova_models",), + "stability": ("stability_models",), + "github_copilot": ("github_copilot_models",), + "chatgpt": ("chatgpt_models",), + "minimax": ("minimax_models",), + "aws_polly": ("aws_polly_models",), + "gigachat": ("gigachat_models",), + "llamagate": ("llamagate_models",), + "reducto": ("reducto_models",), + "bedrock_mantle": ("bedrock_mantle_models",), +} + +_MODEL_LIST_SOURCES: tuple[str, ...] = ( + "open_ai_chat_completion_models", + "open_ai_text_completion_models", + "cohere_models", + "cohere_chat_models", + "anthropic_models", + "replicate_models", + "openrouter_models", + "datarobot_models", + "huggingface_models", + "vertex_chat_models", + "vertex_text_models", + "ai21_models", + "ai21_chat_models", + "together_ai_models", + "baseten_models", + "aleph_alpha_models", + "nlp_cloud_models", + "ollama_models", + "bedrock_models", + "deepinfra_models", + "perplexity_models", + "maritalk_models", + "runwayml_models", + "vertex_language_models", + "watsonx_models", + "gemini_models", + "text_completion_codestral_models", + "text_completion_inception_models", + "xai_models", + "zai_models", + "fal_ai_models", + "deepseek_models", + "modelscope_models", + "azure_ai_models", + "voyage_models", + "infinity_models", + "databricks_models", + "cloudflare_models", + "codestral_models", + "friendliai_models", + "palm_models", + "groq_models", + "azure_models", + "azure_anthropic_models", + "anyscale_models", + "cerebras_models", + "galadriel_models", + "nvidia_nim_models", + "nvidia_riva_models", + "soniox_models", + "sambanova_models", + "azure_text_models", + "novita_models", + "assemblyai_models", + "jina_ai_models", + "snowflake_models", + "gradient_ai_models", + "llama_models", + "featherless_ai_models", + "nscale_models", + "deepgram_models", + "elevenlabs_models", + "dashscope_models", + "moonshot_models", + "publicai_models", + "darkbloom_models", + "v0_models", + "morph_models", + "lambda_ai_models", + "inception_models", + "black_forest_labs_models", + "recraft_models", + "cometapi_models", + "oci_models", + "heroku_models", + "vercel_ai_gateway_models", + "volcengine_models", + "wandb_models", + "ovhcloud_models", + "lemonade_models", + "docker_model_runner_models", + "reducto_models", + "bedrock_mantle_models", + "clarifai_models", +) + +_EMBEDDING_SOURCES: tuple[str, ...] = ( + "open_ai_embedding_models", + "cohere_embedding_models", + "bedrock_embedding_models", + "vertex_embedding_models", + "fireworks_ai_embedding_models", + "nebius_embedding_models", + "sambanova_embedding_models", + "ovhcloud_embedding_models", +) + +_STATIC_SEEDS: Mapping[str, frozenset[str]] = { + "bedrock_converse_models": frozenset(BEDROCK_CONVERSE_MODELS), + "wandb_models": frozenset(WANDB_MODELS), + "empower_models": frozenset(empower_models), + "modelscope_models": frozenset(modelscope_models), +} + +_CONSTANT_SETS: Mapping[str, frozenset[str]] = { + "replicate_models": frozenset(replicate_models), + "clarifai_models": frozenset(clarifai_models), + "huggingface_models": frozenset(huggingface_models), + "together_ai_models": frozenset(together_ai_models), + "baseten_models": frozenset(baseten_models), + "open_ai_embedding_models": frozenset(open_ai_embedding_models), + "cohere_embedding_models": frozenset(cohere_embedding_models), + "bedrock_embedding_models": frozenset(bedrock_embedding_models), +} + + +def _rule_set_names(rule: _ProviderRule) -> tuple[str, ...]: + match rule: + case _LegacyTarget(set_name=set_name): + return (set_name,) + case _ModeSplitTarget(chat_set_name=chat_set_name, default_set_name=default_set_name): + return (chat_set_name, default_set_name) + case _: + assert_never(rule) + + +REGISTRY_SET_NAMES: frozenset[str] = frozenset( + name for rule in _PROVIDER_RULES.values() for name in _rule_set_names(rule) +) + +_OPENROUTER_SET_NAME = _rule_set_names(_PROVIDER_RULES["openrouter"])[0] + +_PROVIDERS_BY_SET: Mapping[str, tuple[str, ...]] = { + set_name: tuple(provider for provider, sources in _PROVIDER_COMPOSITION.items() if set_name in sources) + for set_name in {source for sources in _PROVIDER_COMPOSITION.values() for source in sources} +} + + +def _first(pair: tuple[str, str]) -> str: + return pair[0] + + +def _prefix_roots(provider: str) -> tuple[str, ...]: + return tuple(provider[:index] for index, char in enumerate(provider) if char == "-") + + +def _fallback_provider(provider: str, known_providers: frozenset[str]) -> str | None: + if provider in known_providers: + return provider + return max((root for root in _prefix_roots(provider) if root in known_providers), key=len, default=None) + + +def _classify( + key: str, + value: Mapping[str, object], + known_providers: frozenset[str], +) -> _DerivedMember | None: + """Place one cost-map entry, or drop it. + + Providers with no table rule fall back to provider-bucket-only placement, and only for + prefixed keys: unprefixed keys under such a provider are pricing tiers (``together-ai-up-to-4b``), + not callable models. + """ + provider = value.get("litellm_provider") + if not isinstance(provider, str): + return None + rule = _PROVIDER_RULES.get(provider) + if rule is None: + if "/" not in key: + return None + bucket = _fallback_provider(provider, known_providers) + return None if bucket is None else _ProviderMember(provider=bucket, model=key) + match rule: + case _LegacyTarget(set_name=set_name, strip_vertex_prefix=strip_vertex_prefix, key_filter=key_filter): + if key_filter is not None and not key_filter(key): + return None + model = key.replace(_VERTEX_KEY_PREFIX, "") if strip_vertex_prefix else key + return _LegacyMember(set_name=set_name, model=model) + case _ModeSplitTarget(chat_set_name=chat_set_name, default_set_name=default_set_name): + return _LegacyMember( + set_name=chat_set_name if value.get("mode") == "chat" else default_set_name, + model=key, + ) + case _: + assert_never(rule) + + +def _derive( + model_cost: Mapping[str, Mapping[str, object]], + known_providers: frozenset[str], +) -> tuple[_DerivedMember, ...]: + return tuple( + member for key, value in model_cost.items() if (member := _classify(key, value, known_providers)) is not None + ) + + +def _openrouter_register_aliases( + model_cost_additions: Mapping[str, Mapping[str, object]], +) -> tuple[_DerivedMember, ...]: + """Prefix-stripped twins for openrouter models registered at runtime. + + ``register_model`` has always stored the unprefixed name as well, so ``get_llm_provider`` can + infer openrouter from a bare model name. The shipped cost map carries only verbatim keys, so + this stays out of ``build_snapshot`` and cannot move import-time membership. + """ + return tuple( + _LegacyMember(set_name=_OPENROUTER_SET_NAME, model=key.split("/", 1)[-1]) + for key, value in model_cost_additions.items() + if value.get("litellm_provider") == "openrouter" and "/" in key + ) + + +def _grouped(pairs: Iterable[tuple[str, str]]) -> Mapping[str, frozenset[str]]: + return {name: frozenset(model for _, model in group) for name, group in groupby(sorted(pairs), key=_first)} + + +def _legacy_groups(members: tuple[_DerivedMember, ...]) -> Mapping[str, frozenset[str]]: + return _grouped((member.set_name, member.model) for member in members if isinstance(member, _LegacyMember)) + + +def _provider_groups(members: tuple[_DerivedMember, ...]) -> Mapping[str, frozenset[str]]: + return _grouped((member.provider, member.model) for member in members if isinstance(member, _ProviderMember)) + + +def _union(source_names: Iterable[str], resolved: Mapping[str, frozenset[str]]) -> frozenset[str]: + return frozenset(model for name in source_names for model in resolved.get(name, frozenset())) + + +def _assemble( + legacy_sets: Mapping[str, frozenset[str]], + provider_extras: Mapping[str, frozenset[str]], + static_model_names: Mapping[str, frozenset[str]], +) -> ModelRegistrySnapshot: + resolved = {**_CONSTANT_SETS, **static_model_names, **legacy_sets} + composed = { + provider: _union(sources, resolved) | provider_extras.get(provider, frozenset()) + for provider, sources in _PROVIDER_COMPOSITION.items() + } + extra_only = { + provider: models for provider, models in provider_extras.items() if provider not in _PROVIDER_COMPOSITION + } + model_list_set = _union(_MODEL_LIST_SOURCES, resolved) + return ModelRegistrySnapshot( + legacy_sets=legacy_sets, + models_by_provider={**composed, **extra_only}, + model_list=tuple(sorted(model_list_set)), + model_list_set=model_list_set, + all_embedding_models=_union(_EMBEDDING_SOURCES, resolved), + ) + + +def build_snapshot( + model_cost: Mapping[str, Mapping[str, object]], + known_providers: frozenset[str], + static_model_names: Mapping[str, frozenset[str]], +) -> ModelRegistrySnapshot: + """Derive a complete snapshot from ``model_cost``, dropping models it no longer contains.""" + members = _derive(model_cost, known_providers) + derived = _legacy_groups(members) + legacy_sets = { + name: _STATIC_SEEDS.get(name, frozenset()) | derived.get(name, frozenset()) for name in REGISTRY_SET_NAMES + } + return _assemble(legacy_sets, _provider_groups(members), static_model_names) + + +def extend_snapshot( + snapshot: ModelRegistrySnapshot, + model_cost_additions: Mapping[str, Mapping[str, object]], + known_providers: frozenset[str], +) -> ModelRegistrySnapshot: + """Add ``model_cost_additions`` to ``snapshot`` using the same rules as ``build_snapshot``. + + Additive counterpart of a full rebuild, for callers that register a handful of models at a + time (``litellm.register_model``) and cannot pay for a full re-derivation per call. + """ + members = _derive(model_cost_additions, known_providers) + _openrouter_register_aliases(model_cost_additions) + if not members: + return snapshot + added = _legacy_groups(members) + provider_extras = _provider_groups(members) + legacy_sets = { + **snapshot.legacy_sets, + **{name: snapshot.legacy_sets.get(name, frozenset()) | models for name, models in added.items()}, + } + touched_providers = frozenset( + provider for name in added for provider in _PROVIDERS_BY_SET.get(name, ()) + ) | frozenset(provider_extras) + models_by_provider = { + **snapshot.models_by_provider, + **{ + provider: snapshot.models_by_provider.get(provider, frozenset()) + | _union(_PROVIDER_COMPOSITION.get(provider, ()), added) + | provider_extras.get(provider, frozenset()) + for provider in touched_providers + }, + } + listed = _union(_MODEL_LIST_SOURCES, added) - snapshot.model_list_set + model_list_set = snapshot.model_list_set | listed + return ModelRegistrySnapshot( + legacy_sets=legacy_sets, + models_by_provider=models_by_provider, + model_list=tuple(sorted(model_list_set)) if listed else snapshot.model_list, + model_list_set=model_list_set, + all_embedding_models=snapshot.all_embedding_models | _union(_EMBEDDING_SOURCES, added), + ) diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index ecc56e6f110..ffb5f4cd670 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -125,7 +125,7 @@ class OllamaModelInfo(BaseLLMModelInfo): from litellm import models_by_provider static = models_by_provider.get("ollama", []) or [] - return [f"ollama/{m}" for m in static] + return sorted({name if name.startswith("ollama/") else f"ollama/{name}" for name in static}) except Exception as e1: verbose_logger.warning("Error retrieving static ollama models as fallback: %s", e1) return [] diff --git a/litellm/utils.py b/litellm/utils.py index 5dbff5070aa..0588f7b2253 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -44,7 +44,7 @@ from httpx import Proxy from httpx._utils import get_environment_proxies from openai.lib import _parsing, _pydantic from openai.types.chat.completion_create_params import ResponseFormat -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from tiktoken import Encoding from tokenizers import Tokenizer @@ -2680,6 +2680,9 @@ def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None: return None +_registered_model_cost_adapter: TypeAdapter[dict[str, dict[str, object]]] = TypeAdapter(dict[str, dict[str, object]]) + + def register_model(model_cost: str | dict): """ Register new / Override existing models (and their pricing) to specific providers. @@ -2770,53 +2773,8 @@ def register_model(model_cost: str | dict): _invalidate_model_cost_lowercase_map() verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key) - # add new model names to provider lists - if value.get("litellm_provider") == "openai": - if key not in litellm.open_ai_chat_completion_models: - litellm.open_ai_chat_completion_models.add(key) - elif value.get("litellm_provider") == "text-completion-openai": - if key not in litellm.open_ai_text_completion_models: - litellm.open_ai_text_completion_models.add(key) - elif value.get("litellm_provider") == "cohere": - if key not in litellm.cohere_models: - litellm.cohere_models.add(key) - elif value.get("litellm_provider") == "anthropic": - if key not in litellm.anthropic_models: - litellm.anthropic_models.add(key) - elif value.get("litellm_provider") == "openrouter": - split_string = key.split("/", 1) - if split_string[-1] not in litellm.openrouter_models: - litellm.openrouter_models.add(split_string[-1]) - elif value.get("litellm_provider") == "vercel_ai_gateway": - if key not in litellm.vercel_ai_gateway_models: - litellm.vercel_ai_gateway_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-text-models": - if key not in litellm.vertex_text_models: - litellm.vertex_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-text-models": - if key not in litellm.vertex_code_text_models: - litellm.vertex_code_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-chat-models": - if key not in litellm.vertex_chat_models: - litellm.vertex_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-chat-models": - if key not in litellm.vertex_code_chat_models: - litellm.vertex_code_chat_models.add(key) - elif value.get("litellm_provider") == "ai21": - if key not in litellm.ai21_models: - litellm.ai21_models.add(key) - elif value.get("litellm_provider") == "nlp_cloud": - if key not in litellm.nlp_cloud_models: - litellm.nlp_cloud_models.add(key) - elif value.get("litellm_provider") == "aleph_alpha": - if key not in litellm.aleph_alpha_models: - litellm.aleph_alpha_models.add(key) - elif value.get("litellm_provider") == "bedrock": - if key not in litellm.bedrock_models: - litellm.bedrock_models.add(key) - elif value.get("litellm_provider") == "novita": - if key not in litellm.novita_models: - litellm.novita_models.add(key) + + litellm.extend_known_models(_registered_model_cost_adapter.validate_python(loaded_model_cost)) return model_cost @@ -7070,7 +7028,7 @@ def get_valid_models( ) ) else: - models_for_provider = copy.deepcopy(litellm.models_by_provider.get(provider, [])) + models_for_provider = list(litellm.models_by_provider.get(provider, [])) valid_models.extend(models_for_provider) return valid_models diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 96428267a45..99d1640ab44 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -69,7 +69,7 @@ "limit": 4 }, "C405": { - "limit": 21 + "limit": 10 }, "C408": { "limit": 14 @@ -81,7 +81,7 @@ "limit": 4 }, "C901": { - "limit": 311 + "limit": 310 }, "D419": { "limit": 9 @@ -264,7 +264,7 @@ "limit": 61 }, "SIM102": { - "limit": 324 + "limit": 323 }, "SIM103": { "limit": 129 diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py index 7ae18828d3f..ddaf2982632 100644 --- a/tests/llm_translation/test_lambda_ai.py +++ b/tests/llm_translation/test_lambda_ai.py @@ -12,6 +12,16 @@ from litellm import completion from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig +@pytest.fixture +def local_cost_map_registry(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.add_known_models() + yield + monkeypatch.undo() + litellm.add_known_models() + + def test_lambda_ai_config_initialization(): """Test LambdaAIChatConfig initializes correctly""" config = LambdaAIChatConfig() @@ -103,18 +113,10 @@ async def test_lambda_ai_completion_call(): raise -def test_lambda_ai_models_configuration(): +def test_lambda_ai_models_configuration(local_cost_map_registry): """Test that Lambda AI models are configured correctly""" from litellm import get_model_info - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # Clear and repopulate lambda_ai_models list after reloading model_cost - litellm.lambda_ai_models = set() - litellm.add_known_models() - # Some Lambda AI models to test lambda_ai_models = [ "lambda_ai/deepseek-llama3.3-70b", @@ -145,16 +147,8 @@ def test_lambda_ai_models_configuration(): ), f"{model} should support vision" -def test_lambda_ai_model_list_populated(): +def test_lambda_ai_model_list_populated(local_cost_map_registry): """Test that lambda_ai_models list is populated correctly""" - # Ensure we're using local model cost map and repopulate models - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # Clear and repopulate all model lists after reloading model_cost - litellm.lambda_ai_models = set() - litellm.add_known_models() - # This should be populated by the add_known_models function assert ( len(litellm.lambda_ai_models) > 0 diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 9ebdb4b7e97..d16cb30de89 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -346,7 +346,6 @@ def test_azure_ai_mistral_optional_params(): def test_vertex_ai_llama_3_optional_params(): - litellm.vertex_llama3_models = ["meta/llama3-405b-instruct-maas"] litellm.drop_params = True optional_params = get_optional_params( model="meta/llama3-405b-instruct-maas", @@ -359,7 +358,7 @@ def test_vertex_ai_llama_3_optional_params(): def test_vertex_ai_mistral_optional_params(): - litellm.vertex_mistral_models = ["mistral-large@2407"] + assert "mistral-large@2407" in litellm.vertex_mistral_models litellm.drop_params = True optional_params = get_optional_params( model="mistral-large@2407", diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index cf0c645615d..9b036f8ad66 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -22,7 +22,6 @@ from litellm import ( cost_per_token, get_max_tokens, model_cost, - open_ai_chat_completion_models, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler import json diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index e58e6c9694b..17006e56ec0 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -305,6 +305,17 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w print(e) +@pytest.fixture +def reload_cost_map(monkeypatch): + def _reload(new_cost_map: dict) -> None: + monkeypatch.setattr(litellm, "model_cost", new_cost_map) + litellm.add_known_models(model_cost_map=new_cost_map) + + yield _reload + monkeypatch.undo() + litellm.add_known_models() + + @pytest.mark.parametrize( "key_models, model, expect_to_work", [ @@ -316,7 +327,7 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w ], ) @pytest.mark.asyncio -async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_to_work): +async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_to_work, reload_cost_map): """ Regression test: after a cost-map hot-reload, calling add_known_models(model_cost_map=new_map) must update litellm.anthropic_models @@ -334,22 +345,21 @@ async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_t # Build a new cost map that includes the brand-new model — exactly what # proxy_server.py receives from get_model_cost_map() during a reload. - new_cost_map = dict(litellm.model_cost) - new_cost_map[model] = { - "litellm_provider": "anthropic", - "max_tokens": 8192, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, + new_cost_map = { + **litellm.model_cost, + model: { + "litellm_provider": "anthropic", + "max_tokens": 8192, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + }, } - original_model_cost = litellm.model_cost - litellm.model_cost = new_cost_map - # Confirm the model is NOT yet in the provider set before reload propagation. assert model not in litellm.anthropic_models # Simulate what proxy_server.py now does after every reload. - litellm.add_known_models(model_cost_map=new_cost_map) + reload_cost_map(new_cost_map) # After add_known_models(), the model must be in the set. assert model in litellm.anthropic_models @@ -369,29 +379,25 @@ async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_t router = litellm.Router(model_list=llm_model_list) user_api_key_object = UserAPIKeyAuth(models=key_models) - try: - if expect_to_work: + if expect_to_work: + await can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=user_api_key_object, + llm_router=router, + ) + else: + with pytest.raises(Exception): await can_key_call_model( model=model, llm_model_list=llm_model_list, valid_token=user_api_key_object, llm_router=router, ) - else: - with pytest.raises(Exception): - await can_key_call_model( - model=model, - llm_model_list=llm_model_list, - valid_token=user_api_key_object, - llm_router=router, - ) - finally: - litellm.model_cost = original_model_cost - litellm.anthropic_models.discard(model) @pytest.mark.asyncio -async def test_add_known_models_explicit_map_updates_provider_sets(): +async def test_add_known_models_explicit_map_updates_provider_sets(reload_cost_map): """ Regression test: after a cost-map hot-reload, calling add_known_models(model_cost_map=new_map) with the new map passed explicitly @@ -407,28 +413,23 @@ async def test_add_known_models_explicit_map_updates_provider_sets(): # Baseline: the model must not be in the sets before we do anything. assert fake_new_model not in litellm.anthropic_models - new_cost_map = dict(litellm.model_cost) - new_cost_map[fake_new_model] = { - "litellm_provider": "anthropic", - "max_tokens": 8192, - "input_cost_per_token": 0.000003, - "output_cost_per_token": 0.000015, + new_cost_map = { + **litellm.model_cost, + fake_new_model: { + "litellm_provider": "anthropic", + "max_tokens": 8192, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + }, } # Simulate what proxy_server.py does on reload. - original_model_cost = litellm.model_cost - litellm.model_cost = new_cost_map - litellm.add_known_models(model_cost_map=new_cost_map) + reload_cost_map(new_cost_map) - try: - assert fake_new_model in litellm.anthropic_models, ( - "add_known_models(model_cost_map=...) did not add the new model to " - "litellm.anthropic_models — wildcard access checks would fail." - ) - finally: - # Clean up: restore original state. - litellm.model_cost = original_model_cost - litellm.anthropic_models.discard(fake_new_model) + assert fake_new_model in litellm.anthropic_models, ( + "add_known_models(model_cost_map=...) did not add the new model to " + "litellm.anthropic_models — wildcard access checks would fail." + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index f4aa1926d21..90a922c0273 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -20,6 +20,7 @@ import asyncio import litellm from litellm._logging import ALL_LOGGERS +from litellm.litellm_core_utils.model_registry import REGISTRY_SET_NAMES from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, ) @@ -28,6 +29,8 @@ from litellm.llms.custom_httpx.async_client_cleanup import ( ) from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module +_REGISTRY_SERVED_NAMES = REGISTRY_SET_NAMES | frozenset(litellm._REGISTRY_VIEWS) + def _reset_module_level_aws_auth_caches(): """ @@ -230,8 +233,6 @@ def isolate_litellm_state(): "cost_discount_config", "disable_hf_tokenizer_download", "disable_copilot_system_to_assistant", - "cohere_models", - "anthropic_models", "token_counter", "initialized_langfuse_clients", ): @@ -253,6 +254,7 @@ def isolate_litellm_state(): # Store singleton registries that are lazily initialized during tests and # can change endpoint behavior later in the suite. original_tool_policy_registry = tool_registry_writer_module._tool_policy_registry + original_registry_snapshot = litellm._model_registry_snapshot had_module_level_client = "module_level_client" in litellm.__dict__ had_module_level_aclient = "module_level_aclient" in litellm.__dict__ original_module_level_client = litellm.__dict__.get("module_level_client") @@ -301,9 +303,16 @@ def isolate_litellm_state(): # Restore all callback lists to original state for attr_name, original_value in original_state.items(): + if attr_name in _REGISTRY_SERVED_NAMES: + continue if hasattr(litellm, attr_name): setattr(litellm, attr_name, original_value) + for _shadowed in _REGISTRY_SERVED_NAMES & litellm.__dict__.keys(): + del litellm.__dict__[_shadowed] + if litellm._model_registry_snapshot is not original_registry_snapshot: + litellm.add_known_models() + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/litellm_core_utils/test_model_registry.py b/tests/test_litellm/litellm_core_utils/test_model_registry.py new file mode 100644 index 00000000000..a10d2f1d50a --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_model_registry.py @@ -0,0 +1,344 @@ +from collections.abc import Iterable + +import pytest + +from litellm import constants +from litellm.litellm_core_utils.model_registry import ( + _PROVIDER_COMPOSITION, + ModelRegistrySnapshot, + build_snapshot, +) +from litellm.types.utils import LlmProviders + +KNOWN_PROVIDERS = frozenset(provider.value for provider in LlmProviders) + + +def snapshot( + model_cost: dict[str, dict[str, object]], + static_model_names: dict[str, frozenset[str]] | None = None, +) -> ModelRegistrySnapshot: + return build_snapshot( + model_cost=model_cost, + known_providers=KNOWN_PROVIDERS, + static_model_names=static_model_names or {}, + ) + + +def entry(provider: str, mode: str = "chat") -> dict[str, object]: + return {"litellm_provider": provider, "mode": mode} + + +@pytest.mark.parametrize( + "provider, key, expected_set, expected_member", + [ + ("vertex_ai-anthropic_models", "vertex_ai/claude-sonnet-4-5", "vertex_anthropic_models", "claude-sonnet-4-5"), + ("vertex_ai-llama_models", "vertex_ai/llama-3.1-8b", "vertex_llama3_models", "llama-3.1-8b"), + ("vertex_ai-deepseek_models", "vertex_ai/deepseek-v3", "vertex_deepseek_models", "deepseek-v3"), + ("vertex_ai-mistral_models", "vertex_ai/mistral-large", "vertex_mistral_models", "mistral-large"), + ("vertex_ai-ai21_models", "vertex_ai/jamba-large", "vertex_ai_ai21_models", "jamba-large"), + ("vertex_ai-image-models", "vertex_ai/imagen-4", "vertex_ai_image_models", "imagen-4"), + ("vertex_ai-video-models", "vertex_ai/veo-3", "vertex_ai_video_models", "veo-3"), + ("vertex_ai-openai_models", "vertex_ai/gpt-oss-120b", "vertex_openai_models", "gpt-oss-120b"), + ("vertex_ai-minimax_models", "vertex_ai/minimax-m2", "vertex_minimax_models", "minimax-m2"), + ("vertex_ai-moonshot_models", "vertex_ai/kimi-k2", "vertex_moonshot_models", "kimi-k2"), + ("vertex_ai-zai_models", "vertex_ai/glm-4.6", "vertex_zai_models", "glm-4.6"), + ], +) +def test_vertex_families_strip_the_vertex_ai_key_prefix( + provider: str, key: str, expected_set: str, expected_member: str +) -> None: + snap = snapshot({key: entry(provider)}) + + assert snap.legacy_sets[expected_set] == frozenset({expected_member}) + + +@pytest.mark.parametrize( + "provider, expected_set", + [ + ("vertex_ai-language-models", "vertex_language_models"), + ("vertex_ai-chat-models", "vertex_chat_models"), + ("vertex_ai-text-models", "vertex_text_models"), + ("vertex_ai-embedding-models", "vertex_embedding_models"), + ], +) +def test_vertex_families_without_a_strip_rule_keep_the_whole_key(provider: str, expected_set: str) -> None: + snap = snapshot({"vertex_ai/gemini-3-pro": entry(provider)}) + + assert snap.legacy_sets[expected_set] == frozenset({"vertex_ai/gemini-3-pro"}) + + +def test_openai_finetune_keys_are_priced_but_not_enumerable() -> None: + snap = snapshot( + { + "gpt-5.2": entry("openai"), + "ft:gpt-5.2": entry("openai"), + "ft:gpt-5.2:acme::abc123": entry("openai"), + } + ) + + assert snap.legacy_sets["open_ai_chat_completion_models"] == frozenset({"gpt-5.2", "ft:gpt-5.2:acme::abc123"}) + + +def test_bedrock_pricing_only_keys_are_not_enumerable() -> None: + snap = snapshot( + { + "anthropic.claude-sonnet-4-5-v1:0": entry("bedrock"), + "bedrock/us-east-1/anthropic.claude-sonnet-4-5-v1:0": entry("bedrock"), + "anthropic.claude-sonnet-4-5-v1:0-month-commitment": entry("bedrock"), + } + ) + + assert snap.legacy_sets["bedrock_models"] == frozenset({"anthropic.claude-sonnet-4-5-v1:0"}) + + +def test_fireworks_pricing_tier_keys_are_not_enumerable() -> None: + snap = snapshot( + { + "accounts/fireworks/models/kimi-k2": entry("fireworks_ai"), + "fireworks-ai-up-to-16b": entry("fireworks_ai"), + "fireworks-ai-default": entry("fireworks_ai"), + "nomic-ai/nomic-embed-text-v1.5": entry("fireworks_ai-embedding-models"), + "fireworks-ai-up-to-150m": entry("fireworks_ai-embedding-models"), + } + ) + + assert snap.legacy_sets["fireworks_ai_models"] == frozenset({"accounts/fireworks/models/kimi-k2"}) + assert snap.legacy_sets["fireworks_ai_embedding_models"] == frozenset({"nomic-ai/nomic-embed-text-v1.5"}) + + +def test_ai21_splits_on_chat_mode() -> None: + snap = snapshot( + { + "jamba-large-1.7": entry("ai21", mode="chat"), + "j2-ultra": entry("ai21", mode="completion"), + } + ) + + assert snap.legacy_sets["ai21_chat_models"] == frozenset({"jamba-large-1.7"}) + assert snap.legacy_sets["ai21_models"] == frozenset({"j2-ultra"}) + + +def test_meta_llama_provider_feeds_the_llama_models_set() -> None: + snap = snapshot({"Llama-4-Maverick-17B": entry("meta_llama")}) + + assert snap.legacy_sets["llama_models"] == frozenset({"Llama-4-Maverick-17B"}) + assert snap.models_by_provider["meta_llama"] == frozenset({"Llama-4-Maverick-17B"}) + + +def test_fallback_buckets_prefixed_keys_whose_provider_is_a_known_provider() -> None: + snap = snapshot({"vertex_ai/orphaned-model": entry("vertex_ai")}) + + assert "vertex_ai/orphaned-model" in snap.models_by_provider["vertex_ai"] + + +def test_fallback_routes_a_suffixed_provider_to_its_known_root() -> None: + snap = snapshot({"vertex_ai/qwen3-coder": entry("vertex_ai-qwen_models")}) + + assert "vertex_ai/qwen3-coder" in snap.models_by_provider["vertex_ai"] + + +def test_fallback_drops_providers_that_are_not_known_providers() -> None: + snap = snapshot({"tavily/search": entry("tavily"), "serper/search": entry("serper")}) + + assert "tavily" not in snap.models_by_provider + assert "serper" not in snap.models_by_provider + + +def test_fallback_drops_slashless_keys_because_they_are_pricing_tiers() -> None: + snap = snapshot({"together-ai-up-to-4b": entry("together_ai"), "nebius-up-to-8b": entry("nebius")}) + + assert "together-ai-up-to-4b" not in snap.models_by_provider["together_ai"] + assert "nebius-up-to-8b" not in snap.models_by_provider["nebius"] + + +def test_fallback_entries_never_reach_model_list() -> None: + snap = snapshot( + { + "vertex_ai/orphaned-model": entry("vertex_ai"), + "nebius/Qwen/Qwen3-4B": entry("nebius"), + "claude-sonnet-4-5": entry("anthropic"), + } + ) + + assert "vertex_ai/orphaned-model" in snap.models_by_provider["vertex_ai"] + assert "nebius/Qwen/Qwen3-4B" in snap.models_by_provider["nebius"] + assert "vertex_ai/orphaned-model" not in snap.model_list_set + assert "nebius/Qwen/Qwen3-4B" not in snap.model_list_set + assert "vertex_ai/orphaned-model" not in snap.model_list + assert "nebius/Qwen/Qwen3-4B" not in snap.model_list + assert "claude-sonnet-4-5" in snap.model_list_set + + +def test_rebuild_prunes_models_dropped_from_the_cost_map() -> None: + before = snapshot({"claude-keep": entry("anthropic"), "claude-drop": entry("anthropic")}) + after = snapshot({"claude-keep": entry("anthropic")}) + + assert {"claude-keep", "claude-drop"} <= before.legacy_sets["anthropic_models"] + assert "claude-drop" in before.model_list_set + assert after.legacy_sets["anthropic_models"] == frozenset({"claude-keep"}) + assert "claude-drop" not in after.models_by_provider["anthropic"] + assert "claude-drop" not in after.model_list_set + assert "claude-drop" not in after.model_list + + +def test_rebuild_prunes_fallback_bucketed_models_too() -> None: + before = snapshot({"vertex_ai/keep": entry("vertex_ai"), "vertex_ai/drop": entry("vertex_ai")}) + after = snapshot({"vertex_ai/keep": entry("vertex_ai")}) + + assert "vertex_ai/drop" in before.models_by_provider["vertex_ai"] + assert "vertex_ai/drop" not in after.models_by_provider["vertex_ai"] + + +@pytest.mark.parametrize( + "set_name, seed", + [ + ("bedrock_converse_models", constants.BEDROCK_CONVERSE_MODELS), + ("wandb_models", constants.WANDB_MODELS), + ("empower_models", constants.empower_models), + ("modelscope_models", constants.modelscope_models), + ], +) +def test_static_seeds_survive_a_build_from_an_empty_cost_map(set_name: str, seed: Iterable[str]) -> None: + snap = snapshot({}) + + assert snap.legacy_sets[set_name] == frozenset(seed) + + +def test_an_empty_cost_map_leaves_unseeded_sets_empty() -> None: + snap = snapshot({}) + + assert snap.legacy_sets["anthropic_models"] == frozenset() + assert snap.legacy_sets["groq_models"] == frozenset() + + +def test_building_does_not_mutate_the_constants_sets() -> None: + tracked = ( + "empower_models", + "modelscope_models", + "WANDB_MODELS", + "BEDROCK_CONVERSE_MODELS", + "replicate_models", + "clarifai_models", + "huggingface_models", + "together_ai_models", + "baseten_models", + "open_ai_embedding_models", + "cohere_embedding_models", + "bedrock_embedding_models", + ) + before = {name: frozenset(getattr(constants, name)) for name in tracked} + + snapshot( + { + "empower/empower-brand-new": entry("empower"), + "Qwen/Qwen4-Brand-New": entry("modelscope"), + "openai/gpt-oss-brand-new": entry("wandb"), + "anthropic.claude-brand-new-v1:0": entry("bedrock_converse"), + } + ) + + assert {name: frozenset(getattr(constants, name)) for name in tracked} == before + + +def test_static_model_names_are_injected_into_their_provider_buckets() -> None: + snap = snapshot({}, static_model_names={"petals_models": frozenset({"petals-team/StableBeluga2"})}) + + assert snap.models_by_provider["petals"] == frozenset({"petals-team/StableBeluga2"}) + + +def test_models_by_provider_unions_every_source_set_for_a_provider() -> None: + snap = snapshot( + { + "gpt-5.2": entry("openai"), + "gpt-3.5-turbo-instruct": entry("text-completion-openai"), + "anthropic.claude-sonnet-4-5-v1:0": entry("bedrock"), + "anthropic.claude-converse-v1:0": entry("bedrock_converse"), + } + ) + + assert {"gpt-5.2", "gpt-3.5-turbo-instruct"} <= snap.models_by_provider["openai"] + assert snap.models_by_provider["text-completion-openai"] == frozenset({"gpt-3.5-turbo-instruct"}) + assert {"anthropic.claude-sonnet-4-5-v1:0", "anthropic.claude-converse-v1:0"} <= snap.models_by_provider["bedrock"] + + +def test_all_embedding_models_unions_every_embedding_source() -> None: + snap = snapshot( + { + "vertex_ai/text-embedding-005": entry("vertex_ai-embedding-models", mode="embedding"), + "nomic-ai/nomic-embed-text-v1.5": entry("fireworks_ai-embedding-models", mode="embedding"), + "BAAI/bge-en-icl": entry("nebius-embedding-models", mode="embedding"), + "E5-Mistral-7B-Instruct": entry("sambanova-embedding-models", mode="embedding"), + "bge-multilingual-gemma2": entry("ovhcloud-embedding-models", mode="embedding"), + } + ) + + assert { + "vertex_ai/text-embedding-005", + "nomic-ai/nomic-embed-text-v1.5", + "BAAI/bge-en-icl", + "E5-Mistral-7B-Instruct", + "bge-multilingual-gemma2", + } <= snap.all_embedding_models + assert frozenset(constants.open_ai_embedding_models) <= snap.all_embedding_models + + +def test_model_list_is_sorted_and_matches_model_list_set() -> None: + snap = snapshot({"zzz-model": entry("anthropic"), "aaa-model": entry("anthropic")}) + + assert snap.model_list == tuple(sorted(snap.model_list)) + assert frozenset(snap.model_list) == snap.model_list_set + + +def test_snapshot_collections_are_immutable() -> None: + snap = snapshot({"claude-sonnet-4-5": entry("anthropic")}) + + assert isinstance(snap.legacy_sets["anthropic_models"], frozenset) + assert isinstance(snap.models_by_provider["anthropic"], frozenset) + assert isinstance(snap.model_list_set, frozenset) + assert isinstance(snap.all_embedding_models, frozenset) + assert isinstance(snap.model_list, tuple) + + +def test_fallback_only_provider_buckets_are_frozensets() -> None: + """These buckets skip the composition union, so nothing else re-freezes them.""" + assert "sagemaker" not in _PROVIDER_COMPOSITION + + snap = snapshot({"sagemaker/my-endpoint": entry("sagemaker")}) + + assert snap.models_by_provider["sagemaker"] == frozenset({"sagemaker/my-endpoint"}) + assert isinstance(snap.models_by_provider["sagemaker"], frozenset) + assert all(isinstance(models, frozenset) for models in snap.models_by_provider.values()) + + +@pytest.mark.parametrize("mutator", ["add", "discard", "clear", "update"]) +def test_legacy_sets_reject_in_place_mutation(mutator: str) -> None: + """The pre-refactor sets were mutable; silent in-place edits are what went stale.""" + snap = snapshot({"claude-sonnet-4-5": entry("anthropic")}) + + with pytest.raises(AttributeError): + getattr(snap.legacy_sets["anthropic_models"], mutator) + + +def test_served_names_are_frozensets_and_model_list_is_a_fresh_list() -> None: + """litellm.model_list stayed a list for compat, but callers must not be able to poison it.""" + import litellm + + assert isinstance(litellm.anthropic_models, frozenset) + assert isinstance(litellm.models_by_provider["anthropic"], frozenset) + assert isinstance(litellm.model_list_set, frozenset) + + borrowed = litellm.model_list + assert isinstance(borrowed, list) + borrowed.append("not-a-real-model") + assert "not-a-real-model" not in litellm.model_list + assert "not-a-real-model" not in litellm.model_list_set + + +def test_entries_without_a_string_provider_are_dropped() -> None: + broken = {"no-provider": {"mode": "chat"}, "null-provider": {"litellm_provider": None, "mode": "chat"}} + baseline = snapshot({}) + snap = snapshot({**broken, "claude-sonnet-4-5": entry("anthropic")}) + + assert snap.model_list_set - baseline.model_list_set == frozenset({"claude-sonnet-4-5"}) + assert not any(models & broken.keys() for models in snap.legacy_sets.values()) + assert not any(models & broken.keys() for models in snap.models_by_provider.values()) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 71e686563a5..f08f9919c8e 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -563,6 +563,8 @@ def test_token_counter(): import unittest +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from unittest.mock import MagicMock, patch from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding @@ -571,6 +573,16 @@ from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding _select_tokenizer_helper.cache_clear() +@contextmanager +def registered_models(model_cost_additions: Mapping[str, Mapping[str, str]]) -> Iterator[None]: + try: + with patch.dict(litellm.model_cost, model_cost_additions): + litellm.add_known_models() + yield + finally: + litellm.add_known_models() + + class TestTokenizerSelection(unittest.TestCase): def setUp(self): """Clear the LRU cache before each test method. @@ -602,10 +614,11 @@ class TestTokenizerSelection(unittest.TestCase): mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") # Add Cohere model to the list for testing - litellm.cohere_models = ["command-r-v1"] + with registered_models({"command-r-v1": {"litellm_provider": "cohere", "mode": "chat"}}): + self.assertIn("command-r-v1", litellm.cohere_models) - # Test with Cohere model - result = _select_tokenizer_helper("command-r-v1") + # Test with Cohere model + result = _select_tokenizer_helper("command-r-v1") # Verify the attempt to load Cohere tokenizer mock_from_pretrained.assert_called_once_with( @@ -622,10 +635,11 @@ class TestTokenizerSelection(unittest.TestCase): mock_from_str.side_effect = Exception("Failed to load tokenizer") # Add Claude model to the list for testing - litellm.anthropic_models = ["claude-2"] + with registered_models({"claude-2": {"litellm_provider": "anthropic", "mode": "chat"}}): + self.assertIn("claude-2", litellm.anthropic_models) - # Test with Claude model - result = _select_tokenizer_helper("claude-2") + # Test with Claude model + result = _select_tokenizer_helper("claude-2") # Verify the attempt to load Claude tokenizer mock_from_str.assert_called_once_with(claude_json_str) diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 275fb460b9f..63a63bcca10 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -23,7 +23,6 @@ from litellm.types.utils import LlmProviders @pytest.fixture def local_cost_map(monkeypatch): original_model_cost = litellm.model_cost - original_bedrock_mantle_models = set(litellm.bedrock_mantle_models) try: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") litellm.model_cost = litellm.get_model_cost_map(url="") @@ -32,8 +31,7 @@ def local_cost_map(monkeypatch): yield finally: litellm.model_cost = original_model_cost - litellm.bedrock_mantle_models.clear() - litellm.bedrock_mantle_models.update(original_bedrock_mantle_models) + litellm.add_known_models() litellm.get_model_info.cache_clear() diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index 0750fb9e405..34af936445d 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -7,11 +7,27 @@ import os from unittest import mock import httpx +import pytest import litellm from litellm.llms.inception.chat.transformation import InceptionChatConfig +@pytest.fixture +def local_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + try: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + litellm.add_known_models() + yield + finally: + litellm.model_cost = original_model_cost + litellm.add_known_models() + litellm.get_model_info.cache_clear() + + def test_inception_config_initialization(): config = InceptionChatConfig() assert config.custom_llm_provider == "inception" @@ -231,14 +247,9 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(): +def test_inception_model_configuration(local_cost_map): from litellm import get_model_info - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.inception_models = set() - litellm.add_known_models() - info = get_model_info("inception/mercury-2") assert info.get("litellm_provider") == "inception" assert info.get("mode") == "chat" @@ -251,12 +262,7 @@ def test_inception_model_configuration(): assert info.get("supports_response_schema") is True -def test_inception_model_list_populated(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.inception_models = set() - litellm.add_known_models() - +def test_inception_model_list_populated(local_cost_map): assert "inception/mercury-2" in litellm.inception_models for model in litellm.inception_models: assert model.startswith("inception/") diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 9b7c8dd3742..4c634068a4c 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -15,6 +15,21 @@ from litellm.llms.inception.completion.transformation import ( ) +@pytest.fixture +def local_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + try: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + litellm.add_known_models() + yield + finally: + litellm.model_cost = original_model_cost + litellm.add_known_models() + litellm.get_model_info.cache_clear() + + def _fim_response_bytes(): return json.dumps( { @@ -143,14 +158,9 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(): +def test_inception_fim_model_configuration(local_cost_map): from litellm import get_model_info - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.text_completion_inception_models = set() - litellm.add_known_models() - assert ( "text-completion-inception/mercury-edit-2" in litellm.text_completion_inception_models diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index 8d46151ecce..8708a570c51 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -188,8 +188,9 @@ class TestOllamaModelInfo: def test_get_models_fallback_on_error(self, monkeypatch): """ - If the httpx.get call raises an exception, get_models should - fall back to the static models_by_provider list prefixed by 'ollama/'. + The static fallback bucket mixes bare registry seeds ('llama2') with + already-prefixed cost-map keys ('ollama/llama3'), so get_models must + normalize to exactly one 'ollama/' prefix instead of blindly prepending. """ def mock_get(url, headers): @@ -198,8 +199,16 @@ class TestOllamaModelInfo: monkeypatch.setattr(httpx, "get", mock_get) info = OllamaModelInfo() models = info.get_models() - # Default static ollama_models is ['llama2'], so expect ['ollama/llama2'] - assert models == ["ollama/llama2"] + + static = litellm.models_by_provider["ollama"] + already_prefixed = {name for name in static if name.startswith("ollama/")} + + assert models + assert models == sorted(models) + assert all(name.startswith("ollama/") for name in models) + assert not any(name.removeprefix("ollama/").startswith("ollama/") for name in models) + assert "ollama/llama2" in models + assert already_prefixed and already_prefixed <= set(models) def test_get_models_no_double_prefix(self, monkeypatch): """ diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 5d1cef87fba..b0148e8f052 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -2,6 +2,7 @@ from unittest.mock import AsyncMock, patch import pytest +import litellm from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, Member from litellm.proxy.auth.handle_jwt import JWTAuthManager @@ -709,3 +710,67 @@ def test_expand_wildcard_invalid_litellm_params_passthrough(): # Even if LiteLLM_Params construction fails the deployment should survive result = expand_wildcard_deployments_for_model_info([deployment]) assert result == [deployment] + + +@pytest.fixture +def register_in_cost_map(monkeypatch): + """LIT-4947: mutate the cost map, rebuild the registry, and put the world back afterwards.""" + + def _register(model: str, litellm_provider: str) -> str: + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": litellm_provider, "mode": "chat"}) + litellm.add_known_models() + return model + + yield _register + monkeypatch.undo() + litellm.add_known_models() + + +@pytest.mark.parametrize( + "provider, litellm_provider", + [ + ("vertex_ai", "vertex_ai-language-models"), + ("openai", "openai"), + ("anthropic", "anthropic"), + ("bedrock", "bedrock_converse"), + ("cohere", "cohere_chat"), + ], +) +def test_cost_map_addition_reaches_provider_enumeration(register_in_cost_map, provider, litellm_provider): + """LIT-4947: models_by_provider union entries used to be frozen at import, so reloads went unseen.""" + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model = register_in_cost_map(f"lit-4947-{provider}-brand-new", litellm_provider) + + assert model in litellm.models_by_provider[provider] + assert f"{provider}/{model}" in get_known_models_from_wildcard(wildcard_model=f"{provider}/*") + + +def test_cost_map_addition_reaches_model_list_set(register_in_cost_map): + """LIT-4947: model_list/model_list_set were import-time snapshots that no reload could refresh.""" + model = register_in_cost_map("lit-4947-openai-brand-new", "openai") + + assert model in litellm.model_list_set + assert model in litellm.model_list + assert model in litellm.open_ai_chat_completion_models + + +def test_cost_map_removal_disappears_from_provider_enumeration(monkeypatch): + """LIT-4947: a rebuild must prune, otherwise a shrinking cost map leaves phantom models routable.""" + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + monkeypatch.setitem(litellm.model_cost, "lit-4947-doomed", {"litellm_provider": "anthropic", "mode": "chat"}) + litellm.add_known_models() + try: + assert "lit-4947-doomed" in litellm.models_by_provider["anthropic"] + assert "anthropic/lit-4947-doomed" in get_known_models_from_wildcard(wildcard_model="anthropic/*") + + monkeypatch.delitem(litellm.model_cost, "lit-4947-doomed") + litellm.add_known_models() + + assert "lit-4947-doomed" not in litellm.models_by_provider["anthropic"] + assert "lit-4947-doomed" not in litellm.model_list_set + assert "anthropic/lit-4947-doomed" not in get_known_models_from_wildcard(wildcard_model="anthropic/*") + finally: + monkeypatch.undo() + litellm.add_known_models() diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py index df1d096b3e2..5c582742eb3 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py @@ -55,12 +55,14 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): table = _attach_litellm_config(mock_prisma) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - fake_cost_map = {"gpt-4": {"input_cost": 0.03}, "gpt-3.5": {"input_cost": 0.002}} + fake_cost_map = { + "gpt-4": {"input_cost": 0.03, "litellm_provider": "openai", "mode": "chat"}, + "gpt-3.5": {"input_cost": 0.002, "litellm_provider": "openai", "mode": "chat"}, + } monkeypatch.setattr( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", AsyncMock(return_value=ModelCostMapReloaded(model_cost_map=fake_cost_map)), ) - monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None) monkeypatch.setattr("litellm.model_cost", {}, raising=False) monkeypatch.setattr( "litellm.proxy.proxy_server._invalidate_model_cost_lowercase_map", @@ -85,6 +87,59 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): } assert table.upsert.await_count == 1 + import litellm as litellm_module + + assert litellm_module.open_ai_chat_completion_models == frozenset({"gpt-4", "gpt-3.5"}) + + +def test_reload_model_cost_map_rebuilds_provider_enumeration( + client, auth_as, monkeypatch, mock_prisma +): + """A reload must re-derive provider enumeration: new models appear and dropped models disappear. + + Regression (LIT-4947): the derived collections were import-time snapshots, so a reload + updated pricing while `anthropic/*` kept expanding to the models the pod booted with. + """ + from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + import litellm as litellm_module + + _attach_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + async def _fake_invalidate(name): + return None + + monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate) + + booted_map = {"claude-doomed": {"litellm_provider": "anthropic", "mode": "chat"}} + monkeypatch.setattr("litellm.model_cost", booted_map, raising=False) + litellm_module.add_known_models() + assert "claude-doomed" in litellm_module.models_by_provider["anthropic"] + + reloaded_map = {"claude-fresh": {"litellm_provider": "anthropic", "mode": "chat"}} + monkeypatch.setattr( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", + AsyncMock(return_value=ModelCostMapReloaded(model_cost_map=reloaded_map)), + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post("/reload/model_cost_map") + assert response.status_code == 200 + + assert litellm_module.models_by_provider["anthropic"] == frozenset({"claude-fresh"}) + assert "claude-fresh" in litellm_module.anthropic_models + assert "claude-fresh" in litellm_module.model_list_set + assert "claude-doomed" not in litellm_module.anthropic_models + assert "claude-doomed" not in litellm_module.model_list_set + + expanded = get_known_models_from_wildcard(wildcard_model="anthropic/*") + assert "anthropic/claude-fresh" in expanded + assert "anthropic/claude-doomed" not in expanded + def test_reload_model_cost_map_fetch_failure_502_keeps_map( client, auth_as, monkeypatch, mock_prisma diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b23d3333ea7..859d3bbd71c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2708,25 +2708,33 @@ def test_register_model_with_scientific_notation(): _invalidate_model_cost_lowercase_map() -def test_register_model_openrouter_without_slash(): +@pytest.fixture +def registered_model_cleanup(): + """Register models without leaking them into litellm.model_cost or the registry.""" + from litellm.utils import _invalidate_model_cost_lowercase_map + + original_keys = frozenset(litellm.model_cost) + + yield litellm.register_model + + for key in frozenset(litellm.model_cost) - original_keys: + del litellm.model_cost[key] + _invalidate_model_cost_lowercase_map() + litellm.add_known_models() + + +@pytest.mark.parametrize( + "key", + ["my-custom-alias", "openrouter/some-model", "openrouter/openai/gpt-4-turbo"], +) +def test_register_model_openrouter_stores_the_key_verbatim(registered_model_cleanup, key): """ - Test that register_model handles openrouter models without '/' in the name. - - Fixes https://github.com/BerriAI/litellm/issues/18936 - - Previously, the code did `split_string[1]` which would fail with IndexError - when the model name didn't contain '/'. Now it uses `split_string[-1]` which - always works. + register_model must not raise on openrouter names without '/' (issue #18936) and now + stores whatever key it was handed, matching what the cost map itself contributes. """ - # Clear any existing entries - litellm.openrouter_models.discard("my-custom-alias") - litellm.openrouter_models.discard("gpt-4") - litellm.openrouter_models.discard("openai/gpt-4") - - # Test 1: Model name without '/' (this was the bug - would raise IndexError) - litellm.register_model( + registered_model_cleanup( { - "my-custom-alias": { + key: { "max_tokens": 8192, "input_cost_per_token": 0.00001, "output_cost_per_token": 0.00002, @@ -2735,35 +2743,60 @@ def test_register_model_openrouter_without_slash(): }, } ) - assert "my-custom-alias" in litellm.openrouter_models - # Test 2: Model name with single '/' (openrouter/model format) - litellm.register_model( + assert key in litellm.openrouter_models + assert key in litellm.models_by_provider["openrouter"] + + +def test_register_model_reaches_providers_outside_the_legacy_branch_chain(registered_model_cleanup): + """ + register_model used to hand-roll a 15-provider if/elif chain, so anything outside it + (groq here) was priced but never enumerable. It now runs the same derivation as a rebuild. + """ + registered_model_cleanup( { - "openrouter/gpt-4": { + "groq/brand-new-groq-model": { "max_tokens": 8192, "input_cost_per_token": 0.00001, "output_cost_per_token": 0.00002, - "litellm_provider": "openrouter", + "litellm_provider": "groq", "mode": "chat", }, } ) - assert "gpt-4" in litellm.openrouter_models - # Test 3: Model name with double '/' (openrouter/provider/model format) - litellm.register_model( + assert "groq/brand-new-groq-model" in litellm.groq_models + assert "groq/brand-new-groq-model" in litellm.models_by_provider["groq"] + assert "groq/brand-new-groq-model" in litellm.model_list_set + + +def test_register_model_does_not_detach_names_from_the_registry(registered_model_cleanup): + """ + A later rebuild must still be able to reach every name register_model touched; if + register_model wrote into litellm.__dict__ the rebuild would be invisible forever. + """ + from litellm.litellm_core_utils.model_registry import REGISTRY_SET_NAMES + + registered_model_cleanup( { - "openrouter/openai/gpt-4-turbo": { - "max_tokens": 8192, + "groq/registered-then-rebuilt": { "input_cost_per_token": 0.00001, "output_cost_per_token": 0.00002, - "litellm_provider": "openrouter", + "litellm_provider": "groq", "mode": "chat", }, } ) - assert "openai/gpt-4-turbo" in litellm.openrouter_models + assert not (REGISTRY_SET_NAMES & litellm.__dict__.keys()) + + litellm.model_cost["groq/added-after-register"] = {"litellm_provider": "groq", "mode": "chat"} + try: + litellm.add_known_models() + assert "groq/added-after-register" in litellm.groq_models + assert "groq/registered-then-rebuilt" in litellm.groq_models + finally: + del litellm.model_cost["groq/added-after-register"] + litellm.add_known_models() def test_reasoning_content_preserved_in_text_completion_wrapper(): diff --git a/tests/test_litellm/test_zz_registry_leak_probe.py b/tests/test_litellm/test_zz_registry_leak_probe.py new file mode 100644 index 00000000000..1a8fbd0d42d --- /dev/null +++ b/tests/test_litellm/test_zz_registry_leak_probe.py @@ -0,0 +1,29 @@ +import litellm + +from litellm.litellm_core_utils.model_registry import REGISTRY_SET_NAMES + + +def test_no_registry_names_are_shadowed(): + assert not (REGISTRY_SET_NAMES & litellm.__dict__.keys()) + assert not ({"models_by_provider", "model_list", "model_list_set"} & litellm.__dict__.keys()) + + +def test_cost_map_survived_the_suite(): + assert len(litellm.model_cost) > 2000 + assert "claude-sonnet-4-5-20250929" in litellm.model_cost + + +def test_registry_still_rebuilds(): + litellm.model_cost["zz-probe-model"] = {"litellm_provider": "anthropic", "mode": "chat"} + try: + litellm.add_known_models() + assert "zz-probe-model" in litellm.anthropic_models + assert "zz-probe-model" in litellm.models_by_provider["anthropic"] + finally: + del litellm.model_cost["zz-probe-model"] + litellm.add_known_models() + + +def test_registry_agrees_with_cost_map(): + assert litellm.anthropic_models <= litellm.model_cost.keys() + assert litellm.groq_models <= litellm.model_cost.keys() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 9976be98522..ebba5abfba2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23350 + "limit": 23340 }, "LIT002": { - "limit": 27239 + "limit": 27116 }, "LIT003": { "limit": 292