diff --git a/litellm/__init__.py b/litellm/__init__.py index 0f16fd5625c..ea7bb6a78d8 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 2af6ed8f09e..cf446d15156 100644 --- a/litellm/_lazy_imports_registry.py +++ b/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", diff --git a/litellm/llms/databricks/responses/__init__.py b/litellm/llms/databricks/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/databricks/responses/transformation.py b/litellm/llms/databricks/responses/transformation.py new file mode 100644 index 00000000000..0d9f433bfd2 --- /dev/null +++ b/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, + ) diff --git a/litellm/utils.py b/litellm/utils.py index 5d8d8a16db7..33f60f8d0a7 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/test_litellm/llms/databricks/responses/__init__.py b/tests/test_litellm/llms/databricks/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py b/tests/test_litellm/llms/databricks/responses/test_databricks_responses_transformation.py new file mode 100644 index 00000000000..b4a368be81f --- /dev/null +++ b/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