mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(register_model): handle openrouter models without '/' in name
The issue was in `register_model()` where `split_string[1]` assumed the model name always contained '/'. For custom model names like "glm" or UUIDs, the split would only produce one element, causing an IndexError. Changed `split_string[1]` to `split_string[-1]` which always returns the last element, working correctly for both cases: - "openrouter/gpt-4" → ["openrouter", "gpt-4"] → [-1] = "gpt-4" - "my-custom-alias" → ["my-custom-alias"] → [-1] = "my-custom-alias"
This commit is contained in:
parent
558f01e848
commit
bf3e6bb9f4
2 changed files with 59 additions and 1 deletions
|
|
@ -2780,7 +2780,7 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915
|
|||
elif value.get("litellm_provider") == "openrouter":
|
||||
split_string = key.split("/", 1)
|
||||
if key not in litellm.openrouter_models:
|
||||
litellm.openrouter_models.add(split_string[1])
|
||||
litellm.openrouter_models.add(split_string[-1])
|
||||
elif value.get("litellm_provider") == "vercel_ai_gateway":
|
||||
if key not in litellm.vercel_ai_gateway_models:
|
||||
litellm.vercel_ai_gateway_models.add(key)
|
||||
|
|
|
|||
|
|
@ -2300,6 +2300,64 @@ def test_register_model_with_scientific_notation():
|
|||
assert registered_model["mode"] == "chat"
|
||||
|
||||
|
||||
def test_register_model_openrouter_without_slash():
|
||||
"""
|
||||
Test that register_model handles openrouter models without '/' in the name.
|
||||
|
||||
Fixes https://github.com/BerriAI/litellm/issues/18936
|
||||
|
||||
Previously, the code did `split_string[1]` which would fail with IndexError
|
||||
when the model name didn't contain '/'. Now it uses `split_string[-1]` which
|
||||
always works.
|
||||
"""
|
||||
# Clear any existing entries
|
||||
litellm.openrouter_models.discard("my-custom-alias")
|
||||
litellm.openrouter_models.discard("gpt-4")
|
||||
litellm.openrouter_models.discard("openai/gpt-4")
|
||||
|
||||
# Test 1: Model name without '/' (this was the bug - would raise IndexError)
|
||||
litellm.register_model(
|
||||
{
|
||||
"my-custom-alias": {
|
||||
"max_tokens": 8192,
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"litellm_provider": "openrouter",
|
||||
"mode": "chat",
|
||||
},
|
||||
}
|
||||
)
|
||||
assert "my-custom-alias" in litellm.openrouter_models
|
||||
|
||||
# Test 2: Model name with single '/' (openrouter/model format)
|
||||
litellm.register_model(
|
||||
{
|
||||
"openrouter/gpt-4": {
|
||||
"max_tokens": 8192,
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"litellm_provider": "openrouter",
|
||||
"mode": "chat",
|
||||
},
|
||||
}
|
||||
)
|
||||
assert "gpt-4" in litellm.openrouter_models
|
||||
|
||||
# Test 3: Model name with double '/' (openrouter/provider/model format)
|
||||
litellm.register_model(
|
||||
{
|
||||
"openrouter/openai/gpt-4-turbo": {
|
||||
"max_tokens": 8192,
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
"litellm_provider": "openrouter",
|
||||
"mode": "chat",
|
||||
},
|
||||
}
|
||||
)
|
||||
assert "openai/gpt-4-turbo" in litellm.openrouter_models
|
||||
|
||||
|
||||
def test_reasoning_content_preserved_in_text_completion_wrapper():
|
||||
"""Ensure reasoning_content is copied from delta to text_choices."""
|
||||
chunk = ModelResponseStream(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue