From 6337ea876bd03def56c5d51fc77e2ec8ed93b557 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 5 Jan 2026 12:57:30 +0530
Subject: [PATCH] fix: Add custom llm provider to get_llm_provider when sent
 via UI

---
 litellm/main.py                               | 12 +++++-
 .../litellm_utils_tests/test_health_check.py  | 37 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index a0f3461b45..f4f27eb584 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -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")
 
diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py
index 98b1a35379..986f8cf491 100644
--- a/tests/litellm_utils_tests/test_health_check.py
+++ b/tests/litellm_utils_tests/test_health_check.py
@@ -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)