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 <noreply@anthropic.com>
This commit is contained in:
Tam Nguyen 2026-03-20 10:52:54 +07:00
parent 3093ef844e
commit 1cf3947236
No known key found for this signature in database
GPG key ID: BC498F2C00E1C8B1
3 changed files with 11 additions and 2 deletions

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"

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:

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"