From d82d544ff020aae93baa1773a42b364d6f278742 Mon Sep 17 00:00:00 2001 From: Tam Nguyen Date: Fri, 20 Mar 2026 10:55:24 +0700 Subject: [PATCH] feat(databricks): support per-model client_id/client_secret in litellm_params Allow Databricks OAuth M2M credentials (client_id, client_secret) to be specified per deployment in litellm_params, enabling load balancing across multiple Databricks workspace instances with different service principals. Previously these were only read from global env vars (DATABRICKS_CLIENT_ID, DATABRICKS_CLIENT_SECRET), making multi-instance configs impossible. Co-Authored-By: Claude Opus 4.6 --- .../test_databricks_common_utils.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) 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", + )