diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index c086d4ad755..09a1b390580 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -176,6 +176,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): custom_endpoint=False, headers=headers, custom_user_agent=custom_user_agent, + client_id=litellm_params.get("client_id"), + client_secret=litellm_params.get("client_secret"), ) # Ensure Content-Type header is set headers["Content-Type"] = "application/json" diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index d39d52d2d59..fb223546256 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -314,6 +314,8 @@ class DatabricksBase: custom_endpoint: Optional[bool], headers: Optional[dict], custom_user_agent: Optional[str] = None, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, ) -> Tuple[str, dict]: """ Validate and configure the Databricks environment. @@ -330,6 +332,8 @@ class DatabricksBase: custom_endpoint: Whether using a custom endpoint URL headers: Existing headers dict custom_user_agent: Optional custom user agent to prefix + client_id: OAuth client ID from litellm_params (falls back to env var) + client_secret: OAuth client secret from litellm_params (falls back to env var) Returns: Tuple of (api_base, headers) with authentication configured @@ -337,8 +341,9 @@ class DatabricksBase: from litellm._logging import verbose_logger # Check for OAuth M2M credentials (recommended for production) - client_id = os.getenv("DATABRICKS_CLIENT_ID") - client_secret = os.getenv("DATABRICKS_CLIENT_SECRET") + # Per-model litellm_params take priority over global env vars + client_id = client_id or os.getenv("DATABRICKS_CLIENT_ID") + client_secret = client_secret or os.getenv("DATABRICKS_CLIENT_SECRET") # Determine api_base first if api_base is None: diff --git a/litellm/llms/databricks/responses/transformation.py b/litellm/llms/databricks/responses/transformation.py index 090fef5ac82..afc5be37096 100644 --- a/litellm/llms/databricks/responses/transformation.py +++ b/litellm/llms/databricks/responses/transformation.py @@ -58,6 +58,8 @@ class DatabricksResponsesAPIConfig(DatabricksBase, OpenAIResponsesAPIConfig): endpoint_type="chat_completions", custom_endpoint=False, headers=headers, + client_id=getattr(litellm_params, "client_id", None), + client_secret=getattr(litellm_params, "client_secret", None), ) headers["Content-Type"] = "application/json" diff --git a/tests/test_litellm/llms/databricks/test_databricks_common_utils.py b/tests/test_litellm/llms/databricks/test_databricks_common_utils.py index 7f7ec8e9000..5b38cacb065 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_common_utils.py +++ b/tests/test_litellm/llms/databricks/test_databricks_common_utils.py @@ -30,3 +30,88 @@ def test_databricks_validate_environment(): except Exception: pass mock_get_credentials.assert_called_once() + + +def test_databricks_validate_environment_oauth_m2m_from_litellm_params(): + """Test that client_id/client_secret from litellm_params are used for OAuth M2M auth.""" + databricks_base = DatabricksBase() + + with patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="fake_token" + ) as mock_get_token: + api_base, headers = databricks_base.databricks_validate_environment( + api_key=None, + api_base="https://my-workspace.databricks.com/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + client_id="my-client-id", + client_secret="my-client-secret", + ) + + mock_get_token.assert_called_once_with( + "https://my-workspace.databricks.com/serving-endpoints", + "my-client-id", + "my-client-secret", + ) + assert headers["Authorization"] == "Bearer fake_token" + + +def test_databricks_validate_environment_oauth_m2m_params_override_env(): + """Test that litellm_params client_id/client_secret take priority over env vars.""" + databricks_base = DatabricksBase() + + with patch.dict( + os.environ, + { + "DATABRICKS_CLIENT_ID": "env-client-id", + "DATABRICKS_CLIENT_SECRET": "env-client-secret", + }, + ), patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="fake_token" + ) as mock_get_token: + databricks_base.databricks_validate_environment( + api_key=None, + api_base="https://my-workspace.databricks.com/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + client_id="param-client-id", + client_secret="param-client-secret", + ) + + # litellm_params values should take priority over env vars + mock_get_token.assert_called_once_with( + "https://my-workspace.databricks.com/serving-endpoints", + "param-client-id", + "param-client-secret", + ) + + +def test_databricks_validate_environment_oauth_m2m_falls_back_to_env(): + """Test that env vars are used when litellm_params don't have client_id/client_secret.""" + databricks_base = DatabricksBase() + + with patch.dict( + os.environ, + { + "DATABRICKS_CLIENT_ID": "env-client-id", + "DATABRICKS_CLIENT_SECRET": "env-client-secret", + }, + ), patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="fake_token" + ) as mock_get_token: + databricks_base.databricks_validate_environment( + api_key=None, + api_base="https://my-workspace.databricks.com/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + # Should fall back to env vars + mock_get_token.assert_called_once_with( + "https://my-workspace.databricks.com/serving-endpoints", + "env-client-id", + "env-client-secret", + )