Add support for AI Gateway native provider endpoints with databricks_openai (#390)

Author: smurching

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

@@ -109,15 +109,15 @@ def _fix_empty_assistant_content_in_messages(messages: Any) -> None:
                 message["content"] = " "
 
 
-def _get_ai_gateway_base_url(
+def _discover_ai_gateway_host(
     http_client: Client,
     host: str,
 ) -> str | None:
-    """Check if AI Gateway V2 is enabled and return its base URL.
+    """Discover the AI Gateway host URL (scheme + netloc only).
 
     Calls GET /api/ai-gateway/v2/endpoints. If successful and endpoints exist,
-    extracts the ai_gateway_url from the first endpoint response.
-    Returns None if gateway is not available.
+    extracts the ai_gateway_url from the first endpoint response and returns
+    just the scheme + netloc (no path). Returns None if gateway is not available.
     """
     try:
         response = http_client.get(f"{host}/api/ai-gateway/v2/endpoints")
@@ -131,24 +131,53 @@ def _get_ai_gateway_base_url(
         if not gateway_url:
             return None
         parsed = urlparse(gateway_url)
-        return f"{parsed.scheme}://{parsed.netloc}/mlflow/v1"
+        return f"{parsed.scheme}://{parsed.netloc}"
     except Exception:
         return None
 
 
+def _get_ai_gateway_base_url(
+    http_client: Client,
+    host: str,
+) -> str | None:
+    """Check if AI Gateway V2 is enabled and return its MLflow base URL.
+
+    Returns the AI Gateway base URL with /mlflow/v1 path appended, or None if
+    the gateway is not available.
+    """
+    gateway_host = _discover_ai_gateway_host(http_client, host)
+    return f"{gateway_host}/mlflow/v1" if gateway_host else None
+
+
 def _resolve_base_url(
     workspace_client: WorkspaceClient,
     base_url: str | None,
     use_ai_gateway: bool,
     http_client: Client,
+    use_ai_gateway_native_api: bool,
 ) -> str:
     """Resolve the target base URL for the OpenAI client."""
+    if use_ai_gateway_native_api and base_url is not None:
+        raise ValueError("Cannot specify both 'use_ai_gateway_native_api' and 'base_url'.")
+    if use_ai_gateway_native_api and use_ai_gateway:
+        raise ValueError("Cannot specify both 'use_ai_gateway_native_api' and 'use_ai_gateway'.")
+
     if base_url is not None:
         if _DATABRICKS_APPS_DOMAIN in base_url:
             _validate_oauth_for_apps(workspace_client)
         return base_url
 
-    # Prioritize using AI Gateway endpoints
+    # Native provider API via AI Gateway (e.g. OpenAI-compatible /openai path)
+    if use_ai_gateway_native_api:
+        gateway_host = _discover_ai_gateway_host(http_client, workspace_client.config.host)
+        if gateway_host:
+            return f"{gateway_host}/openai/v1"
+        raise ValueError(
+            "Please ensure AI Gateway V2 is enabled for the workspace "
+            "when use_ai_gateway_native_api is set to True."
+        )
+
+    # MLflow-format AI Gateway endpoints
     if use_ai_gateway:
         gateway_url = _get_ai_gateway_base_url(http_client, workspace_client.config.host)
         if gateway_url:
@@ -362,8 +391,12 @@ class DatabricksOpenAI(OpenAI):
         base_url: Optional base URL to override the default serving endpoints URL. When the URL
             points to a Databricks App (contains "databricksapps"), OAuth authentication is
             required.
+        use_ai_gateway_native_api: If True, auto-detect AI Gateway V2 and route requests through
+            its native OpenAI-compatible API (``<ai_gateway_url>/openai/v1``). This allows use of
+            provider-native features not available through the MLflow API. Cannot be combined
+            with ``base_url`` or ``use_ai_gateway``. Defaults to False.
         use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route
-            requests through it. Defaults to False.
+            requests through it using the MLflow API. Defaults to False.
 
     Example - Query a serving or AI gateway endpoint:
         >>> client = DatabricksOpenAI()
@@ -372,6 +405,13 @@ class DatabricksOpenAI(OpenAI):
         ...     messages=[{"role": "user", "content": "Hello!"}],
         ... )
 
+    Example - Query AI Gateway endpoints via the native OpenAI-compatible API:
+        >>> client = DatabricksOpenAI(use_ai_gateway_native_api=True)
+        >>> response = client.chat.completions.create(
+        ...     model="databricks-meta-llama-3-1-70b-instruct",
+        ...     messages=[{"role": "user", "content": "Hello!"}],
+        ... )
+
     Example - Query a Databricks App directly by URL (requires OAuth):
         >>> # WorkspaceClient must be configured with OAuth authentication
         >>> # See: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-u2m.html
@@ -397,6 +437,7 @@ def __init__(
         self,
         workspace_client: WorkspaceClient | None = None,
         base_url: str | None = None,
+        use_ai_gateway_native_api: bool = False,
         use_ai_gateway: bool = False,
     ):
         if workspace_client is None:
@@ -405,7 +446,9 @@ def __init__(
         self._workspace_client = workspace_client
 
         http_client = _get_authorized_http_client(workspace_client)
-        target_base_url = _resolve_base_url(workspace_client, base_url, use_ai_gateway, http_client)
+        target_base_url = _resolve_base_url(
+            workspace_client, base_url, use_ai_gateway, http_client, use_ai_gateway_native_api
+        )
 
         # Authentication is handled via http_client, not api_key
         super().__init__(
@@ -510,8 +553,12 @@ class AsyncDatabricksOpenAI(AsyncOpenAI):
         base_url: Optional base URL to override the default serving endpoints URL. When the URL
             points to a Databricks App (contains "databricksapps"), OAuth authentication is
             required.
+        use_ai_gateway_native_api: If True, auto-detect AI Gateway V2 and route requests through
+            its native OpenAI-compatible API (``<ai_gateway_url>/openai/v1``). This allows use of
+            provider-native features not available through the MLflow API. Cannot be combined
+            with ``base_url`` or ``use_ai_gateway``. Defaults to False.
         use_ai_gateway: If True, auto-detect AI Gateway V2 availability and route
-            requests through it. Defaults to False.
+            requests through it using the MLflow API. Defaults to False.
 
     Example - Query a serving or AI gateway endpoint:
         >>> client = AsyncDatabricksOpenAI()
@@ -520,6 +567,13 @@ class AsyncDatabricksOpenAI(AsyncOpenAI):
         ...     messages=[{"role": "user", "content": "Hello!"}],
         ... )
 
+    Example - Query AI Gateway endpoints via the native OpenAI-compatible API:
+        >>> client = AsyncDatabricksOpenAI(use_ai_gateway_native_api=True)
+        >>> response = await client.chat.completions.create(
+        ...     model="databricks-meta-llama-3-1-70b-instruct",
+        ...     messages=[{"role": "user", "content": "Hello!"}],
+        ... )
+
     Example - Query a Databricks App directly by URL (requires OAuth):
         >>> # WorkspaceClient must be configured with OAuth authentication
         >>> # See: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-u2m.html
@@ -545,6 +599,7 @@ def __init__(
         self,
         workspace_client: WorkspaceClient | None = None,
         base_url: str | None = None,
+        use_ai_gateway_native_api: bool = False,
         use_ai_gateway: bool = False,
     ):
         if workspace_client is None:
@@ -554,7 +609,7 @@ def __init__(
 
         sync_http_client = _get_authorized_http_client(workspace_client)
         target_base_url = _resolve_base_url(
-            workspace_client, base_url, use_ai_gateway, sync_http_client
+            workspace_client, base_url, use_ai_gateway, sync_http_client, use_ai_gateway_native_api
         )
 
         # Authentication is handled via http_client, not api_key

integrations/openai/tests/integration_tests/test_ai_gateway_native_api.py

@@ -0,0 +1,73 @@
+"""
+Integration tests for DatabricksOpenAI with use_ai_gateway_native_api=True.
+
+Prerequisites:
+- AI Gateway V2 must be enabled on the test workspace.
+- Set DATABRICKS_CONFIG_PROFILE (or use default credentials).
+
+Run with:
+    RUN_AI_GATEWAY_NATIVE_API_TESTS=1 python -m pytest \
+        tests/integration_tests/test_ai_gateway_native_api.py -v
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import pytest_asyncio
+from databricks.sdk import WorkspaceClient
+
+from databricks_openai import AsyncDatabricksOpenAI, DatabricksOpenAI
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("RUN_AI_GATEWAY_NATIVE_API_TESTS") != "1",
+    reason="AI Gateway native API tests disabled. Set RUN_AI_GATEWAY_NATIVE_API_TESTS=1 to enable.",
+)
+
+_TEST_MODEL = os.environ.get("AI_GATEWAY_NATIVE_API_MODEL", "databricks-gpt-5-4")
+_TEST_INPUT = [{"role": "user", "content": "Reply with exactly the word PONG and nothing else."}]
+
+
+@pytest.fixture(scope="module")
+def workspace_client():
+    return WorkspaceClient()
+
+
+@pytest.fixture(scope="module")
+def sync_client(workspace_client):
+    return DatabricksOpenAI(workspace_client=workspace_client, use_ai_gateway_native_api=True)
+
+
+@pytest_asyncio.fixture(scope="module")
+async def async_client(workspace_client):
+    return AsyncDatabricksOpenAI(workspace_client=workspace_client, use_ai_gateway_native_api=True)
+
+
+class TestAIGatewayNativeAPISync:
+    def test_base_url_uses_openai_path(self, sync_client):
+        assert "/openai/v1" in str(sync_client.base_url)
+        assert "ai-gateway" in str(sync_client.base_url)
+
+    def test_responses(self, sync_client):
+        response = sync_client.responses.create(
+            model=_TEST_MODEL,
+            input=_TEST_INPUT,
+            max_output_tokens=50,
+        )
+        assert response.output_text is not None
+
+
+@pytest.mark.asyncio
+class TestAIGatewayNativeAPIAsync:
+    async def test_base_url_uses_openai_path(self, async_client):
+        assert "/openai/v1" in str(async_client.base_url)
+        assert "ai-gateway" in str(async_client.base_url)
+
+    async def test_responses(self, async_client):
+        response = await async_client.responses.create(
+            model=_TEST_MODEL,
+            input=_TEST_INPUT,
+            max_output_tokens=50,
+        )
+        assert response.output_text is not None

integrations/openai/tests/unit_tests/test_clients.py

@@ -893,3 +893,83 @@ 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)
+
+
+class TestAIGatewayNativeAPI:
+    """Tests for use_ai_gateway_native_api parameter in DatabricksOpenAI and AsyncDatabricksOpenAI."""
+
+    @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
+    def test_native_api_uses_openai_base_path(self, client_cls_name, mock_workspace_client):
+        client_cls = (
+            DatabricksOpenAI if client_cls_name == "DatabricksOpenAI" else AsyncDatabricksOpenAI
+        )
+        with patch(
+            "databricks_openai.utils.clients._discover_ai_gateway_host",
+            return_value="https://12345.ai-gateway.cloud.databricks.com",
+        ):
+            client = client_cls(
+                workspace_client=mock_workspace_client,
+                use_ai_gateway_native_api=True,
+            )
+            assert "12345.ai-gateway.cloud.databricks.com" in str(client.base_url)
+            assert "/openai/v1" in str(client.base_url)
+
+    @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
+    def test_native_api_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._discover_ai_gateway_host",
+            return_value=None,
+        ):
+            with pytest.raises(ValueError, match="Please ensure AI Gateway V2 is enabled"):
+                client_cls(
+                    workspace_client=mock_workspace_client,
+                    use_ai_gateway_native_api=True,
+                )
+
+    @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
+    def test_native_api_and_base_url_raises(self, client_cls_name, mock_workspace_client):
+        client_cls = (
+            DatabricksOpenAI if client_cls_name == "DatabricksOpenAI" else AsyncDatabricksOpenAI
+        )
+        with pytest.raises(
+            ValueError, match="Cannot specify both 'use_ai_gateway_native_api' and 'base_url'"
+        ):
+            client_cls(
+                workspace_client=mock_workspace_client,
+                use_ai_gateway_native_api=True,
+                base_url="https://custom.example.com/v1",
+            )
+
+    @pytest.mark.parametrize("client_cls_name", ["DatabricksOpenAI", "AsyncDatabricksOpenAI"])
+    def test_native_api_and_use_ai_gateway_raises(self, client_cls_name, mock_workspace_client):
+        client_cls = (
+            DatabricksOpenAI if client_cls_name == "DatabricksOpenAI" else AsyncDatabricksOpenAI
+        )
+        with pytest.raises(
+            ValueError,
+            match="Cannot specify both 'use_ai_gateway_native_api' and 'use_ai_gateway'",
+        ):
+            client_cls(
+                workspace_client=mock_workspace_client,
+                use_ai_gateway_native_api=True,
+                use_ai_gateway=True,
+            )
+
+    def test_discover_ai_gateway_host_strips_path(self):
+        """_discover_ai_gateway_host returns only scheme+netloc, stripping any path."""
+        from unittest.mock import MagicMock
+
+        from databricks_openai.utils.clients import _discover_ai_gateway_host
+
+        mock_http = MagicMock()
+        mock_http.get.return_value.status_code = 200
+        mock_http.get.return_value.json.return_value = {
+            "endpoints": [
+                {"ai_gateway_url": "https://12345.ai-gateway.cloud.databricks.com/some/path"}
+            ]
+        }
+        result = _discover_ai_gateway_host(mock_http, "https://test.databricks.com")
+        assert result == "https://12345.ai-gateway.cloud.databricks.com"