mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-11 16:26:07 +00:00
fix: Add custom llm provider to get_llm_provider when sent via UI
This commit is contained in:
+10
-2
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user