Merge pull request #21460 from TomeHirata/databricks/responses-api

Add native Responses API support for Databricks GPT models

litellm/__init__.py

@@ -1422,6 +1422,7 @@ if TYPE_CHECKING:
     from .llms.volcengine.responses.transformation import VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig
     from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig
     from .llms.perplexity.responses.transformation import PerplexityResponsesConfig as PerplexityResponsesConfig
+    from .llms.databricks.responses.transformation import DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig
     from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig
     from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config
     from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig

litellm/_lazy_imports_registry.py

@@ -227,6 +227,7 @@ LLM_CONFIG_NAMES = (
     "LiteLLMProxyResponsesAPIConfig",
     "VolcEngineResponsesAPIConfig",
     "PerplexityResponsesConfig",
+    "DatabricksResponsesAPIConfig",
     "GoogleAIStudioInteractionsConfig",
     "OpenAIOSeriesConfig",
     "AnthropicSkillsConfig",
@@ -906,6 +907,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
         ".llms.perplexity.responses.transformation",
         "PerplexityResponsesConfig",
     ),
+    "DatabricksResponsesAPIConfig": (
+        ".llms.databricks.responses.transformation",
+        "DatabricksResponsesAPIConfig",
+    ),
     "GoogleAIStudioInteractionsConfig": (
         ".llms.gemini.interactions.transformation",
         "GoogleAIStudioInteractionsConfig",

litellm/llms/databricks/responses/transformation.py

@@ -0,0 +1,100 @@
+"""
+Databricks Responses API configuration.
+
+Inherits from OpenAIResponsesAPIConfig since Databricks' Responses API
+is compatible with OpenAI's for GPT models.
+
+Reference: https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/api-reference
+"""
+
+import os
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from litellm.llms.databricks.common_utils import DatabricksBase
+from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+from litellm.types.llms.openai import ResponseInputParam
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+class DatabricksResponsesAPIConfig(DatabricksBase, OpenAIResponsesAPIConfig):
+    """
+    Configuration for Databricks Responses API.
+
+    Inherits from OpenAIResponsesAPIConfig since Databricks' Responses API
+    is largely compatible with OpenAI's for GPT models.
+
+    Note: The Responses API on Databricks is only compatible with OpenAI GPT models.
+    """
+
+    @property
+    def custom_llm_provider(self) -> LlmProviders:
+        return LlmProviders.DATABRICKS
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        litellm_params: Optional[GenericLiteLLMParams],
+    ) -> dict:
+        litellm_params = litellm_params or GenericLiteLLMParams()
+        api_key = litellm_params.api_key or os.getenv("DATABRICKS_API_KEY")
+        api_base = litellm_params.api_base or os.getenv("DATABRICKS_API_BASE")
+
+        # Reuse Databricks auth logic (OAuth M2M, PAT, SDK fallback).
+        # custom_endpoint=False allows SDK auth fallback; the appended
+        # /chat/completions suffix is harmless since we discard api_base
+        # here and build the URL separately in get_complete_url().
+        _, headers = self.databricks_validate_environment(
+            api_key=api_key,
+            api_base=api_base,
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=headers,
+        )
+
+        headers["Content-Type"] = "application/json"
+        return headers
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        api_base = api_base or os.getenv("DATABRICKS_API_BASE")
+        api_base = self._get_api_base(api_base)
+        api_base = api_base.rstrip("/")
+        return f"{api_base}/responses"
+
+    def transform_responses_api_request(
+        self,
+        model: str,
+        input: Union[str, ResponseInputParam],
+        response_api_optional_request_params: Dict,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,
+    ) -> Dict:
+        """
+        Transform request for Databricks Responses API.
+
+        Strips the 'databricks/' prefix from model name if present,
+        then delegates to OpenAI's transformation.
+        """
+        # Strip provider prefix if present (e.g., "databricks/databricks-gpt-5-nano" -> "databricks-gpt-5-nano")
+        if model.startswith("databricks/"):
+            model = model[len("databricks/") :]
+
+        return super().transform_responses_api_request(
+            model=model,
+            input=input,
+            response_api_optional_request_params=response_api_optional_request_params,
+            litellm_params=litellm_params,
+            headers=headers,
+        )

litellm/utils.py

@@ -8268,6 +8268,11 @@ class ProviderConfigManager:
             return litellm.ManusResponsesAPIConfig()
         elif litellm.LlmProviders.PERPLEXITY == provider:
             return litellm.PerplexityResponsesConfig()
+        elif litellm.LlmProviders.DATABRICKS == provider:
+            # Databricks Responses API is only compatible with OpenAI GPT models
+            if model and "gpt" in model.lower():
+                return litellm.DatabricksResponsesAPIConfig()
+            return None
         return None
 
     @staticmethod

tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py

@@ -0,0 +1,149 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+from unittest.mock import patch
+
+import litellm
+from litellm.llms.databricks.responses.transformation import (
+    DatabricksResponsesAPIConfig,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+
+class TestDatabricksResponsesAPIConfig:
+    """Tests for DatabricksResponsesAPIConfig"""
+
+    def test_custom_llm_provider(self):
+        config = DatabricksResponsesAPIConfig()
+        assert config.custom_llm_provider == LlmProviders.DATABRICKS
+
+    def test_get_complete_url(self):
+        config = DatabricksResponsesAPIConfig()
+        url = config.get_complete_url(
+            api_base="https://my-workspace.cloud.databricks.com/serving-endpoints",
+            litellm_params={},
+        )
+        assert (
+            url
+            == "https://my-workspace.cloud.databricks.com/serving-endpoints/responses"
+        )
+
+    def test_get_complete_url_strips_trailing_slash(self):
+        config = DatabricksResponsesAPIConfig()
+        url = config.get_complete_url(
+            api_base="https://my-workspace.cloud.databricks.com/serving-endpoints/",
+            litellm_params={},
+        )
+        assert (
+            url
+            == "https://my-workspace.cloud.databricks.com/serving-endpoints/responses"
+        )
+
+    def test_transform_request_strips_provider_prefix(self):
+        config = DatabricksResponsesAPIConfig()
+        request = config.transform_responses_api_request(
+            model="databricks/databricks-gpt-5-nano",
+            input="Hello!",
+            response_api_optional_request_params={},
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+        assert request["model"] == "databricks-gpt-5-nano"
+
+    def test_transform_request_no_prefix(self):
+        config = DatabricksResponsesAPIConfig()
+        request = config.transform_responses_api_request(
+            model="databricks-gpt-5-nano",
+            input="Hello!",
+            response_api_optional_request_params={},
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+        assert request["model"] == "databricks-gpt-5-nano"
+
+    def test_transform_request_preserves_text_param(self):
+        """Verify that the text/format param (response schema) is passed through to the request."""
+        config = DatabricksResponsesAPIConfig()
+        text_param = {
+            "format": {
+                "type": "json_schema",
+                "name": "Color",
+                "schema": {
+                    "type": "object",
+                    "properties": {"color": {"type": "string"}},
+                    "required": ["color"],
+                    "additionalProperties": False,
+                },
+            }
+        }
+        request = config.transform_responses_api_request(
+            model="databricks-gpt-5-nano",
+            input="Hello!",
+            response_api_optional_request_params={"text": text_param},
+            litellm_params=GenericLiteLLMParams(),
+            headers={},
+        )
+        assert request["text"] == text_param
+
+    def test_validate_environment_with_api_key(self):
+        config = DatabricksResponsesAPIConfig()
+        headers = config.validate_environment(
+            headers={},
+            model="databricks-gpt-5-nano",
+            litellm_params=GenericLiteLLMParams(
+                api_key="dapi_test_key",
+                api_base="https://my-workspace.cloud.databricks.com/serving-endpoints",
+            ),
+        )
+        assert headers["Authorization"] == "Bearer dapi_test_key"
+        assert headers["Content-Type"] == "application/json"
+
+
+class TestProviderConfigManagerDatabricks:
+    """Tests for Databricks registration in ProviderConfigManager"""
+
+    def test_gpt_model_returns_responses_config(self):
+        config = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.DATABRICKS,
+            model="databricks-gpt-5-nano",
+        )
+        assert config is not None
+        assert isinstance(config, DatabricksResponsesAPIConfig)
+
+    def test_gpt_model_case_insensitive(self):
+        config = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.DATABRICKS,
+            model="databricks-GPT-5-nano",
+        )
+        assert config is not None
+        assert isinstance(config, DatabricksResponsesAPIConfig)
+
+    def test_claude_model_returns_none(self):
+        """Claude models should fall back to completion transformation, not use Responses API."""
+        config = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.DATABRICKS,
+            model="databricks-claude-3-5-sonnet",
+        )
+        assert config is None
+
+    def test_llama_model_returns_none(self):
+        """Non-GPT models should fall back to completion transformation."""
+        config = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.DATABRICKS,
+            model="databricks-meta-llama-3-1-70b-instruct",
+        )
+        assert config is None
+
+    def test_no_model_returns_none(self):
+        config = ProviderConfigManager.get_provider_responses_api_config(
+            provider=LlmProviders.DATABRICKS,
+            model=None,
+        )
+        assert config is None