mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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>
This commit is contained in:
shivam
parent
cad32fd9bc
commit
165ea273de
4 changed files with 216 additions and 107 deletions
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
69
tests/test_litellm/test_add_known_models_wildcard_refresh.py
Normal file
69
tests/test_litellm/test_add_known_models_wildcard_refresh.py
Normal file
|
|
@ -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/*")
|
||||
|
|
@ -129,6 +129,12 @@ const PriceDataReload: React.FC<PriceDataReloadProps> = ({
|
|||
// 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");
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue