scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge d82d544ff0 into 3e1479c052

Browse source
This commit is contained in:
Tam Nguyen 2026-04-29 23:13:54 +00:00 • committed by GitHub
commit 65ac837f78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 96 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/databricks/chat/transformation.py
View file

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

9
litellm/llms/databricks/common_utils.py
View file

@ -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:

2
litellm/llms/databricks/responses/transformation.py
View file

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

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