diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..dfbad5ada08 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -9,7 +9,7 @@ import threading import os import datetime -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import ( Any, Callable, @@ -547,9 +547,13 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete +from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map as get_model_cost_map, + load_model_cost_map_at_import, + mark_litellm_import_complete, +) -model_cost = get_model_cost_map(url=model_cost_map_url) +model_cost = load_model_cost_map_at_import(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount apply_search_free_tier: bool = False cost_margin_config: Dict[ @@ -590,126 +594,126 @@ 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() +open_ai_chat_completion_models: Set[str] = set() +open_ai_text_completion_models: Set[str] = set() +cohere_models: Set[str] = set() +cohere_chat_models: Set[str] = set() +mistral_chat_models: Set[str] = set() +text_completion_codestral_models: Set[str] = set() +text_completion_inception_models: Set[str] = set() +anthropic_models: Set[str] = set() +openrouter_models: Set[str] = set() +datarobot_models: Set[str] = set() +vertex_language_models: Set[str] = set() +vertex_vision_models: Set[str] = set() +vertex_chat_models: Set[str] = set() +vertex_code_chat_models: Set[str] = set() +vertex_ai_image_models: Set[str] = set() +vertex_ai_video_models: Set[str] = set() +vertex_text_models: Set[str] = set() +vertex_code_text_models: Set[str] = set() +vertex_embedding_models: Set[str] = set() +vertex_anthropic_models: Set[str] = set() +vertex_llama3_models: Set[str] = set() +vertex_deepseek_models: Set[str] = set() +vertex_ai_ai21_models: Set[str] = set() +vertex_mistral_models: Set[str] = set() +vertex_openai_models: Set[str] = set() +vertex_minimax_models: Set[str] = set() +vertex_moonshot_models: Set[str] = set() +vertex_zai_models: Set[str] = set() +ai21_models: Set[str] = set() +ai21_chat_models: Set[str] = set() +nlp_cloud_models: Set[str] = set() +aleph_alpha_models: Set[str] = set() +bedrock_models: Set[str] = set() +bedrock_converse_models: Set[str] = set(BEDROCK_CONVERSE_MODELS) +fal_ai_models: Set[str] = set() +fireworks_ai_models: Set[str] = set() baseten_models: Set[str] = 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() -nadir_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() +fireworks_ai_embedding_models: Set[str] = set() +deepinfra_models: Set[str] = set() +perplexity_models: Set[str] = set() +watsonx_models: Set[str] = set() +gemini_models: Set[str] = set() +xai_models: Set[str] = set() +zai_models: Set[str] = set() +deepseek_models: Set[str] = set() +tencent_models: Set[str] = set() +runwayml_models: Set[str] = set() +azure_ai_models: Set[str] = set() +jina_ai_models: Set[str] = set() +voyage_models: Set[str] = set() +infinity_models: Set[str] = set() +heroku_models: Set[str] = set() +databricks_models: Set[str] = set() +cloudflare_models: Set[str] = set() +codestral_models: Set[str] = set() +friendliai_models: Set[str] = set() +featherless_ai_models: Set[str] = set() +palm_models: Set[str] = set() +groq_models: Set[str] = set() +azure_models: Set[str] = set() +azure_anthropic_models: Set[str] = set() +azure_text_models: Set[str] = set() +anyscale_models: Set[str] = set() +cerebras_models: Set[str] = set() +nadir_models: Set[str] = set() +galadriel_models: Set[str] = set() +nvidia_nim_models: Set[str] = set() +nvidia_riva_models: Set[str] = set() +soniox_models: Set[str] = set() +sambanova_models: Set[str] = set() scaledown_models: Final[Set[str]] = set() # Price reloads update this registry through existing references. -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() -qwencloud_models: Set = set() -qwen_ai_platform_models: Set = set() -moonshot_models: Set = set() -publicai_models: Set = set() -darkbloom_models: Set = set() +sambanova_embedding_models: Set[str] = set() +novita_models: Set[str] = set() +assemblyai_models: Set[str] = set() +snowflake_models: Set[str] = set() +gradient_ai_models: Set[str] = set() +llama_models: Set[str] = set() +nscale_models: Set[str] = set() +nebius_models: Set[str] = set() +nebius_embedding_models: Set[str] = set() +aiml_models: Set[str] = set() +deepgram_models: Set[str] = set() +elevenlabs_models: Set[str] = set() +dashscope_models: Set[str] = set() +qwencloud_models: Set[str] = set() +qwen_ai_platform_models: Set[str] = set() +moonshot_models: Set[str] = set() +publicai_models: Set[str] = set() +darkbloom_models: Set[str] = set() openinfer_models: Final[Set[str]] = 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() -edenai_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() +v0_models: Set[str] = set() +morph_models: Set[str] = set() +lambda_ai_models: Set[str] = set() +inception_models: Set[str] = set() +hyperbolic_models: Set[str] = set() +black_forest_labs_models: Set[str] = set() +recraft_models: Set[str] = set() +cometapi_models: Set[str] = set() +oci_models: Set[str] = set() +vercel_ai_gateway_models: Set[str] = set() +edenai_models: Set[str] = set() +volcengine_models: Set[str] = set() +wandb_models: Set[str] = set(WANDB_MODELS) +ovhcloud_models: Set[str] = set() +ovhcloud_embedding_models: Set[str] = set() +lemonade_models: Set[str] = set() +docker_model_runner_models: Set[str] = set() +amazon_nova_models: Set[str] = set() +stability_models: Set[str] = set() +github_copilot_models: Set[str] = set() microsoft_365_copilot_models: Set[str] = set() -chatgpt_models: Set = set() -minimax_models: Set = set() -aws_polly_models: Set = set() -transcribe_models: Set = set() -gigachat_models: Set = set() -llamagate_models: Set = set() +chatgpt_models: Set[str] = set() +minimax_models: Set[str] = set() +aws_polly_models: Set[str] = set() +transcribe_models: Set[str] = set() +gigachat_models: Set[str] = set() +llamagate_models: Set[str] = set() coralbricks_models: Set[str] = set() -reducto_models: Set = set() -bedrock_mantle_models: Set = set() +reducto_models: Set[str] = set() +bedrock_mantle_models: Set[str] = set() def is_bedrock_pricing_only_model(key: str) -> bool: @@ -755,273 +759,268 @@ def _is_past_deprecation_date(value: Mapping[str, object]) -> bool: return False +def _provider_model_set(key: str, value: Mapping[str, object]) -> Optional[Tuple[Set[str], str]]: + provider: Final = value.get("litellm_provider") + if provider == "openai" and not is_openai_finetune_model(key): + return open_ai_chat_completion_models, key + if provider == "text-completion-openai": + return open_ai_text_completion_models, key + if provider == "azure_text": + return azure_text_models, key + if provider == "cohere": + return cohere_models, key + if provider == "cohere_chat": + return cohere_chat_models, key + if provider == "mistral": + return mistral_chat_models, key + if provider == "anthropic": + return anthropic_models, key + if provider == "empower": + return empower_models, key + if provider == "openrouter": + return openrouter_models, key + if provider == "vercel_ai_gateway": + return vercel_ai_gateway_models, key + if provider == "edenai": + return edenai_models, key + if provider == "datarobot": + return datarobot_models, key + if provider == "vertex_ai-text-models": + return vertex_text_models, key + if provider == "vertex_ai-code-text-models": + return vertex_code_text_models, key + if provider == "vertex_ai-language-models": + return vertex_language_models, key + if provider == "vertex_ai-vision-models": + return vertex_vision_models, key + if provider == "vertex_ai-chat-models": + return vertex_chat_models, key + if provider == "vertex_ai-code-chat-models": + return vertex_code_chat_models, key + if provider == "vertex_ai-embedding-models": + return vertex_embedding_models, key + if provider == "vertex_ai-anthropic_models": + return vertex_anthropic_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-llama_models": + return vertex_llama3_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-deepseek_models": + return vertex_deepseek_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-mistral_models": + return vertex_mistral_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-ai21_models": + return vertex_ai_ai21_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-image-models": + return vertex_ai_image_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-video-models": + return vertex_ai_video_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-openai_models": + return vertex_openai_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-minimax_models": + return vertex_minimax_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-moonshot_models": + return vertex_moonshot_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-zai_models": + return vertex_zai_models, key.replace("vertex_ai/", "") + if provider == "ai21": + if value.get("mode") == "chat": + return ai21_chat_models, key + else: + return ai21_models, key + if provider == "nlp_cloud": + return nlp_cloud_models, key + if provider == "aleph_alpha": + return aleph_alpha_models, key + if provider == "bedrock" and value.get("mode") == "guardrail": + return None + if provider == "bedrock" and not is_bedrock_pricing_only_model(key): + return bedrock_models, key + if provider == "bedrock_converse": + return bedrock_converse_models, key + if provider == "deepinfra": + return deepinfra_models, key + if provider == "perplexity": + return perplexity_models, key + if provider == "watsonx": + return watsonx_models, key + if provider == "gemini": + return gemini_models, key + if 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: + return fireworks_ai_models, key + if provider == "baseten": + return None if _is_past_deprecation_date(value) else (baseten_models, key) + if 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: + return fireworks_ai_embedding_models, key + if provider == "text-completion-codestral": + return text_completion_codestral_models, key + if provider == "text-completion-inception": + return text_completion_inception_models, key + if provider == "xai": + return xai_models, key + if provider == "zai": + return zai_models, key + if provider == "fal_ai": + return fal_ai_models, key + if provider == "deepseek": + return deepseek_models, key + if provider == "tencent": + return tencent_models, key + if provider == "runwayml": + return runwayml_models, key + if provider == "meta_llama": + return llama_models, key + if provider == "nscale": + return nscale_models, key + if provider == "azure_ai": + return azure_ai_models, key + if provider == "voyage": + return voyage_models, key + if provider == "infinity": + return infinity_models, key + if provider == "databricks": + return databricks_models, key + if provider == "cloudflare": + return cloudflare_models, key + if provider == "codestral": + return codestral_models, key + if provider == "friendliai": + return friendliai_models, key + if provider == "palm": + return palm_models, key + if provider == "groq": + return groq_models, key + if provider == "azure": + return azure_models, key + if provider == "azure_anthropic": + return azure_anthropic_models, key + if provider == "anyscale": + return anyscale_models, key + if provider == "cerebras": + return cerebras_models, key + if provider == "nadir": + return nadir_models, key + if provider == "galadriel": + return galadriel_models, key + if provider == "nvidia_nim": + return nvidia_nim_models, key + if provider == "nvidia_riva": + return nvidia_riva_models, key + if provider == "soniox": + return soniox_models, key + if provider == "sambanova": + return sambanova_models, key + if provider == "scaledown": + return scaledown_models, key + if provider == "sambanova-embedding-models": + return sambanova_embedding_models, key + if provider == "novita": + return novita_models, key + if provider == "nebius-chat-models": + return nebius_models, key + if provider == "nebius-embedding-models": + return nebius_embedding_models, key + if provider == "aiml": + return aiml_models, key + if provider == "assemblyai": + return assemblyai_models, key + if provider == "jina_ai": + return jina_ai_models, key + if provider == "snowflake": + return snowflake_models, key + if provider == "gradient_ai": + return gradient_ai_models, key + if provider == "featherless_ai": + return featherless_ai_models, key + if provider == "deepgram": + return deepgram_models, key + if provider == "elevenlabs": + return elevenlabs_models, key + if provider == "heroku": + return heroku_models, key + if provider == "dashscope": + return dashscope_models, key + if provider == "qwencloud": + return qwencloud_models, key + if provider == "qwen_ai_platform": + return qwen_ai_platform_models, key + if provider == "modelscope": + return modelscope_models, key + if provider == "moonshot": + return moonshot_models, key + if provider == "publicai": + return publicai_models, key + if provider == "darkbloom": + return darkbloom_models, key + if provider == "openinfer": + return openinfer_models, key + if provider == "v0": + return v0_models, key + if provider == "morph": + return morph_models, key + if provider == "lambda_ai": + return lambda_ai_models, key + if provider == "inception": + return inception_models, key + if provider == "hyperbolic": + return hyperbolic_models, key + if provider == "black_forest_labs": + return black_forest_labs_models, key + if provider == "recraft": + return recraft_models, key + if provider == "cometapi": + return cometapi_models, key + if provider == "oci": + return oci_models, key + if provider == "volcengine": + return volcengine_models, key + if provider == "wandb": + return wandb_models, key + if provider == "ovhcloud": + return ovhcloud_models, key + if provider == "ovhcloud-embedding-models": + return ovhcloud_embedding_models, key + if provider == "lemonade": + return lemonade_models, key + if provider == "docker_model_runner": + return docker_model_runner_models, key + if provider == "amazon_nova": + return amazon_nova_models, key + if provider == "stability": + return stability_models, key + if provider == "github_copilot": + return github_copilot_models, key + if provider == "microsoft_365_copilot": + return microsoft_365_copilot_models, key + if provider == "chatgpt": + return chatgpt_models, key + if provider == "minimax": + return minimax_models, key + if provider == "aws_polly": + return aws_polly_models, key + if provider == "transcribe": + return transcribe_models, key + if provider == "gigachat": + return gigachat_models, key + if provider == "llamagate": + return llamagate_models, key + if provider == "coralbricks": + return coralbricks_models, key + if provider == "reducto": + return reducto_models, key + if provider == "bedrock_mantle": + return bedrock_mantle_models, key + return None + + def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: for key, value in tuple(model_cost_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") == "edenai": - edenai_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 value.get("mode") == "guardrail": - pass - 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) + if (membership := _provider_model_set(key, value)) is not None: + membership[0].add(membership[1]) elif value.get("litellm_provider") == "baseten": - if _is_past_deprecation_date(value): - baseten_models.discard(key) - else: - baseten_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") == "nadir": - nadir_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") == "scaledown": - scaledown_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") == "qwencloud": - qwencloud_models.add(key) - elif value.get("litellm_provider") == "qwen_ai_platform": - qwen_ai_platform_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") == "openinfer": - openinfer_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") == "microsoft_365_copilot": - microsoft_365_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") == "transcribe": - transcribe_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") == "coralbricks": - coralbricks_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) + baseten_models.discard(key) def add_known_models(model_cost_map: Optional[Dict] = None): @@ -1033,6 +1032,38 @@ def add_known_models(model_cost_map: Optional[Dict] = None): models_by_provider.update(_build_models_by_provider()) +def add_known_models_to_model_list() -> None: + """``add_known_models``, plus listing every model name it newly made known in ``model_list``.""" + known_before: Final = _known_model_names() + add_known_models() + added: Final = _known_model_names() - known_before - model_list_set + model_list.extend(added) + model_list_set.update(added) + + +def _provider_model_set_left( + key: str, before: Mapping[str, object], after: Mapping[str, object] +) -> Optional[Tuple[Set[str], str]]: + old: Final = _provider_model_set(key, before) + new: Final = _provider_model_set(key, after) + if old is None or (new is not None and old[0] is new[0] and old[1] == new[1]): + return None + return old + + +def drop_stale_known_models(changes: Iterable[Tuple[str, Mapping[str, object], Mapping[str, object]]]) -> None: + left: Final = tuple( + membership + for key, before, after in changes + if (membership := _provider_model_set_left(key, before, after)) is not None + ) + if not left: + return + for model_set, name in left: + model_set.discard(name) + models_by_provider.update(_build_models_by_provider()) + + _populate_provider_model_sets(model_cost) # known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary @@ -1063,99 +1094,104 @@ 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 - | nadir_models - | galadriel_models - | nvidia_nim_models - | nvidia_riva_models - | soniox_models - | sambanova_models - | scaledown_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 - | qwencloud_models - | qwen_ai_platform_models - | moonshot_models - | publicai_models - | darkbloom_models - | openinfer_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 - | edenai_models - | volcengine_models - | wandb_models - | ovhcloud_models - | lemonade_models - | docker_model_runner_models - | reducto_models - | bedrock_mantle_models - | coralbricks_models - | set(clarifai_models) -) + +def _known_model_names() -> Set[str]: + return ( + 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 + | nadir_models + | galadriel_models + | nvidia_nim_models + | nvidia_riva_models + | soniox_models + | sambanova_models + | scaledown_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 + | qwencloud_models + | qwen_ai_platform_models + | moonshot_models + | publicai_models + | darkbloom_models + | openinfer_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 + | edenai_models + | volcengine_models + | wandb_models + | ovhcloud_models + | lemonade_models + | docker_model_runner_models + | reducto_models + | bedrock_mantle_models + | coralbricks_models + | set(clarifai_models) + ) + + +model_list = list(_known_model_names()) model_list_set = set(model_list) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 550bbfcde98..bf03793d7e0 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -25,6 +25,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type +from litellm.litellm_core_utils.get_model_cost_map import wait_for_pending_import_model_cost_map from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -465,6 +466,7 @@ def cost_per_token( if model is None: raise Exception("Invalid arg. Model cannot be none.") + wait_for_pending_import_model_cost_map() ## RECONSTRUCT USAGE BLOCK ## if usage_object is not None: @@ -1544,6 +1546,7 @@ def completion_cost( Raises if no candidate model name resolves to pricing. Register the model via `litellm.register_model` or add it to model_prices_and_context_window.json. """ + wait_for_pending_import_model_cost_map() try: call_type = infer_call_type(call_type, completion_response) or "completion" if pricing_provider is not None: diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..25d6cff6054 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -12,6 +12,10 @@ from litellm.constants import ( from litellm.litellm_core_utils.fallback_generalizations import ( match_routing_generalization, ) +from litellm.litellm_core_utils.get_model_cost_map import ( + import_model_cost_map_generation, + wait_for_import_model_cost_map, +) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.secret_managers.main import get_secret, get_secret_str @@ -186,6 +190,38 @@ def get_llm_provider( Return model, custom_llm_provider, dynamic_api_key, api_base """ + generation: Final = import_model_cost_map_generation() + resolved: Final = _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) + if resolved is not None: + return resolved + retried: Final = ( + _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) + if wait_for_import_model_cost_map(since=generation) + else None + ) + if retried is not None: + return retried + if litellm.suppress_debug_info is False: + print() # noqa: T201 + print( # noqa: T201 + "\033[1;31mProvider List: https://docs.litellm.ai/docs/providers\033[0m" + ) + print() # noqa: T201 + raise litellm.exceptions.BadRequestError( + message=f"LLM Provider NOT provided. Pass in the LLM provider you are trying to call. You passed model={model}\n Pass model as E.g. For 'Huggingface' inference endpoints pass in `completion(model='huggingface/starcoder',..)` Learn more: https://docs.litellm.ai/docs/providers", + model=model, + response=None, + llm_provider="", + ) + + +def _resolve_llm_provider( + model: str, + custom_llm_provider: str | None, + api_base: str | None, + api_key: str | None, + litellm_params: GenericLiteLLMParams | None, +) -> tuple[str, str, str | None, str | None] | None: try: # Early validation - model is required if model is None: @@ -545,20 +581,7 @@ def get_llm_provider( custom_llm_provider = match_routing_generalization(model) if not custom_llm_provider: - if litellm.suppress_debug_info is False: - print() # noqa: T201 - print( # noqa: T201 - "\033[1;31mProvider List: https://docs.litellm.ai/docs/providers\033[0m" - ) - print() # noqa: T201 - error_str = f"LLM Provider NOT provided. Pass in the LLM provider you are trying to call. You passed model={model}\n Pass model as E.g. For 'Huggingface' inference endpoints pass in `completion(model='huggingface/starcoder',..)` Learn more: https://docs.litellm.ai/docs/providers" - # maps to openai.NotFoundError, this is raised when openai does not recognize the llm - raise litellm.exceptions.BadRequestError( - message=error_str, - model=model, - response=None, - llm_provider="", - ) + return None if api_base is not None and not isinstance(api_base, str): raise Exception(f"api base needs to be a string. api_base={api_base}") if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): @@ -568,7 +591,6 @@ def get_llm_provider( if isinstance(e, litellm.exceptions.BadRequestError): raise e else: - error_str = f"GetLLMProvider Exception - {e}\n\noriginal model: {model}" raise litellm.exceptions.BadRequestError( message=f"GetLLMProvider Exception - {e}\n\noriginal model: {model}", model=model, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 7b3835dc1ac..8e275205f6c 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -18,18 +18,28 @@ import sys import threading import time import urllib.parse -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, replace +from collections.abc import Awaitable, Callable, Iterator, Mapping, MutableMapping +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from importlib.resources import files from itertools import chain from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol, cast # noqa: TID251 # adapts boto3 and botocore's dynamic client responses +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # adapts boto3 and botocore's dynamic client responses +) import httpx from pydantic import ConfigDict, TypeAdapter -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ( + NotRequired, + ReadOnly, + TypedDict, + TypeIs, # noqa: TID251 # narrows live cost map entries without copying them +) from litellm import verbose_logger from litellm.constants import ( @@ -40,22 +50,31 @@ from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +if TYPE_CHECKING: + from _collections_abc import dict_items, dict_values + FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) +_ABSENT: Final = object() +_IMPORT_FETCH_IN_PROGRESS: Final = "Remote fetch in progress" def _is_cli_process() -> bool: return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES +def _is_uwsgi_process() -> bool: + return "uwsgi" in sys.modules + + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. RESERVED_TOP_LEVEL_KEYS: Final = frozenset({"sample_spec", FALLBACK_GENERALIZATIONS_KEY}) -def _count_model_entries(model_cost: dict) -> int: +def _count_model_entries(model_cost: Mapping[str, object]) -> int: """Count actual model entries, excluding reserved meta keys.""" return sum(1 for key in model_cost if key not in RESERVED_TOP_LEVEL_KEYS) @@ -117,14 +136,14 @@ class GetModelCostMap: def get_backup_model_count(cls) -> int: """Return the number of models in the local backup (cached int).""" if cls._backup_model_count < 0: - backup: Final = cls.load_local_model_cost_map() + backup: Final = cls.load_local_model_cost_map_with_revision().model_cost_map cls._backup_model_count = _count_model_entries(backup) return cls._backup_model_count _get_backup_model_count = get_backup_model_count @staticmethod - def _check_is_valid_dict(fetched_map: dict) -> bool: + def _check_is_valid_dict(fetched_map: object) -> bool: """Check 1: fetched map is a non-empty dict.""" if not isinstance(fetched_map, dict): verbose_logger.warning( @@ -144,7 +163,7 @@ class GetModelCostMap: @classmethod def _check_model_count_not_reduced( cls, - fetched_map: dict, + fetched_map: dict[str, object], backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, @@ -179,7 +198,7 @@ class GetModelCostMap: @classmethod def validate_model_cost_map( cls, - fetched_map: dict, + fetched_map: dict[str, object], backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, @@ -222,7 +241,7 @@ def mark_litellm_import_complete() -> None: @dataclass(frozen=True, slots=True) class ModelCostMapReloaded: - model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict + model_cost_map: dict[str, object] # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract revision: str | None = None etag: str | None = None @@ -424,7 +443,7 @@ def _fetch_remote_model_cost_map_with_retry_sync( url: str, timeout: int, attempts: range, - sleep: Callable[[float], None], + wait_to_retry: Callable[[_FetchAttemptRetryable, float], None], rng: random.Random, client: _SyncGetClient, ) -> ModelCostMapReloadResult: @@ -436,7 +455,7 @@ def _fetch_remote_model_cost_map_with_retry_sync( wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) if isinstance(wait_seconds, ModelCostMapReloadUnavailable): return wait_seconds - sleep(wait_seconds) + wait_to_retry(outcome, wait_seconds) return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") @@ -568,6 +587,19 @@ class AliasedCostEntry(dict[str, object]): self.registered_fields: frozenset[str] = registered_fields +class _SnapshotIteratingDict(dict[str, object]): + __slots__ = () + + def __iter__(self) -> Iterator[str]: + return iter(tuple(super().keys())) + + def values(self) -> "dict_values[str, object]": + return dict(super().items()).values() + + def items(self) -> "dict_items[str, object]": + return dict(super().items()).items() + + def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]: return MappingProxyType( { @@ -645,12 +677,18 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) -def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: +def _record_and_finalize(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: _cost_map_source_info.source_revision = loaded.revision _cost_map_source_info.etag = loaded.etag - finalized: Final = _finalize_model_cost_map(loaded.model_cost_map) - GetModelCostMap._snapshot_loaded_catalog(finalized) # pyright: ignore[reportPrivateUsage] # same module - return replace(loaded, model_cost_map=finalized) + return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) + + +def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: + finalized: Final = _record_and_finalize(loaded) + GetModelCostMap._snapshot_loaded_catalog( # pyright: ignore[reportPrivateUsage] # same module + finalized.model_cost_map + ) + return finalized def adopt_model_cost_map( @@ -667,6 +705,245 @@ def adopt_model_cost_map( return fetched_model_count +@dataclass(frozen=True, slots=True) +class _ImportFetch: + url: str + timeout: int + max_attempts: int + sleep: Callable[[float], None] + rng: random.Random + client: _SyncGetClient + bundled: dict[str, object] # mutable-ok: litellm.model_cost itself, which the remote map is merged into + loaded: Mapping[str, Mapping[str, object]] + local: Callable[[], ModelCostMapReloaded] + clock: Callable[[], float] + first_attempt: int = 1 + first_attempt_at: float = 0.0 + first_attempt_settled: threading.Event = field(default_factory=threading.Event) + + +@dataclass(frozen=True, slots=True) +class _ScheduledAttempt: + attempt: int + at: float + + +class _PendingImportFetch: + fetch: _ImportFetch | None = None + thread: threading.Thread | None = None + pid: int = 0 + next_attempt: _ScheduledAttempt | None = None + adoptions: int = 0 + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _litellm_model_cost_is(expected: Mapping[str, object]) -> bool: + import litellm + + return litellm.model_cost is expected + + +def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: + return fetch is not None and not _litellm_model_cost_is(fetch.bundled) + + +def _report_import_fetch_failure(fetch: _ImportFetch | None, reason: str) -> None: + _litellm_import_complete.wait() + if not _model_cost_replaced_since_import(fetch): + _cost_map_source_info.fallback_reason = reason + + +def _is_cost_entry(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: cost map keys are str + return isinstance(value, Mapping) + + +def _bundled_as_loaded(fetch: _ImportFetch) -> Mapping[str, Mapping[str, object]]: + raw: Final = fetch.local().model_cost_map + catalog: Final = _CATALOG_ADAPTER.validate_python( + {key: value for key, value in raw.items() if key != FALLBACK_GENERALIZATIONS_KEY} + ) + return MappingProxyType(_expand_model_aliases(catalog)) + + +def _fields_written_since_import( + current: Mapping[str, object], loaded: Mapping[str, object], parsed: Mapping[str, object] +) -> frozenset[str]: + rewritten: Final = ( + name + for name, value in tuple(current.items()) + if loaded.get(name, _ABSENT) is not value or (isinstance(value, (dict, list)) and value != parsed.get(name)) + ) + return frozenset(chain(rewritten, loaded.keys() - current.keys())) + + +def _merged_entry( + current: Mapping[str, object], remote: Mapping[str, object], written: frozenset[str] +) -> dict[str, object]: # mutable-ok: cost map entries are dicts callers may edit + trimmed: Final = any(name not in current for name in written) + fields: Final = { + **{name: value for name, value in remote.items() if name not in written and (name in current or not trimmed)}, + **{name: value for name, value in current.items() if name in written}, + } + if isinstance(remote, AliasedCostEntry): + return AliasedCostEntry(fields, remote.names) + if isinstance(current, AliasedCostEntry): + return AliasedCostEntry(fields, current.names) + return fields + + +def _alias_names(entry: Mapping[str, object]) -> tuple[str, ...] | None: + return entry.names if isinstance(entry, AliasedCostEntry) else None + + +def _updatable_in_place(merged: Mapping[str, object], current: Mapping[str, object]) -> bool: + return ( + isinstance(current, MutableMapping) + and merged.keys() == current.keys() + and _alias_names(merged) == _alias_names(current) + ) + + +def _same_entry(merged: Mapping[str, object], current: Mapping[str, object]) -> bool: + return merged == current and _alias_names(merged) == _alias_names(current) + + +def _live_entry_if_equal(live: object, remote: object) -> object: + return live if live == remote else remote + + +def _untouched_since_import(current: object, loaded: Mapping[str, object], parsed: Mapping[str, object]) -> bool: + return _is_cost_entry(current) and not _fields_written_since_import(current, loaded, parsed) + + +def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) -> None: + import litellm + from litellm import utils + + parsed: Final = _bundled_as_loaded(fetch) + remote: Final = tuple(finalized.items()) + shared_once: Final = { + id(current): (current, key, entry) + for key, entry in remote + if key in fetch.loaded and _is_cost_entry(current := fetch.bundled.get(key)) and _is_cost_entry(entry) + } + merged: Final = { + entry_id: _merged_entry( + current, entry, _fields_written_since_import(current, fetch.loaded[key], parsed.get(key, _NO_FIELDS)) + ) + for entry_id, (current, key, entry) in shared_once.items() + } + changed: Final = { + entry_id: entry for entry_id, entry in merged.items() if not _same_entry(entry, shared_once[entry_id][0]) + } + reshaped: Final = { + entry_id: entry + for entry_id, entry in changed.items() + if not _updatable_in_place(entry, shared_once[entry_id][0]) + } + updated_in_place: Final = tuple( + (live, entry) + for entry_id, entry in changed.items() + if entry_id not in reshaped and isinstance(live := shared_once[entry_id][0], MutableMapping) + ) + replaced: Final = {key: reshaped[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in reshaped} + moved: Final = tuple( + (key, MappingProxyType(dict(value)), changed[id(value)]) + for key, value in tuple(fetch.bundled.items()) + if id(value) in changed and _is_cost_entry(value) + ) + kept_for_remote: Final = { + id(entry): reshaped.get(entry_id, current) for entry_id, (current, _, entry) in shared_once.items() + } + added: Final = { + key: kept_for_remote.get(id(entry), entry) + for key, entry in remote + if key not in fetch.loaded and key not in fetch.bundled + } + added_since_import: Final = { + key: _merged_entry(current, entry, frozenset(current.keys())) + for key, entry in remote + if key not in fetch.loaded and _is_cost_entry(current := fetch.bundled.get(key)) and _is_cost_entry(entry) + } + dropped: Final = tuple( + key + for key in fetch.loaded.keys() - finalized.keys() + if _untouched_since_import(fetch.bundled.get(key), fetch.loaded[key], parsed.get(key, _NO_FIELDS)) + ) + regrouped: Final = ( + *moved, + *((key, entry, _NO_FIELDS) for key in dropped if _is_cost_entry(entry := fetch.bundled.get(key))), + ) + for live_entry, entry in updated_in_place: + live_entry.update(entry) + fetch.bundled.update({**replaced, **added, **added_since_import}) + for key in dropped: + fetch.bundled.pop(key, None) + GetModelCostMap._snapshot_loaded_catalog( # pyright: ignore[reportPrivateUsage] # same module + {key: _live_entry_if_equal(fetch.bundled.get(key), entry) for key, entry in remote} + ) + _record_remote_adoption() + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + litellm.add_known_models_to_model_list() + litellm.drop_stale_known_models(regrouped) + utils.reapply_runtime_model_cost_registrations() + _PendingImportFetch.adoptions += 1 + + +def _adopt_remote_model_cost_map( + url: str, + timeout: int, + attempts: range, + wait_to_retry: Callable[[_FetchAttemptRetryable, float], None], + rng: random.Random, + client: _SyncGetClient, + import_fetch: _ImportFetch | None, +) -> None: + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + attempts=attempts, + wait_to_retry=wait_to_retry, + rng=rng, + client=client, + ) + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s (%s); keeping local backup", + url, + result.reason, + ) + _report_import_fetch_failure(import_fetch, f"Remote fetch failed: {result.reason}") + return + _litellm_import_complete.wait() + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap.get_backup_model_count(), + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + _report_import_fetch_failure(import_fetch, "Remote data failed integrity validation") + return + if _model_cost_replaced_since_import(import_fetch): + verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") + return + if import_fetch is None: + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _record_remote_adoption() + adopt_model_cost_map(finalized) + return + _merge_into_import_map(import_fetch, _record_and_finalize(result).model_cost_map) + + +def _record_remote_adoption() -> None: + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + + def _retry_remote_fetch_in_background( url: str, timeout: int, @@ -681,40 +958,211 @@ def _retry_remote_fetch_in_background( if isinstance(first_wait, ModelCostMapReloadUnavailable): return sleep(first_wait) - result: Final = _fetch_remote_model_cost_map_with_retry_sync( + _adopt_remote_model_cost_map( url=url, timeout=timeout, attempts=range(2, max_attempts + 1), - sleep=sleep, + wait_to_retry=lambda _, seconds: sleep(seconds), rng=rng, client=client, + import_fetch=None, ) - if isinstance(result, ModelCostMapReloadUnavailable): - verbose_logger.warning( - "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", - url, - max_attempts, - ) - return - _litellm_import_complete.wait() - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap.get_backup_model_count(), - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, - ) - return - finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map - _cost_map_source_info.source = "remote" - _cost_map_source_info.fallback_reason = None - _cost_map_source_info.loaded_at = datetime.now(timezone.utc) - adopt_model_cost_map(finalized) except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) +def _run_import_fetch(fetch: _ImportFetch) -> None: + def settle_then_sleep(failed: _FetchAttemptRetryable, seconds: float) -> None: + scheduled: Final = _PendingImportFetch.next_attempt + retry: Final = _ScheduledAttempt( + attempt=(scheduled.attempt if scheduled is not None else fetch.first_attempt) + 1, + at=fetch.clock() + seconds, + ) + _PendingImportFetch.next_attempt = retry + _report_import_fetch_failure(fetch, f"Remote fetch failed: {failed.reason}") + fetch.first_attempt_settled.set() + fetch.sleep(max(0.0, retry.at - fetch.clock())) + + try: + if fetch.first_attempt > 1: + fetch.sleep(max(0.0, fetch.first_attempt_at - fetch.clock())) + _adopt_remote_model_cost_map( + url=fetch.url, + timeout=fetch.timeout, + attempts=range(fetch.first_attempt, fetch.max_attempts + 1), + wait_to_retry=settle_then_sleep, + rng=fetch.rng, + client=fetch.client, + import_fetch=fetch, + ) + except Exception as e: # noqa: BLE001 # a failed background fetch must not kill the thread; the backup stays + verbose_logger.warning("LiteLLM: Background model cost map fetch failed: %s", e) + finally: + if _PendingImportFetch.fetch is fetch: + _PendingImportFetch.fetch = None + if ( + _model_cost_replaced_since_import(fetch) + and _cost_map_source_info.fallback_reason == _IMPORT_FETCH_IN_PROGRESS + ): + _cost_map_source_info.fallback_reason = "litellm.model_cost was replaced after import" + fetch.first_attempt_settled.set() + + +def _start_import_fetch(fetch: _ImportFetch) -> None: + thread: Final = threading.Thread( + target=_run_import_fetch, args=(fetch,), name="litellm-model-cost-map-fetch", daemon=True + ) + _PendingImportFetch.fetch = fetch + _PendingImportFetch.thread = thread + _PendingImportFetch.pid = os.getpid() + thread.start() + + +def _settled_event() -> threading.Event: + settled: Final = threading.Event() + settled.set() + return settled + + +def _resume_import_fetch_after_fork() -> None: + pending: Final = _PendingImportFetch.fetch + scheduled: Final = _PendingImportFetch.next_attempt + if pending is None: + return + if scheduled is None: + _start_import_fetch(replace(pending, first_attempt_settled=threading.Event())) + return + rng: Final = random.Random() + _start_import_fetch( + replace( + pending, + rng=rng, + first_attempt=scheduled.attempt, + first_attempt_at=scheduled.at + rng.uniform(0.0, 1.0), + first_attempt_settled=_settled_event(), + ) + ) + + +def import_model_cost_map_generation() -> int: + """How many times the remote map fetched at import was adopted; read it before a lookup that may miss.""" + return _PendingImportFetch.adoptions + + +def wait_for_import_model_cost_map(since: int) -> bool: + """ + Called by a lookup that missed the cost map, with the generation it read before looking up. + While the remote map fetched at import is still in flight, blocks until that fetch's first + attempt settles, so a model the bundled map lacks resolves the way it did when import waited + for the fetch. True means the remote map was adopted since ``since`` and the lookup should be + tried once more. + """ + wait_for_pending_import_model_cost_map() + return _PendingImportFetch.adoptions != since + + +_resume_after_unhooked_fork: Final = threading.Lock() + + +def _resume_import_fetch_after_unhooked_fork() -> None: + with _resume_after_unhooked_fork: + if _PendingImportFetch.fetch is not None and _PendingImportFetch.pid != os.getpid(): + _resume_import_fetch_after_fork() + + +def wait_for_pending_import_model_cost_map() -> None: + if _PendingImportFetch.fetch is not None and _PendingImportFetch.pid != os.getpid(): + _resume_import_fetch_after_unhooked_fork() + pending: Final = _PendingImportFetch.fetch + if ( + pending is None + or not _litellm_import_complete.is_set() + or threading.current_thread() is _PendingImportFetch.thread + ): + return + if not pending.first_attempt_settled.is_set(): + verbose_logger.debug("LiteLLM: waiting for the remote map fetched at import") + pending.first_attempt_settled.wait() + + +def _import_fetch_may_block() -> bool: + pending: Final = _PendingImportFetch.fetch + return pending is not None and ( + _PendingImportFetch.pid != os.getpid() or not pending.first_attempt_settled.is_set() + ) + + +async def await_pending_import_model_cost_map() -> None: + if _import_fetch_may_block(): + await asyncio.to_thread(wait_for_pending_import_model_cost_map) + + +def _settle_import_fetch_before_fork() -> None: + pending: Final = _PendingImportFetch.fetch + thread: Final = _PendingImportFetch.thread + if pending is None or thread is None or thread is threading.current_thread(): + return + wait_for_pending_import_model_cost_map() + if _PendingImportFetch.fetch is None: + thread.join(timeout=pending.timeout) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(before=_settle_import_fetch_before_fork, after_in_child=_resume_import_fetch_after_fork) + + +def load_model_cost_map_at_import( + url: str, + timeout: int = 5, + max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, + sleep: Callable[[float], None] = time.sleep, + rng: random.Random | None = None, + client: "_SyncGetClient | None" = None, + local: Callable[[], ModelCostMapReloaded] = GetModelCostMap.load_local_model_cost_map_with_revision, + fetch_before_returning: Callable[[], bool] = _is_uwsgi_process, + clock: Callable[[], float] = time.monotonic, +) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract + """ + Return the bundled map right away and fetch the remote one in a background thread, + so ``import litellm`` never waits on the network. Once the remote map passes the + integrity check it is merged into the returned dict in place, so references taken + at import see it too, and every field written since import keeps its value. An entry + that lost fields since import keeps its own set of fields. A caller that replaced + ``litellm.model_cost`` outright keeps its own map. Costs and model info looked up + while the first attempt is in flight wait for it, see + ``wait_for_pending_import_model_cost_map``. A process forked while the fetch is + pending, by ``os.fork`` or by a server forking at the C level, resumes it on its own. + Under uWSGI, whose workers forked while a thread holds the GIL deadlock, the remote + map is fetched before returning instead. + """ + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process() or fetch_before_returning(): + return get_model_cost_map( + url=url, timeout=timeout, max_attempts=max_attempts, sleep=sleep, rng=rng, client=client + ) + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + _cost_map_source_info.source = "local" + _cost_map_source_info.url = url + _cost_map_source_info.is_env_forced = False + _cost_map_source_info.fallback_reason = _IMPORT_FETCH_IN_PROGRESS + bundled: Final = _SnapshotIteratingDict(_finalize_loaded_model_cost_map(local()).model_cost_map) + _PendingImportFetch.next_attempt = None + _start_import_fetch( + _ImportFetch( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=rng if rng is not None else random.Random(), + client=client if client is not None else httpx, + bundled=bundled, + loaded=GetModelCostMap.loaded_model_cost_map(), + local=local, + clock=clock, + ) + ) + return bundled + + def get_model_cost_map( url: str, timeout: int = 5, diff --git a/litellm/utils.py b/litellm/utils.py index 7de7ce6188d..25cc3615b6f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -91,7 +91,11 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.litellm_core_utils.get_model_cost_map import ( + AliasedCostEntry, + await_pending_import_model_cost_map, + wait_for_pending_import_model_cost_map, +) from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -2158,6 +2162,7 @@ def client(original_function): @wraps(original_function) async def wrapper_async(*args, **kwargs): + await await_pending_import_model_cost_map() print_args_passed_to_litellm(original_function, args, kwargs) start_time: Final = datetime.datetime.now() call_type: Final = _get_call_type_from_public_name(original_function.__name__) @@ -6582,6 +6587,7 @@ def get_model_info_helper( """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's """ + wait_for_pending_import_model_cost_map() from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider try: @@ -8513,7 +8519,7 @@ def get_valid_models( Returns: A list of valid LLMs """ - + wait_for_pending_import_model_cost_map() try: ################################ # init litellm_params diff --git a/tests/integration/routing/test_stale_cost_map_boot.py b/tests/integration/routing/test_stale_cost_map_boot.py index 8ae7290edbe..7a2e50c6ea7 100644 --- a/tests/integration/routing/test_stale_cost_map_boot.py +++ b/tests/integration/routing/test_stale_cost_map_boot.py @@ -1,5 +1,7 @@ +import functools import json import threading +import time import uuid from pathlib import Path from typing import Final @@ -7,7 +9,7 @@ from typing import Final import httpx import pytest from integration._support.client import Gateway, eventually -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -86,3 +88,50 @@ def test_config_deployment_dropped_by_stale_boot_cost_map_is_restored_after_relo "body": {"model": model, "messages": [{"role": "user", "content": "stale cost map control"}]}, } ] + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_bare_model_only_the_remote_map_knows_is_served_once_a_rate_limited_first_download_retries( + gateway: Gateway, tmp_path: Path +) -> None: + model: Final = "integration-remote-only-" + uuid.uuid4().hex + remote_map: Final = json.dumps( + {model: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0, "output_cost_per_token": 0}} + ).encode() + + @functools.cache + def first_download_at() -> float: + return time.monotonic() + + def respond(request: Request) -> Reply: + assert request.target == "/model_prices.json", request + if time.monotonic() < first_download_at() + 55: + return Reply(status=429, headers={"Retry-After": "30"}) + return Reply(body=remote_map) + + overrides: Final = {"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"} + with ( + wire_server(respond) as peer, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + owned_proxy( + gateway, + tmp_path, + {**overrides, "LITELLM_MODEL_COST_MAP_URL": peer.url + "/model_prices.json"}, + config=_proxy_config(tmp_path, model, gateway.upstream_url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + workers=2, + ) as candidate, + ): + upstream.get("/__observations").raise_for_status() + body: Final = {"model": model, "messages": [{"role": "user", "content": "rate-limited cost map control"}]} + eventually( + lambda: candidate.request("POST", "/v1/chat/completions", body).status_code, + lambda status: status == 200, + seconds=90, + ) + statuses: Final = tuple(candidate.request("POST", "/v1/chat/completions", body).status_code for _ in range(6)) + assert statuses == (200,) * 6 + forwarded: Final = tuple( + (observed["path"], observed["body"]) for observed in upstream.get("/__observations").json()["requests"] + ) + assert forwarded == (("/v1/chat/completions", body),) * 7 diff --git a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py new file mode 100644 index 00000000000..f2b295c0289 --- /dev/null +++ b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py @@ -0,0 +1,221 @@ +import json +import os +import subprocess +import sys +import textwrap +import threading +import time +from collections.abc import Iterator +from queue import SimpleQueue +from dataclasses import dataclass +from importlib.resources import files +from pathlib import Path +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark: Final = pytest.mark.timeout(300) + +_REMOTE_ONLY_MODEL: Final = "zz-remote-only-sdk-model" +_CHILD_SECONDS: Final = 120 + +_SCRIPT: Final = textwrap.dedent( + """ + import os, sys, threading, time + import litellm + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info + + variant, model = sys.argv[1:3] + + def join_fetch(): + for thread in threading.enumerate(): + if thread.name == "litellm-model-cost-map-fetch": + thread.join(timeout=60) + + info = get_model_cost_map_source_info() + print(info["source"], info["fallback_reason"], model in litellm.model_cost, flush=True) + if variant == "adopt": + sys.stdin.readline() + join_fetch() + print(get_model_cost_map_source_info()["source"], model in litellm.model_cost, flush=True) + elif variant == "get_llm_provider": + print(litellm.get_llm_provider(model)[1], flush=True) + elif variant == "get_model_info": + print(litellm.get_model_info(model, custom_llm_provider="openai")["litellm_provider"], flush=True) + elif variant == "fork": + pid = os.fork() + if pid == 0: + join_fetch() + os._exit(0 if model in litellm.model_cost else 1) + print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) + elif variant == "raw_fork": + import ctypes, signal + pid = ctypes.PyDLL(None).fork() + if pid == 0: + signal.alarm(60) + try: + print(litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0], flush=True) + os._exit(0) + finally: + os._exit(1) + print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) + elif variant == "forks_during_retry_wait": + while get_model_cost_map_source_info()["fallback_reason"] == "Remote fetch in progress": + time.sleep(0.01) + children = [] + for _ in range(6): + pid = os.fork() + if pid == 0: + join_fetch() + os._exit(0 if model in litellm.model_cost else 1) + children.append(pid) + join_fetch() + print(*(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]) for pid in children), flush=True) + else: + raise SystemExit(f"unknown variant {variant}") + """ +) + + +def _bundled_map() -> dict: + return json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + + +def _remote_gpt_4o_input_cost() -> float: + return 3 * _bundled_map()["gpt-4o"]["input_cost_per_token"] + + +def _remote_map() -> bytes: + bundled: Final = _bundled_map() + return json.dumps( + { + **bundled, + "gpt-4o": {**bundled["gpt-4o"], "input_cost_per_token": _remote_gpt_4o_input_cost()}, + _REMOTE_ONLY_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + } + ).encode() + + +@dataclass(frozen=True, slots=True) +class _CostMapHost: + wire: Wire + first_request_answered: threading.Event + + +@pytest.fixture +def cost_map_host() -> Iterator[_CostMapHost]: + body: Final = _remote_map() + gate: Final = threading.Event() + first_seen: Final = threading.Event() + + def respond(request: Request) -> Reply: + if not first_seen.is_set(): + first_seen.set() + gate.wait(timeout=_CHILD_SECONDS) + return Reply(body=body) + + with wire_server(respond) as wire: + try: + yield _CostMapHost(wire=wire, first_request_answered=gate) + finally: + gate.set() + + +@dataclass(frozen=True, slots=True) +class _Child: + after_import: str + rest: tuple[str, ...] + + +def _run_child(host: _CostMapHost, variant: str, tmp_path: Path, answer_after_import: bool) -> _Child: + stderr_path: Final = tmp_path / "stderr.log" + with ( + stderr_path.open("w") as stderr, + subprocess.Popen( + [sys.executable, "-P", "-c", _SCRIPT, variant, _REMOTE_ONLY_MODEL], + env={ + **os.environ, + "LITELLM_LOCAL_MODEL_COST_MAP": "False", + "NO_PROXY": "*", + "LITELLM_MODEL_COST_MAP_URL": f"{host.wire.url}/model_prices_and_context_window.json", + }, + cwd=tmp_path, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=stderr, + text=True, + ) as child, + ): + assert child.stdout is not None + after_import: Final = child.stdout.readline().strip() + if answer_after_import: + host.first_request_answered.set() + rest: Final = child.communicate(input="\n", timeout=_CHILD_SECONDS)[0] + assert child.returncode == 0, stderr_path.read_text() + return _Child(after_import=after_import, rest=tuple(rest.split())) + + +def test_import_returns_before_the_remote_map_arrives_and_adopts_it_after( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "adopt", tmp_path, answer_after_import=True) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("remote", "True") + + +@pytest.mark.parametrize("lookup", ["get_llm_provider", "get_model_info"]) +def test_a_lookup_right_after_import_finds_a_model_only_the_remote_map_has( + cost_map_host: _CostMapHost, tmp_path: Path, lookup: str +) -> None: + child: Final = _run_child(cost_map_host, lookup, tmp_path, answer_after_import=True) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("openai",) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_worker_forked_while_the_remote_map_is_in_flight_fetches_it_itself( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "fork", tmp_path, answer_after_import=False) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("0",) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_worker_forked_without_the_fork_hooks_while_the_remote_map_is_in_flight_prices_from_it( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "raw_fork", tmp_path, answer_after_import=False) + + assert child.after_import == "local Remote fetch in progress False" + assert float(child.rest[0]) == pytest.approx(1000 * _remote_gpt_4o_input_cost()) + assert child.rest[1:] == ("0",) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_workers_forked_during_a_rate_limited_retry_wait_spread_their_retries(tmp_path: Path) -> None: + body: Final = _remote_map() + arrivals: Final[SimpleQueue[float]] = SimpleQueue() + + def respond(request: Request) -> Reply: + arrivals.put(time.monotonic()) + if arrivals.qsize() == 1: + return Reply(status=429, headers={"Retry-After": "2"}) + return Reply(body=body) + + with wire_server(respond) as wire: + child: Final = _run_child( + _CostMapHost(wire=wire, first_request_answered=threading.Event()), + "forks_during_retry_wait", + tmp_path, + answer_after_import=False, + ) + retries: Final = sorted(tuple(arrivals.get_nowait() for _ in range(arrivals.qsize()))[1:]) + + assert child.rest == ("0",) * 6 + assert len(retries) == 7 + assert retries[5] - retries[0] > 0.1 diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 84e785ac99a..66b352befd0 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -5,6 +5,7 @@ count actual model entries, not reserved meta keys) and the extraction of the """ import asyncio, importlib, importlib.resources as importlib_get_model, json, litellm +import logging import os import sys import threading @@ -602,17 +603,19 @@ def test_boot_load_success_does_not_start_background_retry(): assert get_model_cost_map_source_info()["source"] == "remote" -def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): +@pytest.fixture +def isolated_litellm_cost_state(monkeypatch): import litellm from litellm import utils as litellm_utils from litellm.litellm_core_utils import get_model_cost_map as module - original_model_cost = litellm.model_cost - monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) for name, provider_models in tuple(vars(litellm).items()): if name.endswith("_models") and isinstance(provider_models, set): monkeypatch.setattr(litellm, name, set(provider_models)) monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + monkeypatch.setattr(litellm, "model_list", list(litellm.model_list)) + monkeypatch.setattr(litellm, "model_list_set", set(litellm.model_list_set)) monkeypatch.setattr( litellm_utils, "_runtime_registered_model_cost", @@ -621,6 +624,12 @@ def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_ source_info = module._cost_map_source_info for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): monkeypatch.setattr(source_info, name, getattr(source_info, name)) + for name in ("_loaded_catalog", "_loaded_catalog_fields"): + monkeypatch.setattr(GetModelCostMap, name, getattr(GetModelCostMap, name)) + + +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(isolated_litellm_cost_state): + import litellm remote_map = _load_root_cost_map() remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} @@ -798,12 +807,1074 @@ def test_boot_load_skips_remote_fetch_for_cli_processes( assert source["source"] == "remote" +def _fetch_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-fetch"] + + +def _held_client(outcome): + """Serves `outcome` only after `release` is set, so a test can observe the caller before the fetch finishes.""" + started = threading.Event() + release = threading.Event() + + def handler(request): + started.set() + return outcome if release.wait(timeout=10) else httpx.Response(503) + + return httpx.Client(transport=httpx.MockTransport(handler)), started, release + + +def _rate_limited_then_held_client(body: bytes): + retry_started = threading.Event() + release = threading.Event() + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + return httpx.Response(429) + retry_started.set() + release.wait(timeout=10) + return httpx.Response(200, content=body) + + return httpx.Client(transport=httpx.MockTransport(handler)), retry_started, release + + +def _finish_fetch(release): + release.set() + threads = _fetch_threads() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + + +def _remote_map_with(model: str) -> bytes: + remote_map = _load_root_cost_map() + remote_map[model] = {"litellm_provider": "anthropic", "mode": "chat"} + return json.dumps(remote_map).encode() + + +def test_import_load_returns_the_bundled_map_before_the_remote_fetch_answers_then_adopts_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + cost_map = load_model_cost_map_at_import(url=_URL, client=client) + assert started.wait(timeout=10) + assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] == "Remote fetch in progress" + monkeypatch.setattr(litellm, "model_cost", cost_map) + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + assert "claude-remote-only-test" in litellm.model_list + assert "claude-remote-only-test" in litellm.model_list_set + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + + +def test_import_load_keeps_custom_pricing_and_edits_made_before_the_remote_map_lands( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + litellm.register_model( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 42.0}}, + persist_across_reloads=False, + ) + litellm.model_cost["my-direct-edit"] = { + "litellm_provider": "openai", + "mode": "chat", + } # test-quality-ok: the edit is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 42.0 + assert "my-direct-edit" in litellm.model_cost + assert get_model_cost_map_source_info()["source"] == "remote" + + +@pytest.mark.parametrize("remote_response", [httpx.Response(200, content=b""), httpx.Response(404)], ids=["ok", "404"]) +def test_import_load_leaves_a_model_cost_the_caller_replaced_and_its_status_alone( + isolated_litellm_cost_state, monkeypatch, remote_response +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + if remote_response.status_code == 200: + remote_response = httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + client, started, release = _held_client(remote_response) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + reloaded_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} + monkeypatch.setattr(litellm, "model_cost", reloaded_map) + monkeypatch.setattr(module._cost_map_source_info, "source", "remote") + monkeypatch.setattr(module._cost_map_source_info, "fallback_reason", None) + finally: + _finish_fetch(release) + + assert litellm.model_cost is reloaded_map + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + + +def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, _, release = _held_client(httpx.Response(404)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + finally: + _finish_fetch(release) + + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") + + +def test_import_load_under_uwsgi_returns_the_remote_map_once_fetched(isolated_litellm_cost_state): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + fetched_on = [] + + def handler(request): + fetched_on.append(threading.current_thread()) + return httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + loaded = load_model_cost_map_at_import(url=_URL, client=client, fetch_before_returning=lambda: True) + + assert fetched_on == [threading.current_thread()] + assert "claude-remote-only-test" in loaded + assert get_model_cost_map_source_info()["source"] == "remote" + + +def test_import_load_reports_why_the_last_attempt_failed_while_a_retry_waits(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module, "_litellm_import_complete", threading.Event()) + sleeps = _BlockingSleeps() + client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(429))) + try: + bundled = module.load_model_cost_map_at_import(url=_URL, client=client, sleep=sleeps, rng=random.Random(0)) + retry_wait_began_during_import = sleeps.called.acquire(timeout=1) + monkeypatch.setattr(litellm, "model_cost", bundled) + module.mark_litellm_import_complete() + assert retry_wait_began_during_import or sleeps.called.acquire(timeout=10) + source = get_model_cost_map_source_info() + finally: + _finish_fetch(sleeps.release) + + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 429") + + +def test_a_retry_scheduled_while_import_is_still_running_is_not_pushed_back_by_the_import( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module, "_litellm_import_complete", threading.Event()) + sleeps = _BlockingSleeps() + clock = _SteppedClock() + client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(429, headers={"Retry-After": "2"})) + ) + try: + bundled = module.load_model_cost_map_at_import( + url=_URL, client=client, sleep=sleeps, rng=random.Random(0), clock=clock + ) + monkeypatch.setattr(litellm, "model_cost", bundled) + assert clock.read.wait(timeout=10) + clock.now = 0.5 + module.mark_litellm_import_complete() + assert sleeps.called.acquire(timeout=10) + finally: + _finish_fetch(sleeps.release) + + assert sleeps.waits[0] == pytest.approx(1.5) + + +class _SteppedClock: + def __init__(self): + self.now = 0.0 + self.read = threading.Event() + + def __call__(self) -> float: + now = self.now + self.read.set() + return now + + +def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_one_hangs( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + inherited_started = threading.Event() + release = threading.Event() + body = _remote_map_with("claude-remote-only-test") + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + inherited_started.set() + release.wait(timeout=10) + return httpx.Response(200, content=body) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert inherited_started.wait(timeout=10) + inherited_thread = module._PendingImportFetch.thread + module._resume_import_fetch_after_fork() + resumed_thread = module._PendingImportFetch.thread + assert resumed_thread is not inherited_thread + resumed_thread.join(timeout=10) + assert not resumed_thread.is_alive() + assert "claude-remote-only-test" in litellm.model_cost + assert inherited_thread.is_alive() + finally: + _finish_fetch(release) + + assert len(served) == 2 + + +def _remote_map_repricing_gpt_4o() -> tuple[bytes, float]: + remote_map = _load_root_cost_map() + remote_price = remote_map["gpt-4o"]["input_cost_per_token"] * 10 + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "input_cost_per_token": remote_price} + return json.dumps(remote_map).encode(), remote_price + + +def _remote_map_repricing_gpt_4o_and_adding(model: str) -> tuple[bytes, float]: + body, remote_price = _remote_map_repricing_gpt_4o() + remote_map = json.loads(body) + remote_map[model] = {"litellm_provider": "anthropic", "mode": "chat"} + return json.dumps(remote_map).encode(), remote_price + + +def test_references_to_the_map_taken_at_import_see_the_remote_map_once_it_lands( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o_and_adding("claude-remote-only-test") + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + held_map = litellm.model_cost + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + held_map["my-direct-edit"] = {"litellm_provider": "openai", "mode": "chat"} + assert held_map["gpt-4o"]["input_cost_per_token"] == remote_price + assert "claude-remote-only-test" in held_map + assert "my-direct-edit" in litellm.model_cost + + +def test_a_model_entry_held_from_import_sees_the_remote_price_and_edits_through_it_reach_the_map( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + held_entry = litellm.model_cost["gpt-4o"] + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + held_entry["output_cost_per_token"] = 9.0 + assert held_entry["input_cost_per_token"] == remote_price + assert litellm.model_cost["gpt-4o"]["output_cost_per_token"] == 9.0 + + +def test_a_model_entry_being_read_field_by_field_when_the_remote_map_lands_reads_to_the_end( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "supports_remote_only_test_field": True} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_fields = len(litellm.model_cost["gpt-4o"]) + fields = iter(litellm.model_cost["gpt-4o"].items()) + first = next(fields) + finally: + _finish_fetch(release) + + assert len([first, *fields]) == bundled_fields + assert litellm.model_cost["gpt-4o"]["supports_remote_only_test_field"] is True + + +@pytest.mark.parametrize( + "start_loop", + [ + pytest.param(iter, id="models"), + pytest.param(lambda model_cost: iter(model_cost.values()), id="values"), + pytest.param(lambda model_cost: iter(model_cost.items()), id="items"), + ], +) +def test_a_loop_over_the_map_running_when_the_remote_map_lands_reads_to_the_end( + isolated_litellm_cost_state, monkeypatch, start_loop +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_models = len(litellm.model_cost) + loop = start_loop(litellm.model_cost) + first = next(loop) + finally: + _finish_fetch(release) + + assert len([first, *loop]) == bundled_models + assert "claude-remote-only-test" in litellm.model_cost + + +def test_an_alias_the_remote_map_no_longer_lists_stops_resolving(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = json.loads(_cost_map_with_aliased_model(5.0)) + del remote_map["my-aliased-model"]["aliases"] + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, + client=client, + local=lambda: ModelCostMapReloaded(model_cost_map=json.loads(_cost_map_with_aliased_model(1.0))), + ), + ) + assert started.wait(timeout=10) + assert "my-aliased-model-alias" in litellm.model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["my-aliased-model"]["input_cost_per_token"] == 5.0 + assert "my-aliased-model-alias" not in litellm.model_cost + + +def test_a_price_pinned_to_its_bundled_value_before_the_remote_map_lands_keeps_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_price = litellm.model_cost["gpt-4o"]["input_cost_per_token"] + litellm.model_cost["gpt-4o"]["input_cost_per_token"] = float(repr(bundled_price)) + finally: + _finish_fetch(release) + + assert remote_price != bundled_price + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == bundled_price + + +def test_a_model_cost_the_caller_replaces_while_the_remote_map_is_merged_stays_theirs( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + caller_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} + local_loads = [] + + def local_replacing_the_map_once_the_merge_starts(): + local_loads.append(None) + if len(local_loads) == 2: + monkeypatch.setattr(litellm, "model_cost", caller_map) + return GetModelCostMap.load_local_model_cost_map_with_revision() + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import(url=_URL, client=client, local=local_replacing_the_map_once_the_merge_starts), + ) + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + assert len(local_loads) == 2 + assert litellm.model_cost is caller_map + assert "claude-remote-only-test" not in litellm.model_cost + + +class _BlockingSleeps: + def __init__(self): + self.waits = [] + self.called = threading.Semaphore(0) + self.release = threading.Event() + + def __call__(self, seconds: float) -> None: + self.waits.append(seconds) + self.called.release() + self.release.wait(timeout=10) + + +def test_a_fetch_resumed_after_fork_during_a_retry_wait_waits_out_the_rest_of_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + return httpx.Response(429, headers={"Retry-After": "30"}) + return httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + + sleeps = _BlockingSleeps() + try: + monkeypatch.setattr( + litellm, + "model_cost", + module.load_model_cost_map_at_import( + url=_URL, + client=httpx.Client(transport=httpx.MockTransport(handler)), + sleep=sleeps, + rng=random.Random(0), + ), + ) + assert sleeps.called.acquire(timeout=10) + module._resume_import_fetch_after_fork() + assert sleeps.called.acquire(timeout=10) + requests_during_the_wait = len(served) + finally: + _finish_fetch(sleeps.release) + + parent_wait, resumed_wait = sleeps.waits[:2] + assert requests_during_the_wait == 1 + assert 0 < resumed_wait <= parent_wait + 1.0 + assert "claude-remote-only-test" in litellm.model_cost + + +def test_a_router_built_before_the_remote_map_lands_prices_its_models_from_the_remote_map( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, retry_started, release = _rate_limited_then_held_client(body) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + } + ] + ) + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == remote_price + assert litellm.get_model_info("openai/gpt-4o")["input_cost_per_token"] == remote_price + + +def test_a_router_built_before_the_remote_map_lands_serves_a_bare_model_only_the_remote_map_knows( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, retry_started, release = _rate_limited_then_held_client(_remote_map_with("claude-remote-only-test")) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + router = litellm.Router( + model_list=[ + { + "model_name": "remote-only", + "litellm_params": {"model": "claude-remote-only-test", "api_key": "sk-test"}, + } + ], + ignore_invalid_deployments=True, + ) + served_before_the_remote_map = len(router.model_list) + finally: + _finish_fetch(release) + + assert served_before_the_remote_map == 0 + assert [deployment["model_name"] for deployment in router.model_list] == ["remote-only"] + assert litellm.get_model_info("claude-remote-only-test")["litellm_provider"] == "anthropic" + + +def _bundled_chat_model_of(provider: str) -> str: + return min( + key + for key, entry in _load_root_cost_map().items() + if entry.get("litellm_provider") == provider and entry.get("mode") == "chat" + ) + + +def _remote_map_editing(model: str, fields: dict) -> bytes: + remote_map = _load_root_cost_map() + remote_map[model] = {**remote_map[model], **fields} + return json.dumps(remote_map).encode() + + +def _land_remote_map(monkeypatch, body: bytes) -> None: + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, retry_started, release = _rate_limited_then_held_client(body) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + finally: + _finish_fetch(release) + + +@pytest.mark.parametrize( + "remote_fields, provider_after", + ( + ({"litellm_provider": "text-completion-openai"}, "text-completion-openai"), + ({"input_cost_per_token": 1.0}, "openai"), + ), + ids=("moved", "repriced"), +) +def test_a_bare_model_the_remote_map_edits_resolves_to_the_remote_provider( + isolated_litellm_cost_state, monkeypatch, remote_fields, provider_after +): + import litellm + + model = _bundled_chat_model_of("openai") + assert litellm.get_llm_provider(model)[1] == "openai" + + _land_remote_map(monkeypatch, _remote_map_editing(model, remote_fields)) + + assert litellm.get_llm_provider(model)[1] == provider_after + + +def test_an_azure_ai_deployment_of_a_model_the_remote_map_drops_stays_on_azure_ai( + isolated_litellm_cost_state, monkeypatch +): + import litellm + + model = _bundled_chat_model_of("openai") + deployment = f"azure_ai/{model}" + api_base = "https://example.invalid" + assert litellm.get_llm_provider(deployment, api_base=api_base)[1] == "azure" + remote_map = _load_root_cost_map() + del remote_map[model] + + _land_remote_map(monkeypatch, json.dumps(remote_map).encode()) + + assert litellm.get_llm_provider(deployment, api_base=api_base)[1] == "azure_ai" + + +def test_a_model_the_remote_map_moves_to_another_provider_leaves_the_old_provider_wildcard_list( + isolated_litellm_cost_state, monkeypatch +): + import litellm + + model = _bundled_chat_model_of("azure") + assert model in litellm.models_by_provider["azure"] + + _land_remote_map(monkeypatch, _remote_map_editing(model, {"litellm_provider": "openai"})) + + assert model not in litellm.models_by_provider["azure"] + assert model in litellm.models_by_provider["openai"] + + +def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_it_changed( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + edited = { + name: value for name, value in litellm.model_cost["gpt-4o"].items() if name != "output_cost_per_token" + } + litellm.model_cost["gpt-4o"] = { + **edited, + "max_tokens": 7, + } # test-quality-ok: the edit is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == remote_price + assert litellm.model_cost["gpt-4o"]["max_tokens"] == 7 + assert "output_cost_per_token" not in litellm.model_cost["gpt-4o"] + + +def test_a_model_entry_replaced_by_one_with_fewer_fields_before_the_remote_map_lands_stays_as_written( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "output_cost_per_reasoning_token": 1.0} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + replacement = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + litellm.model_cost["gpt-4o"] = dict( + replacement + ) # test-quality-ok: the replacement is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"] == replacement + + +def test_a_model_deleted_before_the_remote_map_lands_stays_deleted(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + del litellm.model_cost["gpt-4o-mini"] # test-quality-ok: the deletion is under test; the fixture restores it + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert "gpt-4o-mini" not in litellm.model_cost + + +def _nested_fields(cost_map): + for model, entry in cost_map.items(): + for field, value in entry.items(): + if isinstance(value, dict): + yield model, field + + +def test_a_nested_edit_made_before_the_remote_map_lands_survives_it(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + model, field = next(_nested_fields(litellm.model_cost)) + litellm.model_cost[model][field][ + "nested-edit-test" + ] = 99.0 # test-quality-ok: the edit is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert litellm.model_cost[model][field]["nested-edit-test"] == 99.0 + + +def test_import_load_reports_that_the_caller_replaced_the_model_cost_it_returned( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + monkeypatch.setattr(litellm, "model_cost", {"my-model": {"litellm_provider": "openai", "mode": "chat"}}) + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" not in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] == "litellm.model_cost was replaced after import" + + +def _cost_map_with_aliased_model(input_cost_per_token: float) -> bytes: + cost_map = _load_root_cost_map() + cost_map["my-aliased-model"] = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": input_cost_per_token, + "output_cost_per_token": 1.0, + "aliases": ["my-aliased-model-alias"], + } + return json.dumps(cost_map).encode() + + +def test_an_edit_to_an_aliased_model_made_before_the_remote_map_lands_stays_shared_by_its_aliases( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, retry_started, release = _rate_limited_then_held_client(_cost_map_with_aliased_model(5.0)) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, + client=client, + sleep=lambda _: None, + local=lambda: ModelCostMapReloaded(model_cost_map=json.loads(_cost_map_with_aliased_model(1.0))), + ), + ) + assert retry_started.wait(timeout=10) + litellm.register_model({"my-aliased-model": {"output_cost_per_token": 9.0}}, persist_across_reloads=False) + finally: + _finish_fetch(release) + + entry = litellm.model_cost["my-aliased-model"] + assert litellm.model_cost["my-aliased-model-alias"] is entry + assert entry["input_cost_per_token"] == 5.0 + assert entry["output_cost_per_token"] == 9.0 + + +def test_an_alias_the_remote_map_adds_shares_the_prices_its_model_kept(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + bundled_map = json.loads(_cost_map_with_aliased_model(1.0)) + del bundled_map["my-aliased-model"]["aliases"] + client, started, release = _held_client(httpx.Response(200, content=_cost_map_with_aliased_model(5.0))) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, client=client, local=lambda: ModelCostMapReloaded(model_cost_map=bundled_map) + ), + ) + assert started.wait(timeout=10) + litellm.register_model({"my-aliased-model": {"output_cost_per_token": 9.0}}, persist_across_reloads=False) + finally: + _finish_fetch(release) + + assert litellm.model_cost["my-aliased-model-alias"]["output_cost_per_token"] == 9.0 + assert litellm.model_cost["my-aliased-model-alias"]["input_cost_per_token"] == 5.0 + litellm.register_model({"my-aliased-model": {"input_cost_per_token": 7.0}}, persist_across_reloads=False) + assert litellm.model_cost["my-aliased-model-alias"]["input_cost_per_token"] == 7.0 + + +class _ReleaseWhenALookupWaits(logging.Handler): + """Lets the held fetch answer only once a lookup is blocked waiting on it, so the lookup provably missed first.""" + + def __init__(self, release: threading.Event): + super().__init__(level=logging.DEBUG) + self.fetch_release = release + self.waited = threading.Event() + + def emit(self, record): + if "waiting for the remote map fetched at import" in record.getMessage(): + self.waited.set() + self.fetch_release.set() + + +@pytest.fixture() +def release_when_a_lookup_waits(): + from litellm import verbose_logger + + installed = [] + + def install(release: threading.Event) -> _ReleaseWhenALookupWaits: + handler = _ReleaseWhenALookupWaits(release) + verbose_logger.addHandler(handler) + installed.append(handler) + return handler + + previous_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) + yield install + verbose_logger.setLevel(previous_level) + for handler in installed: + verbose_logger.removeHandler(handler) + + +@pytest.mark.parametrize( + "lookup", + [ + pytest.param(lambda: litellm.get_llm_provider("zz-remote-only-test")[1], id="get_llm_provider"), + pytest.param( + lambda: litellm.get_model_info("zz-remote-only-test", custom_llm_provider="anthropic")["litellm_provider"], + id="get_model_info", + ), + ], +) +def test_a_lookup_that_misses_the_bundled_map_waits_for_the_remote_map_fetched_at_import( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits, lookup +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("zz-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + assert "zz-remote-only-test" not in litellm.model_cost + provider = lookup() + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert provider == "anthropic" + + +def test_a_lookup_that_misses_both_maps_still_raises_after_waiting( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("zz-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + with pytest.raises(litellm.exceptions.BadRequestError, match="LLM Provider NOT provided"): + litellm.get_llm_provider("zz-in-neither-map") + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + + +@pytest.mark.parametrize("remote_status", [200, 404], ids=["adopted", "failed"]) +def test_a_lookup_that_missed_before_the_fetch_settled_retries_only_once_the_remote_map_was_adopted( + isolated_litellm_cost_state, monkeypatch, remote_status +): + from litellm.litellm_core_utils import get_model_cost_map as module + + body = _remote_map_with("zz-remote-only-test") if remote_status == 200 else b"" + client, started, release = _held_client(httpx.Response(remote_status, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + generation_at_the_miss = module.import_model_cost_map_generation() + finally: + _finish_fetch(release) + + assert module._PendingImportFetch.fetch is None + assert module.wait_for_import_model_cost_map(since=generation_at_the_miss) is (remote_status == 200) + assert module.wait_for_import_model_cost_map(since=module.import_model_cost_map_generation()) is False + + +@pytest.mark.parametrize( + "price_a_thousand_prompt_tokens", + [ + pytest.param( + lambda: litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0], + id="cost_per_token", + ), + pytest.param( + lambda: litellm.completion_cost( + completion_response=litellm.ModelResponse( + model="gpt-4o", usage=litellm.Usage(prompt_tokens=1000, completion_tokens=0, total_tokens=1000) + ), + model="gpt-4o", + ), + id="completion_cost", + ), + pytest.param( + lambda: 1000 * litellm.get_model_info("gpt-4o")["input_cost_per_token"], + id="get_model_info", + ), + ], +) +def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits, price_a_thousand_prompt_tokens +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + cost = price_a_thousand_prompt_tokens() + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(1000 * remote_price) + + +async def test_an_async_call_priced_while_the_remote_map_is_in_flight_keeps_the_event_loop_running( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_price = remote_map["text-embedding-3-small"]["input_cost_per_token"] * 10 + remote_map["text-embedding-3-small"] = {**remote_map["text-embedding-3-small"], "input_cost_per_token": remote_price} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + + def embed(): + return litellm.aembedding(model="text-embedding-3-small", input=["hi"], mock_response=[0.1]) + + await embed() + + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, max_attempts=1) + ) + assert started.wait(timeout=10) + lookup = release_when_a_lookup_waits(threading.Event()) + priced = asyncio.ensure_future(embed()) + assert await asyncio.to_thread(lookup.waited.wait, 10) + release.set() + response = await priced + finally: + _finish_fetch(release) + + assert response._hidden_params["response_cost"] == pytest.approx(10 * remote_price) + + +def test_a_cost_computed_while_the_first_attempt_outlasts_the_fetch_timeout_still_uses_the_remote_price( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, timeout=0)) + assert started.wait(timeout=10) + cost = litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0] + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(1000 * remote_price) + + +def test_an_image_priced_while_the_remote_map_is_in_flight_finds_a_model_only_the_remote_map_has( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + from litellm.types.utils import ImageObject, ImageResponse + + remote_map = _load_root_cost_map() + remote_map["zz-remote-only-image-test"] = { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image": 0.25, + } + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + cost = litellm.completion_cost( + completion_response=ImageResponse(created=1, data=[ImageObject(url="https://example.com/image.png")]), + model="zz-remote-only-image-test", + custom_llm_provider="openai", + call_type="image_generation", + ) + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(0.25) + + +def test_models_listed_while_the_remote_map_is_in_flight_include_a_model_only_the_remote_map_has( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + valid_models = litellm.get_valid_models(custom_llm_provider="anthropic") + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert "claude-remote-only-test" in valid_models + + +def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + def _fail(request): + raise AssertionError("no HTTP request should be made when local map is forced") + + cost_map = load_model_cost_map_at_import(url=_URL, client=httpx.Client(transport=httpx.MockTransport(_fail))) + + assert cost_map + assert _fetch_threads() == [] + assert get_model_cost_map_source_info()["is_env_forced"] is True + + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -813,6 +1884,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -846,6 +1918,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -860,6 +1933,7 @@ _SCALAR_DEFAULTS = { "cohere_key": getattr(litellm, "cohere_key", None), } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckIsValidDict: """Unit tests for _check_is_valid_dict.""" @@ -884,6 +1958,7 @@ class TestCheckIsValidDict: """Non-empty dict should pass.""" assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckModelCountNotReduced: """Unit tests for _check_model_count_not_reduced.""" @@ -938,6 +2013,7 @@ class TestCheckModelCountNotReduced: is True ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestValidateModelCostMap: """Unit tests for validate_model_cost_map (combines both checks).""" @@ -974,6 +2050,7 @@ class TestValidateModelCostMap: is True ) + def _fetch_with_single_outcome(monkeypatch, outcome): from litellm.litellm_core_utils import get_model_cost_map as module @@ -1049,6 +2126,7 @@ class TestGetModelCostMapFallback: assert isinstance(result, dict) assert len(result) > 0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBackupModelCostMapExists: """Validates the local backup file is always present and valid.""" @@ -1064,6 +2142,7 @@ class TestBackupModelCostMapExists: backup = GetModelCostMap.load_local_model_cost_map() assert len(backup) > 100, f"Backup has only {len(backup)} models, expected > 100" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBadHostedModelCostMap: """ @@ -1128,12 +2207,14 @@ class TestBadHostedModelCostMap: finally: litellm.model_cost = original + @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -1186,6 +2267,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS_local_testing = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -1206,6 +2288,7 @@ _SCALAR_DEFAULTS_local_testing = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -1228,6 +2311,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -1239,6 +2323,7 @@ def test_get_model_cost_map(): except Exception as e: pytest.fail(f"An exception occurred: {e}") + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 701b79d0ba9..11b735fc398 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -719,7 +719,7 @@ def unregistered_providers(rows: Mapping[str, object]) -> list[str]: def test_every_cost_map_provider_is_registered(path: Path): assert unregistered_providers(json.loads(path.read_text())) == [], ( f"{path.name} carries a litellm_provider whose models a `/*` grant does not list. A new provider " - "needs a `_models` set in litellm/__init__.py, filled in _populate_provider_model_sets and listed " + "needs a `_models` set in litellm/__init__.py, returned by _provider_model_set and listed " "in _build_models_by_provider. A new `-` label needs its rows added to a set that " "`models_by_provider[]` includes" ) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cdff03af3b4..dc31cc7e972 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2499,6 +2499,24 @@ def test_registered_chat_model_lists_under_its_provider_wildcard( litellm.add_known_models(model_cost_map={}) +def test_a_registered_azure_o_series_deployment_name_keeps_reasoning_effort_supported(): + model_name: Final = "my-o3-deployment-regression" + try: + litellm.register_model( + {model_name: {"litellm_provider": "azure", "mode": "chat"}}, persist_across_reloads=False + ) + + supported: Final = litellm.get_supported_openai_params(model=model_name, custom_llm_provider="azure") + + assert supported is not None + assert "reasoning_effort" in supported + finally: + litellm.model_cost.pop(model_name, None) + litellm.azure_models.discard(model_name) + litellm.model_list_set.discard(model_name) + litellm.add_known_models(model_cost_map={}) + + def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing(): from litellm.proxy.auth.model_checks import get_known_models_from_wildcard