throw error when use_ai_gateway=True but AI Gateway V2 is not available for this workspace, remove logs

Author: jennsun

integrations/openai/src/databricks_openai/utils/clients.py

@@ -1,4 +1,3 @@
-import logging
 import os
 from typing import Any, Generator
 from urllib.parse import urlparse
@@ -11,8 +10,6 @@
 from openai.resources.responses import AsyncResponses, Responses
 from typing_extensions import override
 
-logger = logging.getLogger(__name__)
-
 # Prefix for routing requests to Databricks Apps
 _APPS_ENDPOINT_PREFIX = "apps/"
 # Domain pattern indicating a Databricks App URL
@@ -117,28 +114,19 @@ def _get_ai_gateway_base_url(
     Returns None if gateway is not available.
     """
     try:
-        url = f"{host}/api/ai-gateway/v2/endpoints"
-        logger.debug("Probing AI Gateway V2 at %s", url)
-        response = http_client.get(url)
+        response = http_client.get(f"{host}/api/ai-gateway/v2/endpoints")
         if response.status_code != 200:
-            logger.debug("AI Gateway V2 probe returned status %s", response.status_code)
             return None
         data = response.json()
         endpoints = data.get("endpoints", [])
         if not endpoints:
-            logger.debug("AI Gateway V2 returned no endpoints")
             return None
-        # Extract gateway base URL from the first endpoint's ai_gateway_url
         gateway_url = endpoints[0].get("ai_gateway_url")
         if not gateway_url:
-            logger.debug("AI Gateway V2 endpoint missing ai_gateway_url field")
             return None
         parsed = urlparse(gateway_url)
-        base_url = f"{parsed.scheme}://{parsed.netloc}/mlflow/v1"
-        logger.info("AI Gateway V2 detected, using base URL: %s", base_url)
-        return base_url
+        return f"{parsed.scheme}://{parsed.netloc}/mlflow/v1"
     except Exception:
-        logger.debug("AI Gateway V2 check failed, falling back to serving endpoints", exc_info=True)
         return None
 
 
@@ -159,6 +147,10 @@ def _resolve_base_url(
         gateway_url = _get_ai_gateway_base_url(http_client, workspace_client.config.host)
         if gateway_url:
             return gateway_url
+        raise ValueError(
+            "use_ai_gateway=True but AI Gateway V2 is not available for this workspace. "
+            "Set use_ai_gateway=False to use serving endpoints directly."
+        )
 
     # Fallback to using serving endpoints
     return f"{workspace_client.config.host}/serving-endpoints"

integrations/openai/tests/unit_tests/test_clients.py

@@ -776,18 +776,16 @@ def test_gateway_available_uses_gateway_url(self, client_cls_name, mock_workspac
             assert "12345.ai-gateway.cloud.databricks.com" in str(client.base_url)
 
     @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
-    def test_gateway_unavailable_falls_back_to_serving(
-        self, client_cls_name, mock_workspace_client
-    ):
+    def test_gateway_unavailable_raises_error(self, client_cls_name, mock_workspace_client):
         client_cls = (
             DatabricksOpenAI if client_cls_name == "DatabricksOpenAI" else AsyncDatabricksOpenAI
         )
         with patch(
             "databricks_openai.utils.clients._get_ai_gateway_base_url",
             return_value=None,
         ):
-            client = client_cls(workspace_client=mock_workspace_client, use_ai_gateway=True)
-            assert "/serving-endpoints/" in str(client.base_url)
+            with pytest.raises(ValueError, match="use_ai_gateway=True but AI Gateway V2"):
+                client_cls(workspace_client=mock_workspace_client, use_ai_gateway=True)
 
     @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
     def test_gateway_disabled_no_api_call(self, client_cls_name, mock_workspace_client):
@@ -815,15 +813,3 @@ def test_explicit_base_url_skips_gateway_check(self, client_cls_name, mock_works
             )
             mock_gateway.assert_not_called()
             assert "custom.example.com" in str(client.base_url)
-
-    @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
-    def test_gateway_detection_failure_falls_back(self, client_cls_name, mock_workspace_client):
-        client_cls = (
-            DatabricksOpenAI if client_cls_name == "DatabricksOpenAI" else AsyncDatabricksOpenAI
-        )
-        with patch(
-            "databricks_openai.utils.clients._get_ai_gateway_base_url",
-            return_value=None,
-        ):
-            client = client_cls(workspace_client=mock_workspace_client, use_ai_gateway=True)
-            assert "/serving-endpoints/" in str(client.base_url)