diff --git a/litellm/router.py b/litellm/router.py index 9185e437a3a..a1ae374ecb1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7754,6 +7754,8 @@ class Router: ) elif litellm_model_name_model_info is not None: model_info = litellm_model_name_model_info + elif custom_model_info is not None: + model_info = cast(ModelInfo, custom_model_info) return model_info diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index dc9b2c525c2..4238b724836 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1955,6 +1955,25 @@ def test_get_deployment_model_info_base_model_flow(): # Should return None when no model info is found assert result is None + # Test Case 6: Only custom model info exists (no model-name info) + mock_custom_only_info = { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "mode": "chat", + "litellm_provider": "openai", + } + + with patch.object(litellm, "model_cost", {"custom-only-id": mock_custom_only_info}): + with patch.object(litellm, "get_model_info", return_value=None): + result = router.get_deployment_model_info( + model_id="custom-only-id", model_name="missing-model-name" + ) + + assert result is not None + assert result["input_cost_per_token"] == 0.0 + assert result["output_cost_per_token"] == 0.0 + assert result["litellm_provider"] == "openai" + print("✓ All base model flow test cases passed!")