scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(databricks): support per-model client_id/client_secret in litellm_params

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Tam Nguyen 2026-03-20 10:55:24 +07:00
parent 1cf3947236
commit d82d544ff0
No known key found for this signature in database
GPG key ID: BC498F2C00E1C8B1
1 changed files with 85 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

85
tests/test_litellm/llms/databricks/test_databricks_common_utils.py
View file

@ -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",
)