fix: Add custom llm provider to get_llm_provider when sent via UI

This commit is contained in:
Sameer Kankute
2026-01-05 12:57:30 +05:30
parent fcabc059ca
commit 6337ea876b
2 changed files with 47 additions and 2 deletions
+10 -2
View File
@@ -110,7 +110,6 @@ from litellm.types.utils import (
RawRequestTypedDict,
StreamingChoices,
)
from litellm.utils import (
Choices,
CustomStreamWrapper,
@@ -6656,7 +6655,16 @@ async def ahealth_check(
if model in litellm.model_cost and mode is None:
mode = litellm.model_cost[model].get("mode")
model, custom_llm_provider, _, _ = get_llm_provider(model=model)
custom_llm_provider_from_params = model_params.get("custom_llm_provider", None)
api_base_from_params = model_params.get("api_base", None)
api_key_from_params = model_params.get("api_key", None)
model, custom_llm_provider, _, _ = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider_from_params,
api_base=api_base_from_params,
api_key=api_key_from_params,
)
if model in litellm.model_cost and mode is None:
mode = litellm.model_cost[model].get("mode")
@@ -637,3 +637,40 @@ async def test_image_generation_health_check_prompt(monkeypatch):
assert len(health_check_calls) == 1
assert health_check_calls[0]["prompt"] == override_prompt
@pytest.mark.asyncio
async def test_health_check_with_custom_llm_provider():
"""
Test that ahealth_check correctly uses custom_llm_provider from model_params.
This test verifies the fix for the issue where the UI's "Test connect" button
failed with "LLM Provider NOT provided" error for OpenAI-compatible self-hosted
providers, even when a provider was selected in the dropdown.
The fix ensures that when custom_llm_provider is passed in model_params,
it's properly forwarded to get_llm_provider() to identify the correct provider.
"""
from unittest.mock import MagicMock
# Mock the completion call to avoid making real API calls
mock_response = MagicMock()
mock_response._hidden_params = {"headers": {"x-ratelimit-remaining-tokens": "1000"}}
with patch("litellm.acompletion", return_value=mock_response):
# Test with a custom model name that wouldn't be recognized without custom_llm_provider
response = await litellm.ahealth_check(
model_params={
"model": "deepseek-r1-distill-qwen-1.5B-q4",
"custom_llm_provider": "openai",
"api_base": "https://example.com/v1",
"api_key": "fake-key",
},
mode="chat",
)
print(f"response: {response}")
# Should succeed without "LLM Provider NOT provided" error
assert "error" not in response
assert isinstance(response, dict)