From 165ea273def41a1a49bd9aaa43de6a15d809abff Mon Sep 17 00:00:00 2001 From: shivam Date: Wed, 29 Jul 2026 21:52:13 +0000 Subject: [PATCH] fix(proxy): refresh models_by_provider on cost map reload so wildcard expansion picks up new models add_known_models mutates the per-provider sets, but models_by_provider was a frozen snapshot built once at import via | set unions, so newly added models never reached get_valid_models / get_known_models_from_wildcard. Rebuild models_by_provider from the live sets at the end of add_known_models. Also add a missing add_known_models branch for the bare vertex_ai provider so those models bucket into vertex_ai_models and expand under vertex_ai/*, and make the reload endpoint report a no-op (instead of success) when LITELLM_LOCAL_MODEL_COST_MAP forces the bundled backup map. Resolves LIT-4947 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 224 +++++++++--------- litellm/proxy/proxy_server.py | 24 +- .../test_add_known_models_wildcard_refresh.py | 69 ++++++ .../src/components/price_data_reload.tsx | 6 + 4 files changed, 216 insertions(+), 107 deletions(-) create mode 100644 tests/test_litellm/test_add_known_models_wildcard_refresh.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f8c742c5a2..f82a2e5b6b8 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -583,6 +583,7 @@ vertex_openai_models: Set = set() vertex_minimax_models: Set = set() vertex_moonshot_models: Set = set() vertex_zai_models: Set = set() +vertex_ai_models: Set = set() ai21_models: Set = set() ai21_chat_models: Set = set() nlp_cloud_models: Set = set() @@ -772,6 +773,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): elif value.get("litellm_provider") == "vertex_ai-zai_models": key = key.replace("vertex_ai/", "") vertex_zai_models.add(key) + elif value.get("litellm_provider") == "vertex_ai": + vertex_ai_models.add(key) elif value.get("litellm_provider") == "ai21": if value.get("mode") == "chat": ai21_chat_models.add(key) @@ -944,6 +947,10 @@ def add_known_models(model_cost_map: Optional[Dict] = None): elif value.get("litellm_provider") == "bedrock_mantle": bedrock_mantle_models.add(key) + global models_by_provider + if "models_by_provider" in globals(): + models_by_provider = _build_models_by_provider() + add_known_models() # known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary @@ -1067,112 +1074,117 @@ model_list_set = set(model_list) # provider_list is lazy-loaded via __getattr__ to avoid importing LlmProviders at import time -models_by_provider: dict = { - "openai": open_ai_chat_completion_models | open_ai_text_completion_models, - "text-completion-openai": open_ai_text_completion_models, - "cohere": cohere_models | cohere_chat_models, - "cohere_chat": cohere_chat_models, - "anthropic": anthropic_models, - "replicate": replicate_models, - "huggingface": huggingface_models, - "together_ai": together_ai_models, - "baseten": baseten_models, - "openrouter": openrouter_models, - "vercel_ai_gateway": vercel_ai_gateway_models, - "datarobot": datarobot_models, - "vertex_ai": vertex_chat_models - | vertex_text_models - | vertex_anthropic_models - | vertex_vision_models - | vertex_language_models - | vertex_deepseek_models - | vertex_minimax_models - | vertex_moonshot_models - | vertex_zai_models, - "ai21": ai21_models, - "bedrock": bedrock_models | bedrock_converse_models, - "petals": petals_models, - "ollama": ollama_models, - "ollama_chat": ollama_models, - "deepinfra": deepinfra_models, - "perplexity": perplexity_models, - "maritalk": maritalk_models, - "watsonx": watsonx_models, - "gemini": gemini_models, - "fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models, - "aleph_alpha": aleph_alpha_models, - "text-completion-codestral": text_completion_codestral_models, - "text-completion-inception": text_completion_inception_models, - "xai": xai_models, - "zai": zai_models, - "fal_ai": fal_ai_models, - "deepseek": deepseek_models, - "tencent": tencent_models, - "runwayml": runwayml_models, - "mistral": mistral_chat_models, - "azure_ai": azure_ai_models, - "voyage": voyage_models, - "infinity": infinity_models, - "databricks": databricks_models, - "cloudflare": cloudflare_models, - "codestral": codestral_models, - "nlp_cloud": nlp_cloud_models, - "friendliai": friendliai_models, - "palm": palm_models, - "groq": groq_models, - "azure": azure_models | azure_text_models, - "azure_anthropic": azure_anthropic_models, - "azure_text": azure_text_models, - "anyscale": anyscale_models, - "cerebras": cerebras_models, - "galadriel": galadriel_models, - "nvidia_nim": nvidia_nim_models, - "nvidia_riva": nvidia_riva_models, - "soniox": soniox_models, - "sambanova": sambanova_models | sambanova_embedding_models, - "novita": novita_models, - "nebius": nebius_models | nebius_embedding_models, - "aiml": aiml_models, - "assemblyai": assemblyai_models, - "jina_ai": jina_ai_models, - "snowflake": snowflake_models, - "gradient_ai": gradient_ai_models, - "meta_llama": llama_models, - "nscale": nscale_models, - "featherless_ai": featherless_ai_models, - "deepgram": deepgram_models, - "elevenlabs": elevenlabs_models, - "heroku": heroku_models, - "dashscope": dashscope_models, - "modelscope": modelscope_models, - "moonshot": moonshot_models, - "publicai": publicai_models, - "darkbloom": darkbloom_models, - "v0": v0_models, - "morph": morph_models, - "lambda_ai": lambda_ai_models, - "inception": inception_models, - "hyperbolic": hyperbolic_models, - "black_forest_labs": black_forest_labs_models, - "recraft": recraft_models, - "cometapi": cometapi_models, - "oci": oci_models, - "volcengine": volcengine_models, - "wandb": wandb_models, - "ovhcloud": ovhcloud_models | ovhcloud_embedding_models, - "lemonade": lemonade_models, - "clarifai": clarifai_models, - "amazon_nova": amazon_nova_models, - "stability": stability_models, - "github_copilot": github_copilot_models, - "chatgpt": chatgpt_models, - "minimax": minimax_models, - "aws_polly": aws_polly_models, - "gigachat": gigachat_models, - "llamagate": llamagate_models, - "reducto": reducto_models, - "bedrock_mantle": bedrock_mantle_models, -} +def _build_models_by_provider() -> dict: + return { + "openai": open_ai_chat_completion_models | open_ai_text_completion_models, + "text-completion-openai": open_ai_text_completion_models, + "cohere": cohere_models | cohere_chat_models, + "cohere_chat": cohere_chat_models, + "anthropic": anthropic_models, + "replicate": replicate_models, + "huggingface": huggingface_models, + "together_ai": together_ai_models, + "baseten": baseten_models, + "openrouter": openrouter_models, + "vercel_ai_gateway": vercel_ai_gateway_models, + "datarobot": datarobot_models, + "vertex_ai": vertex_chat_models + | vertex_text_models + | vertex_anthropic_models + | vertex_vision_models + | vertex_language_models + | vertex_deepseek_models + | vertex_minimax_models + | vertex_moonshot_models + | vertex_zai_models + | vertex_ai_models, + "ai21": ai21_models, + "bedrock": bedrock_models | bedrock_converse_models, + "petals": petals_models, + "ollama": ollama_models, + "ollama_chat": ollama_models, + "deepinfra": deepinfra_models, + "perplexity": perplexity_models, + "maritalk": maritalk_models, + "watsonx": watsonx_models, + "gemini": gemini_models, + "fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models, + "aleph_alpha": aleph_alpha_models, + "text-completion-codestral": text_completion_codestral_models, + "text-completion-inception": text_completion_inception_models, + "xai": xai_models, + "zai": zai_models, + "fal_ai": fal_ai_models, + "deepseek": deepseek_models, + "tencent": tencent_models, + "runwayml": runwayml_models, + "mistral": mistral_chat_models, + "azure_ai": azure_ai_models, + "voyage": voyage_models, + "infinity": infinity_models, + "databricks": databricks_models, + "cloudflare": cloudflare_models, + "codestral": codestral_models, + "nlp_cloud": nlp_cloud_models, + "friendliai": friendliai_models, + "palm": palm_models, + "groq": groq_models, + "azure": azure_models | azure_text_models, + "azure_anthropic": azure_anthropic_models, + "azure_text": azure_text_models, + "anyscale": anyscale_models, + "cerebras": cerebras_models, + "galadriel": galadriel_models, + "nvidia_nim": nvidia_nim_models, + "nvidia_riva": nvidia_riva_models, + "soniox": soniox_models, + "sambanova": sambanova_models | sambanova_embedding_models, + "novita": novita_models, + "nebius": nebius_models | nebius_embedding_models, + "aiml": aiml_models, + "assemblyai": assemblyai_models, + "jina_ai": jina_ai_models, + "snowflake": snowflake_models, + "gradient_ai": gradient_ai_models, + "meta_llama": llama_models, + "nscale": nscale_models, + "featherless_ai": featherless_ai_models, + "deepgram": deepgram_models, + "elevenlabs": elevenlabs_models, + "heroku": heroku_models, + "dashscope": dashscope_models, + "modelscope": modelscope_models, + "moonshot": moonshot_models, + "publicai": publicai_models, + "darkbloom": darkbloom_models, + "v0": v0_models, + "morph": morph_models, + "lambda_ai": lambda_ai_models, + "inception": inception_models, + "hyperbolic": hyperbolic_models, + "black_forest_labs": black_forest_labs_models, + "recraft": recraft_models, + "cometapi": cometapi_models, + "oci": oci_models, + "volcengine": volcengine_models, + "wandb": wandb_models, + "ovhcloud": ovhcloud_models | ovhcloud_embedding_models, + "lemonade": lemonade_models, + "clarifai": clarifai_models, + "amazon_nova": amazon_nova_models, + "stability": stability_models, + "github_copilot": github_copilot_models, + "chatgpt": chatgpt_models, + "minimax": minimax_models, + "aws_polly": aws_polly_models, + "gigachat": gigachat_models, + "llamagate": llamagate_models, + "reducto": reducto_models, + "bedrock_mantle": bedrock_mantle_models, + } + + +models_by_provider = _build_models_by_provider() # mapping for those models which have larger equivalents longer_context_model_fallback_dict: dict = { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4486cd7de59..e884e4649f5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15661,10 +15661,14 @@ async def reload_model_cost_map( raise HTTPException(status_code=500, detail="Database connection not available") # Immediately reload the model cost map in the current pod - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + get_model_cost_map_source_info, + ) model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) + source_info = get_model_cost_map_source_info() litellm.model_cost = new_model_cost_map # Invalidate case-insensitive lookup map since model_cost was replaced _invalidate_model_cost_lowercase_map() @@ -15698,12 +15702,30 @@ async def reload_model_cost_map( await invalidate_config_param("model_cost_map_reload_config") models_count = len(new_model_cost_map) if new_model_cost_map else 0 + + if source_info["is_env_forced"]: + no_op_message = ( + "Reload was a no-op: LITELLM_LOCAL_MODEL_COST_MAP=true forces the bundled backup " + "cost map, so no fresh pricing data was fetched. Unset that env var to reload live data." + ) + verbose_proxy_logger.warning(no_op_message) + return { + "message": no_op_message, + "status": "no_op", + "models_count": models_count, + "source": source_info["source"], + "is_env_forced": True, + "timestamp": current_time.isoformat(), + } + verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}") return { "message": f"Price data reloaded successfully! {models_count} models updated.", "status": "success", "models_count": models_count, + "source": source_info["source"], + "is_env_forced": False, "timestamp": current_time.isoformat(), } except Exception as e: diff --git a/tests/test_litellm/test_add_known_models_wildcard_refresh.py b/tests/test_litellm/test_add_known_models_wildcard_refresh.py new file mode 100644 index 00000000000..c9c0e3d675e --- /dev/null +++ b/tests/test_litellm/test_add_known_models_wildcard_refresh.py @@ -0,0 +1,69 @@ +"""Regression tests for LIT-4947. + +`add_known_models` must refresh `litellm.models_by_provider` so that a model newly +added to the cost map (e.g. after "Reload Price Data") shows up everywhere wildcard +expansion is used, without requiring a proxy restart. +""" + +import copy + +import pytest + +import litellm +from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + +@pytest.fixture +def restore_model_registry(): + """Snapshot and restore the global model registry mutated by add_known_models.""" + original_model_cost = litellm.model_cost + original_models_by_provider = litellm.models_by_provider + original_sets = { + name: set(value) + for name in dir(litellm) + if name.endswith("_models") and isinstance((value := getattr(litellm, name)), set) + } + + yield + + litellm.model_cost = original_model_cost + litellm.models_by_provider = original_models_by_provider + for name, snapshot in original_sets.items(): + live_set = getattr(litellm, name) + live_set.clear() + live_set.update(snapshot) + + +def _reload_with_new_model(model: str, litellm_provider: str) -> None: + new_map = copy.deepcopy(litellm.model_cost) + new_map[model] = {"litellm_provider": litellm_provider, "mode": "chat"} + litellm.model_cost = new_map + litellm.add_known_models(model_cost_map=new_map) + + +class TestAddKnownModelsRefreshesWildcardExpansion: + def test_new_vertex_language_model_reaches_models_by_provider_and_wildcard(self, restore_model_registry): + model = "vertex_ai/gemini-9.9-flash-lite" + + _reload_with_new_model(model, "vertex_ai-language-models") + + assert model in litellm.vertex_language_models + assert model in litellm.models_by_provider["vertex_ai"] + assert model in get_known_models_from_wildcard("vertex_ai/*") + + def test_bare_vertex_ai_provider_model_is_bucketed(self, restore_model_registry): + model = "vertex_ai/some-new-native-thing" + + _reload_with_new_model(model, "vertex_ai") + + assert model in litellm.vertex_ai_models + assert model in litellm.models_by_provider["vertex_ai"] + assert model in get_known_models_from_wildcard("vertex_ai/*") + + def test_new_openai_model_reaches_wildcard(self, restore_model_registry): + model = "gpt-99-turbo" + + _reload_with_new_model(model, "openai") + + assert model in litellm.models_by_provider["openai"] + assert f"openai/{model}" in get_known_models_from_wildcard("openai/*") diff --git a/ui/litellm-dashboard/src/components/price_data_reload.tsx b/ui/litellm-dashboard/src/components/price_data_reload.tsx index e0ec0fb794d..72fb64d5396 100644 --- a/ui/litellm-dashboard/src/components/price_data_reload.tsx +++ b/ui/litellm-dashboard/src/components/price_data_reload.tsx @@ -129,6 +129,12 @@ const PriceDataReload: React.FC = ({ // Refresh status and source info after successful reload await fetchReloadStatus(); await fetchSourceInfo(); + } else if (response.status === "no_op") { + NotificationsManager.fromBackend( + response.message || "Reload was a no-op because LITELLM_LOCAL_MODEL_COST_MAP forces the bundled cost map.", + ); + await fetchReloadStatus(); + await fetchSourceInfo(); } else { NotificationsManager.fromBackend("Failed to reload price data"); }