Merge pull request #18349 from prasadkona/feat/databricks-partner-integration

feat(databricks): Add enhanced authentication, security features, and custom user-agent support

docs/my-website/docs/providers/databricks.md

@@ -11,6 +11,99 @@ LiteLLM supports all models on Databricks
 
 :::
 
+## Authentication
+
+LiteLLM supports multiple authentication methods for Databricks, listed in order of preference:
+
+### OAuth M2M (Recommended for Production)
+
+OAuth Machine-to-Machine authentication using Service Principal credentials is the **recommended method for production** deployments per Databricks Partner requirements.
+
+```python
+import os
+from litellm import completion
+
+# Set OAuth credentials (Service Principal)
+os.environ["DATABRICKS_CLIENT_ID"] = "your-service-principal-application-id"
+os.environ["DATABRICKS_CLIENT_SECRET"] = "your-service-principal-secret"
+os.environ["DATABRICKS_API_BASE"] = "https://adb-xxx.azuredatabricks.net/serving-endpoints"
+
+response = completion(
+    model="databricks/databricks-dbrx-instruct",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+```
+
+### Personal Access Token (PAT)
+
+PAT authentication is supported for development and testing scenarios.
+
+```python
+import os
+from litellm import completion
+
+os.environ["DATABRICKS_API_KEY"] = "dapi..."  # Your Personal Access Token
+os.environ["DATABRICKS_API_BASE"] = "https://adb-xxx.azuredatabricks.net/serving-endpoints"
+
+response = completion(
+    model="databricks/databricks-dbrx-instruct",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+```
+
+### Databricks SDK Authentication (Automatic)
+
+If no credentials are provided, LiteLLM will use the Databricks SDK for automatic authentication. This supports OAuth, Azure AD, and other unified auth methods configured in your environment.
+
+```python
+from litellm import completion
+
+# No environment variables needed - uses Databricks SDK unified auth
+# Requires: pip install databricks-sdk
+response = completion(
+    model="databricks/databricks-dbrx-instruct",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+```
+
+## Custom User-Agent for Partner Attribution
+
+If you're building a product on top of LiteLLM that integrates with Databricks, you can pass your own partner identifier for proper attribution in Databricks telemetry.
+
+The partner name will be prefixed to the LiteLLM user agent:
+
+```python
+# Via parameter
+response = completion(
+    model="databricks/databricks-dbrx-instruct",
+    messages=[{"role": "user", "content": "Hello!"}],
+    user_agent="mycompany/1.0.0",
+)
+# Resulting User-Agent: mycompany_litellm/1.79.1
+
+# Via environment variable
+os.environ["DATABRICKS_USER_AGENT"] = "mycompany/1.0.0"
+# Resulting User-Agent: mycompany_litellm/1.79.1
+```
+
+| Input | Resulting User-Agent |
+|-------|---------------------|
+| (none) | `litellm/1.79.1` |
+| `mycompany/1.0.0` | `mycompany_litellm/1.79.1` |
+| `partner_product/2.5.0` | `partner_product_litellm/1.79.1` |
+| `acme` | `acme_litellm/1.79.1` |
+
+**Note:** The version from your custom user agent is ignored; LiteLLM's version is always used.
+
+## Security
+
+LiteLLM automatically redacts sensitive information (tokens, secrets, API keys) from all debug logs to prevent credential leakage. This includes:
+
+- Authorization headers
+- API keys and tokens
+- Client secrets
+- Personal access tokens (PATs)
+
 ## Usage
 
 <Tabs>
@@ -51,6 +144,7 @@ response = completion(
         model: databricks/databricks-dbrx-instruct
         api_key: os.environ/DATABRICKS_API_KEY
         api_base: os.environ/DATABRICKS_API_BASE
+        user_agent: "mycompany/1.0.0"  # Optional: for partner attribution
   ```
 
 

litellm/llms/databricks/chat/transformation.py

@@ -2,6 +2,7 @@
 Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completions`
 """
 
+import os
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -26,7 +27,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
     _should_convert_tool_call_to_json_mode,
 )
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
-    strip_name_from_message
+    strip_name_from_message,
 )
 from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
 from litellm.types.llms.anthropic import AllAnthropicToolsValues
@@ -124,12 +125,24 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
     ) -> dict:
+        # Check for custom user agent in optional_params or environment
+        # This allows partners building on LiteLLM to set their own telemetry
+        # Use pop() to remove these keys so they don't get sent to the API
+        custom_user_agent = (
+            optional_params.pop("user_agent", None)
+            or optional_params.pop("databricks_user_agent", None)
+            or litellm_params.get("user_agent")
+            or os.getenv("LITELLM_USER_AGENT")
+            or os.getenv("DATABRICKS_USER_AGENT")
+        )
+
         api_base, headers = self.databricks_validate_environment(
             api_base=api_base,
             api_key=api_key,
             endpoint_type="chat_completions",
             custom_endpoint=False,
             headers=headers,
+            custom_user_agent=custom_user_agent,
         )
         # Ensure Content-Type header is set
         headers["Content-Type"] = "application/json"
@@ -173,9 +186,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         # Build DatabricksFunction explicitly to avoid parameter conflicts
         function_params: DatabricksFunction = {
             "name": tool["name"],
-            "parameters": cast(dict, tool.get("input_schema") or {})
+            "parameters": cast(dict, tool.get("input_schema") or {}),
         }
-        
+
         # Only add description if it exists
         description = tool.get("description")
         if description is not None:
@@ -229,7 +242,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         Databricks supports Anthropic-style cache control for Claude models.
         Databricks ignores the cache_control flag with other models.
         """
-        # TODO: Think about how to best design the request transformation so that 
+        # TODO: Think about how to best design the request transformation so that
         # every request doesn't have to be transformed for to OpenAI and Anthropic request formats.
         return messages, tools
 
@@ -347,15 +360,17 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
                 messages=new_messages, model=model, is_async=cast(Literal[False], False)
             )
 
-    def _move_cache_control_into_string_content_block(self, message: AllMessageValues) -> AllMessageValues:
+    def _move_cache_control_into_string_content_block(
+        self, message: AllMessageValues
+    ) -> AllMessageValues:
         """
         Moves message-level cache_control into a content block when content is a string.
-        
+
         Transforms:
             {"role": "user", "content": "text", "cache_control": {...}}
         Into:
             {"role": "user", "content": [{"type": "text", "text": "text", "cache_control": {...}}]}
-        
+
         This is required for Anthropic's prompt caching API when cache_control is specified
         at the message level but content is a simple string (not already an array of content blocks).
         """
@@ -371,7 +386,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
             }
         ]
         return cast(AllMessageValues, transformed_message)
-        
 
     @staticmethod
     def extract_content_str(
@@ -509,9 +523,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
                     reasoning_content=reasoning_content,
                     thinking_blocks=thinking_blocks,
                     tool_calls=choice["message"].get("tool_calls"),
-                    provider_specific_fields={"citations": citations}
-                    if citations is not None
-                    else None,
+                    provider_specific_fields=(
+                        {"citations": citations} if citations is not None else None
+                    ),
                 )
 
             if finish_reason is None:
@@ -543,12 +557,15 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         api_key: Optional[str] = None,
         json_mode: Optional[bool] = None,
     ) -> ModelResponse:
-        ## LOGGING
+        # Redact sensitive data before logging to prevent credential leakage
+        redacted_request_data = self.redact_sensitive_data(request_data)
+
+        ## LOGGING - Never log actual API keys
         logging_obj.post_call(
             input=messages,
-            api_key=api_key,
+            api_key="[REDACTED]",
             original_response=raw_response.text,
-            additional_args={"complete_input_dict": request_data},
+            additional_args={"complete_input_dict": redacted_request_data},
         )
 
         ## RESPONSE OBJECT

litellm/llms/databricks/common_utils.py

@@ -1,4 +1,18 @@
-from typing import Literal, Optional, Tuple
+"""
+Databricks integration utilities for LiteLLM.
+
+This module provides authentication, telemetry, and security utilities
+for the Databricks LLM provider integration.
+
+Authentication priority:
+1. OAuth M2M (DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET) - Recommended for production
+2. PAT (DATABRICKS_API_KEY) - Supported for development
+3. Databricks SDK automatic auth - Fallback (uses unified auth)
+"""
+
+import os
+import re
+from typing import Any, Dict, Literal, Optional, Tuple
 
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 
@@ -8,17 +22,175 @@ class DatabricksException(BaseLLMException):
 
 
 class DatabricksBase:
+    """
+    Base class for Databricks integration with authentication,
+    telemetry, and security utilities.
+    """
+
+    # Patterns to redact in logs
+    SENSITIVE_PATTERNS = [
+        (re.compile(r"(Bearer\s+)[A-Za-z0-9\-_\.]+", re.IGNORECASE), r"\1[REDACTED]"),
+        (re.compile(r"(Authorization:\s*)[^\s,}]+", re.IGNORECASE), r"\1[REDACTED]"),
+        (
+            re.compile(r'(api[_-]?key["\s:=]+)[^\s,}"\']+', re.IGNORECASE),
+            r"\1[REDACTED]",
+        ),
+        (
+            re.compile(r'(client[_-]?secret["\s:=]+)[^\s,}"\']+', re.IGNORECASE),
+            r"\1[REDACTED]",
+        ),
+        (re.compile(r"(dapi[a-zA-Z0-9]{32,})", re.IGNORECASE), r"[REDACTED_PAT]"),
+        (
+            re.compile(r'(access[_-]?token["\s:=]+)[^\s,}"\']+', re.IGNORECASE),
+            r"\1[REDACTED]",
+        ),
+    ]
+
+    @classmethod
+    def redact_sensitive_data(cls, data: Any) -> Any:
+        """
+        Redact sensitive information (tokens, secrets) from data before logging.
+
+        Handles strings, dicts, and lists recursively. Keys containing sensitive
+        terms (authorization, api_key, token, secret, password, credential) are
+        fully redacted.
+
+        Args:
+            data: String, dict, or other data structure to redact
+
+        Returns:
+            Redacted version of the data safe for logging
+        """
+        if data is None:
+            return None
+
+        if isinstance(data, str):
+            result = data
+            for pattern, replacement in cls.SENSITIVE_PATTERNS:
+                result = pattern.sub(replacement, result)
+            return result
+
+        if isinstance(data, dict):
+            redacted = {}
+            for key, value in data.items():
+                lower_key = key.lower()
+                if any(
+                    sensitive in lower_key
+                    for sensitive in [
+                        "authorization",
+                        "api_key",
+                        "apikey",
+                        "token",
+                        "secret",
+                        "password",
+                        "credential",
+                    ]
+                ):
+                    redacted[key] = "[REDACTED]"
+                else:
+                    redacted[key] = cls.redact_sensitive_data(value)
+            return redacted
+
+        if isinstance(data, list):
+            return [cls.redact_sensitive_data(item) for item in data]
+
+        return data
+
+    @classmethod
+    def redact_headers_for_logging(cls, headers: Dict[str, str]) -> Dict[str, str]:
+        """
+        Create a copy of headers with sensitive values redacted for safe logging.
+
+        Shows first 8 characters of sensitive values for debugging purposes,
+        with the rest redacted.
+
+        Args:
+            headers: HTTP headers dictionary
+
+        Returns:
+            New dictionary with sensitive headers redacted
+        """
+        if not headers:
+            return {}
+
+        redacted = {}
+        sensitive_headers = {
+            "authorization",
+            "x-api-key",
+            "api-key",
+            "x-databricks-token",
+        }
+
+        for key, value in headers.items():
+            if key.lower() in sensitive_headers:
+                if len(value) > 10:
+                    redacted[key] = f"{value[:8]}...[REDACTED]"
+                else:
+                    redacted[key] = "[REDACTED]"
+            else:
+                redacted[key] = value
+
+        return redacted
+
+    @staticmethod
+    def _build_user_agent(custom_user_agent: Optional[str] = None) -> str:
+        """
+        Build the User-Agent string for Databricks API calls.
+
+        If a custom user agent is provided, the partner name (part before /)
+        is extracted and prefixed to the litellm user agent with an underscore.
+        The custom version is ignored; LiteLLM's version is always used.
+
+        Args:
+            custom_user_agent: Optional custom user agent string (e.g., "mycompany/1.0.0")
+
+        Returns:
+            User-Agent string in format:
+            - Default: "litellm/{version}"
+            - With custom: "{partner}_litellm/{version}"
+
+        Examples:
+            - None -> "litellm/1.79.1"
+            - "mycompany/1.0.0" -> "mycompany_litellm/1.79.1"
+            - "partner_product/2.0.0" -> "partner_product_litellm/1.79.1"
+            - "acme" -> "acme_litellm/1.79.1"
+        """
+        try:
+            from litellm._version import version
+        except Exception:
+            version = "0.0.0"
+
+        if custom_user_agent:
+            custom_user_agent = custom_user_agent.strip()
+
+            # Extract partner name (part before / if present)
+            if "/" in custom_user_agent:
+                partner_name = custom_user_agent.split("/")[0].strip()
+            else:
+                partner_name = custom_user_agent
+
+            # Validate partner name: alphanumeric, underscore, hyphen only
+            if (
+                partner_name
+                and partner_name.replace("_", "").replace("-", "").isalnum()
+            ):
+                return f"{partner_name}_litellm/{version}"
+
+        # Default: just litellm
+        return f"litellm/{version}"
+
     def _get_api_base(self, api_base: Optional[str]) -> str:
+        """
+        Get the Databricks API base URL.
+
+        If not provided, attempts to get it from the Databricks SDK.
+        """
         if api_base is None:
             try:
                 from databricks.sdk import WorkspaceClient
 
                 databricks_client = WorkspaceClient()
-
-                api_base = (
-                    api_base or f"{databricks_client.config.host}/serving-endpoints"
-                )
-
+                api_base = f"{databricks_client.config.host}/serving-endpoints"
                 return api_base
             except ImportError:
                 raise DatabricksException(
@@ -30,12 +202,87 @@ class DatabricksBase:
                 )
         return api_base
 
+    def _get_oauth_m2m_token(
+        self,
+        api_base: str,
+        client_id: str,
+        client_secret: str,
+    ) -> str:
+        """
+        Obtain an OAuth M2M access token using client credentials flow.
+
+        This is the recommended authentication method for production integrations
+        per Databricks Partner requirements.
+
+        Args:
+            api_base: Databricks workspace URL
+            client_id: OAuth client ID (Service Principal application ID)
+            client_secret: OAuth client secret
+
+        Returns:
+            Access token string
+
+        Raises:
+            DatabricksException: If token request fails
+        """
+        import requests
+
+        # Extract workspace URL from api_base
+        workspace_url = api_base.rstrip("/")
+        if "/serving-endpoints" in workspace_url:
+            workspace_url = workspace_url.replace("/serving-endpoints", "")
+
+        token_url = f"{workspace_url}/oidc/v1/token"
+
+        try:
+            response = requests.post(
+                token_url,
+                data={
+                    "grant_type": "client_credentials",
+                    "scope": "all-apis",
+                },
+                auth=(client_id, client_secret),
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=30,
+            )
+        except requests.RequestException as e:
+            raise DatabricksException(
+                status_code=500,
+                message=f"OAuth M2M token request failed: {str(e)}",
+            )
+
+        if response.status_code != 200:
+            raise DatabricksException(
+                status_code=response.status_code,
+                message=f"OAuth M2M token request failed: {response.text}",
+            )
+
+        token_data = response.json()
+        return token_data["access_token"]
+
     def _get_databricks_credentials(
         self, api_key: Optional[str], api_base: Optional[str], headers: Optional[dict]
     ) -> Tuple[str, dict]:
+        """
+        Get Databricks credentials using the Databricks SDK.
+
+        Also registers LiteLLM as a partner for proper telemetry attribution
+        in Databricks system.access.audit table.
+
+        Args:
+            api_key: Optional API key (PAT)
+            api_base: Optional API base URL
+            headers: Optional existing headers
+
+        Returns:
+            Tuple of (api_base, headers)
+        """
         headers = headers or {"Content-Type": "application/json"}
         try:
-            from databricks.sdk import WorkspaceClient
+            from databricks.sdk import WorkspaceClient, useragent
+
+            # Register LiteLLM as partner for Databricks telemetry attribution
+            useragent.with_partner("litellm")
 
             databricks_client = WorkspaceClient()
 
@@ -66,14 +313,53 @@ class DatabricksBase:
         endpoint_type: Literal["chat_completions", "embeddings"],
         custom_endpoint: Optional[bool],
         headers: Optional[dict],
+        custom_user_agent: Optional[str] = None,
     ) -> Tuple[str, dict]:
-        if api_key is None and not headers:  # handle empty headers
+        """
+        Validate and configure the Databricks environment.
+
+        Authentication priority:
+        1. OAuth M2M (DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET) - Recommended
+        2. PAT (DATABRICKS_API_KEY) - Supported for development
+        3. Databricks SDK automatic auth - Fallback (uses unified auth)
+
+        Args:
+            api_key: Personal access token (PAT)
+            api_base: Databricks workspace URL with /serving-endpoints
+            endpoint_type: Type of endpoint (chat_completions or embeddings)
+            custom_endpoint: Whether using a custom endpoint URL
+            headers: Existing headers dict
+            custom_user_agent: Optional custom user agent to prefix
+
+        Returns:
+            Tuple of (api_base, headers) with authentication configured
+        """
+        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")
+
+        # Determine api_base first
+        if api_base is None:
+            api_base = os.getenv("DATABRICKS_API_BASE")
+
+        if client_id and client_secret and api_base:
+            # Use OAuth M2M flow (preferred for production)
+            verbose_logger.debug("Using OAuth M2M authentication for Databricks")
+            access_token = self._get_oauth_m2m_token(api_base, client_id, client_secret)
+            headers = headers or {}
+            headers["Authorization"] = f"Bearer {access_token}"
+            headers["Content-Type"] = "application/json"
+        elif api_key is None and not headers:
             if custom_endpoint is True:
                 raise DatabricksException(
                     status_code=400,
                     message="Missing API Key - A call is being made to LLM Provider but no key is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params",
                 )
             else:
+                # Fallback to Databricks SDK (registers partner telemetry)
+                verbose_logger.debug("Using Databricks SDK for authentication")
                 api_base, headers = self._get_databricks_credentials(
                     api_base=api_base, api_key=api_key, headers=headers
                 )
@@ -101,8 +387,17 @@ class DatabricksBase:
         if api_key is not None:
             headers["Authorization"] = f"Bearer {api_key}"
 
+        # Set User-Agent with optional custom prefix
+        headers["User-Agent"] = self._build_user_agent(custom_user_agent)
+
+        # Debug logging with redaction (never log actual tokens)
+        verbose_logger.debug(
+            f"Databricks request headers: {self.redact_headers_for_logging(headers)}"
+        )
+
         if endpoint_type == "chat_completions" and custom_endpoint is not True:
             api_base = "{}/chat/completions".format(api_base)
         elif endpoint_type == "embeddings" and custom_endpoint is not True:
             api_base = "{}/embeddings".format(api_base)
+
         return api_base, headers

litellm/llms/databricks/embed/handler.py

@@ -2,6 +2,7 @@
 Calling logic for Databricks embeddings
 """
 
+import os
 from typing import Optional
 
 from litellm.utils import EmbeddingResponse
@@ -26,12 +27,23 @@ class DatabricksEmbeddingHandler(OpenAILikeEmbeddingHandler, DatabricksBase):
         custom_endpoint: Optional[bool] = None,
         headers: Optional[dict] = None,
     ) -> EmbeddingResponse:
+        # Check for custom user agent in optional_params or environment
+        # This allows partners building on LiteLLM to set their own telemetry
+        # Use pop() to remove these keys so they don't get sent to the API
+        custom_user_agent = (
+            optional_params.pop("user_agent", None)
+            or optional_params.pop("databricks_user_agent", None)
+            or os.getenv("LITELLM_USER_AGENT")
+            or os.getenv("DATABRICKS_USER_AGENT")
+        )
+
         api_base, headers = self.databricks_validate_environment(
             api_base=api_base,
             api_key=api_key,
             endpoint_type="embeddings",
             custom_endpoint=custom_endpoint,
             headers=headers,
+            custom_user_agent=custom_user_agent,
         )
         return super().embedding(
             model=model,

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
 
 [[package]]
 name = "aiofiles"

tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py

@@ -94,12 +94,13 @@ def test_transform_choices_without_signature():
     assert thinking_block["type"] == "thinking"
     assert thinking_block["thinking"] == "i'm thinking without signature."
 
+
 def test_convert_anthropic_tool_to_databricks_tool_with_description():
     config = DatabricksConfig()
     anthropic_tool = {
         "name": "test_tool",
         "description": "test description",
-        "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}}
+        "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}},
     }
 
     databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool)
@@ -113,7 +114,7 @@ def test_convert_anthropic_tool_to_databricks_tool_without_description():
     config = DatabricksConfig()
     anthropic_tool = {
         "name": "test_tool",
-        "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}}
+        "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}},
     }
 
     databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool)
@@ -122,6 +123,7 @@ def test_convert_anthropic_tool_to_databricks_tool_without_description():
     assert databricks_tool["type"] == "function"
     assert databricks_tool["function"].get("description") is None
 
+
 def test_transform_choices_with_citations():
     config = DatabricksConfig()
     databricks_choices = [

tests/test_litellm/llms/databricks/databricks_config.template.txt

@@ -0,0 +1,78 @@
+# Databricks Configuration Template for LiteLLM Testing
+# =====================================================
+#
+# Copy this file to your preferred location and fill in your credentials:
+#   cp databricks_config.template.txt /path/to/databricks_config.txt
+#
+# Then update the CONFIG_FILE path in test_databricks_integration.py
+#
+# Lines starting with # are comments and will be ignored
+# Only lines with KEY=VALUE format (where VALUE is not empty) will be read
+
+# ==============================================================================
+# DATABRICKS WORKSPACE CONFIGURATION (Required)
+# ==============================================================================
+
+# Your Databricks workspace URL (without /serving-endpoints suffix)
+# Example: https://adb-1234567890123456.7.azuredatabricks.net
+DATABRICKS_HOST=
+
+# API Base URL for serving endpoints (usually {host}/serving-endpoints)
+# Example: https://adb-1234567890123456.7.azuredatabricks.net/serving-endpoints
+DATABRICKS_API_BASE=
+
+# ==============================================================================
+# AUTHENTICATION METHOD 1: OAuth M2M (Recommended for Production)
+# Use Service Principal credentials
+# ==============================================================================
+
+# Service Principal Application/Client ID
+# Example: 12345678-1234-1234-1234-123456789012
+DATABRICKS_CLIENT_ID=
+
+# Service Principal Secret
+# Example: your-client-secret-value
+DATABRICKS_CLIENT_SECRET=
+
+# ==============================================================================
+# AUTHENTICATION METHOD 2: Personal Access Token (PAT)
+# For development and testing
+# ==============================================================================
+
+# Personal Access Token (starts with 'dapi')
+# Example: dapi_your_token_here
+DATABRICKS_API_KEY=
+
+# ==============================================================================
+# MODEL CONFIGURATION
+# ==============================================================================
+
+# Model to use for testing chat completions
+# Example: databricks-gpt-oss-120b, databricks-meta-llama-3-1-70b-instruct
+TEST_CHAT_MODEL=databricks-gpt-oss-120b
+
+# Model to use for testing embeddings (optional)
+# Example: databricks-bge-large-en
+TEST_EMBEDDING_MODEL=databricks-bge-large-en
+
+# ==============================================================================
+# OPTIONAL: Custom User-Agent for Partner Attribution Testing
+# ==============================================================================
+
+# Custom user agent string to test partner attribution
+# Example: mycompany/1.0.0
+# This will result in User-Agent: mycompany_litellm/{version}
+# Leave empty to use default: litellm/{version}
+CUSTOM_USER_AGENT=
+
+# ==============================================================================
+# TEST SETTINGS
+# ==============================================================================
+
+# Which authentication method to test: oauth, pat, sdk, or all
+# oauth = Use DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET
+# pat = Use DATABRICKS_API_KEY
+# sdk = Use Databricks SDK automatic authentication (~/.databrickscfg)
+# all = Test all three methods (oauth, pat, sdk) in sequence
+TEST_AUTH_METHOD=pat
+

tests/test_litellm/llms/databricks/test_databricks_partner_integration.py

@@ -0,0 +1,662 @@
+"""
+Unit Tests for Databricks Partner Integration Features
+=======================================================
+
+These tests are designed for automated CI/CD pipelines and do NOT require
+real Databricks credentials. All external calls are mocked.
+
+For integration tests that use real Databricks credentials, see:
+    test_databricks_integration.py
+
+Features Tested:
+    - User-Agent building with partner prefixing (Databricks partner telemetry)
+    - Token/sensitive data redaction for secure logging
+    - OAuth M2M (Machine-to-Machine) authentication flow
+    - Databricks SDK partner telemetry registration
+    - Authentication priority (OAuth M2M > PAT > SDK)
+
+Run with:
+    pytest test_databricks_partner_integration.py -v
+
+These tests align with Databricks Partner Architecture best practices:
+    https://github.com/databrickslabs/partner-architecture
+"""
+
+import json
+import os
+import sys
+
+import pytest
+from unittest.mock import MagicMock, patch, Mock
+
+sys.path.insert(
+    0, os.path.abspath("../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.databricks.common_utils import DatabricksBase, DatabricksException
+
+
+class TestBuildUserAgent:
+    """Test cases for User-Agent string building."""
+
+    def test_default_user_agent(self):
+        """No custom user agent returns litellm/{version}."""
+        ua = DatabricksBase._build_user_agent(None)
+        assert ua.startswith("litellm/")
+        assert "_" not in ua.split("/")[0]
+
+    def test_custom_user_agent_with_version(self):
+        """Custom user agent with version extracts partner name."""
+        ua = DatabricksBase._build_user_agent("mycompany/1.0.0")
+        assert ua.startswith("mycompany_litellm/")
+        # Verify the version is litellm's, not the custom one
+        assert "/1.0.0" not in ua or "mycompany_litellm/1.0.0" not in ua
+
+    def test_custom_user_agent_without_version(self):
+        """Custom user agent without version still works."""
+        ua = DatabricksBase._build_user_agent("mycompany")
+        assert ua.startswith("mycompany_litellm/")
+
+    def test_custom_user_agent_with_underscore(self):
+        """Partner names with underscores are preserved."""
+        ua = DatabricksBase._build_user_agent("my_company/2.0.0")
+        assert ua.startswith("my_company_litellm/")
+
+    def test_custom_user_agent_with_hyphen(self):
+        """Partner names with hyphens are preserved."""
+        ua = DatabricksBase._build_user_agent("my-company/2.0.0")
+        assert ua.startswith("my-company_litellm/")
+
+    def test_custom_user_agent_ignores_custom_version(self):
+        """Custom version is ignored, litellm version is used."""
+        ua = DatabricksBase._build_user_agent("partner/99.99.99")
+        parts = ua.split("/")
+        assert parts[0] == "partner_litellm"
+        assert parts[1] != "99.99.99"
+
+    def test_empty_string_returns_default(self):
+        """Empty string returns default user agent."""
+        ua = DatabricksBase._build_user_agent("")
+        assert ua.startswith("litellm/")
+        assert "_" not in ua.split("/")[0]
+
+    def test_whitespace_only_returns_default(self):
+        """Whitespace-only string returns default user agent."""
+        ua = DatabricksBase._build_user_agent("   ")
+        assert ua.startswith("litellm/")
+        assert "_" not in ua.split("/")[0]
+
+    def test_invalid_partner_name_returns_default(self):
+        """Invalid partner names (special chars) return default."""
+        ua = DatabricksBase._build_user_agent("my@company/1.0.0")
+        assert ua.startswith("litellm/")
+
+    def test_partner_with_numbers(self):
+        """Partner names with numbers work."""
+        ua = DatabricksBase._build_user_agent("company123/1.0.0")
+        assert ua.startswith("company123_litellm/")
+
+
+class TestRedactSensitiveData:
+    """Test cases for sensitive data redaction."""
+
+    def test_redact_bearer_token_in_string(self):
+        """Bearer tokens are redacted in strings."""
+        result = DatabricksBase.redact_sensitive_data("Bearer dapi12345abcdef")
+        assert "dapi12345abcdef" not in result
+        assert "[REDACTED]" in result
+
+    def test_redact_dict_with_authorization(self):
+        """Dict with authorization key is redacted."""
+        data = {"Authorization": "Bearer secret123", "other": "value"}
+        result = DatabricksBase.redact_sensitive_data(data)
+        assert result["Authorization"] == "[REDACTED]"
+        assert result["other"] == "value"
+
+    def test_redact_nested_dict(self):
+        """Nested dicts with sensitive keys are redacted."""
+        data = {"config": {"api_key": "secret", "name": "test"}}
+        result = DatabricksBase.redact_sensitive_data(data)
+        assert result["config"]["api_key"] == "[REDACTED]"
+        assert result["config"]["name"] == "test"
+
+    def test_redact_pat_token(self):
+        """Databricks PAT tokens are redacted."""
+        result = DatabricksBase.redact_sensitive_data(
+            "Using token dapi_fake_test_token_value"
+        )
+        assert "dapi_fake_test_token_value" not in result
+        assert "[REDACTED_PAT]" in result
+
+    def test_redact_client_secret(self):
+        """Client secrets are redacted."""
+        data = {"client_secret": "my-super-secret-value"}
+        result = DatabricksBase.redact_sensitive_data(data)
+        assert result["client_secret"] == "[REDACTED]"
+
+    def test_redact_list_of_dicts(self):
+        """Lists containing dicts with sensitive data are redacted."""
+        data = [{"api_key": "secret1"}, {"name": "test"}]
+        result = DatabricksBase.redact_sensitive_data(data)
+        assert result[0]["api_key"] == "[REDACTED]"
+        assert result[1]["name"] == "test"
+
+    def test_redact_none_returns_none(self):
+        """None input returns None."""
+        assert DatabricksBase.redact_sensitive_data(None) is None
+
+    def test_redact_preserves_non_sensitive_data(self):
+        """Non-sensitive data is preserved."""
+        data = {"model": "dbrx", "temperature": 0.7, "messages": ["hello"]}
+        result = DatabricksBase.redact_sensitive_data(data)
+        assert result == data
+
+
+class TestRedactHeadersForLogging:
+    """Test cases for header redaction."""
+
+    def test_authorization_header_partially_shown(self):
+        """Authorization header shows first 8 chars then redacts."""
+        headers = {"Authorization": "Bearer dapi123456789abcdef"}
+        result = DatabricksBase.redact_headers_for_logging(headers)
+        assert result["Authorization"].startswith("Bearer d")
+        assert "[REDACTED]" in result["Authorization"]
+
+    def test_short_authorization_header_fully_redacted(self):
+        """Short authorization values are fully redacted."""
+        headers = {"Authorization": "short"}
+        result = DatabricksBase.redact_headers_for_logging(headers)
+        assert result["Authorization"] == "[REDACTED]"
+
+    def test_non_sensitive_headers_preserved(self):
+        """Non-sensitive headers are not modified."""
+        headers = {"Content-Type": "application/json", "User-Agent": "test/1.0"}
+        result = DatabricksBase.redact_headers_for_logging(headers)
+        assert result["Content-Type"] == "application/json"
+        assert result["User-Agent"] == "test/1.0"
+
+    def test_empty_headers_returns_empty(self):
+        """Empty headers dict returns empty dict."""
+        assert DatabricksBase.redact_headers_for_logging({}) == {}
+
+    def test_none_headers_returns_empty(self):
+        """None headers returns empty dict."""
+        assert DatabricksBase.redact_headers_for_logging(None) == {}
+
+    def test_x_api_key_header_redacted(self):
+        """X-API-Key header is redacted."""
+        headers = {"X-API-Key": "my-api-key-12345"}
+        result = DatabricksBase.redact_headers_for_logging(headers)
+        assert "[REDACTED]" in result["X-API-Key"]
+
+
+class TestOAuthM2M:
+    """Test cases for OAuth M2M authentication."""
+
+    def test_oauth_m2m_token_success(self):
+        """OAuth M2M token is successfully obtained."""
+        databricks_base = DatabricksBase()
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"access_token": "test-access-token"}
+
+        with patch("requests.post", return_value=mock_response) as mock_post:
+            token = databricks_base._get_oauth_m2m_token(
+                api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
+                client_id="test-client-id",
+                client_secret="test-client-secret",
+            )
+
+            assert token == "test-access-token"
+            mock_post.assert_called_once()
+            call_args = mock_post.call_args
+            assert "oidc/v1/token" in call_args[0][0]
+            assert call_args[1]["data"]["grant_type"] == "client_credentials"
+
+    def test_oauth_m2m_token_failure(self):
+        """OAuth M2M raises exception on failure."""
+        databricks_base = DatabricksBase()
+
+        mock_response = Mock()
+        mock_response.status_code = 401
+        mock_response.text = "Unauthorized"
+
+        with patch("requests.post", return_value=mock_response):
+            with pytest.raises(DatabricksException) as exc_info:
+                databricks_base._get_oauth_m2m_token(
+                    api_base="https://adb-123.azuredatabricks.net",
+                    client_id="bad-client-id",
+                    client_secret="bad-secret",
+                )
+            assert exc_info.value.status_code == 401
+
+    def test_oauth_m2m_strips_serving_endpoints(self):
+        """OAuth M2M correctly strips /serving-endpoints from URL."""
+        databricks_base = DatabricksBase()
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"access_token": "token"}
+
+        with patch("requests.post", return_value=mock_response) as mock_post:
+            databricks_base._get_oauth_m2m_token(
+                api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
+                client_id="id",
+                client_secret="secret",
+            )
+
+            call_url = mock_post.call_args[0][0]
+            assert "/serving-endpoints" not in call_url
+            assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
+
+
+class TestValidateEnvironmentWithOAuth:
+    """Test OAuth M2M is used when credentials are available."""
+
+    def test_oauth_used_when_credentials_set(self, monkeypatch):
+        """OAuth M2M is used when client_id and client_secret are set."""
+        monkeypatch.setenv("DATABRICKS_CLIENT_ID", "test-client-id")
+        monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "test-secret")
+        monkeypatch.setenv(
+            "DATABRICKS_API_BASE", "https://adb-123.net/serving-endpoints"
+        )
+
+        databricks_base = DatabricksBase()
+
+        with patch.object(
+            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+        ) as mock_oauth:
+            api_base, headers = databricks_base.databricks_validate_environment(
+                api_key=None,
+                api_base=None,
+                endpoint_type="chat_completions",
+                custom_endpoint=False,
+                headers=None,
+            )
+
+            mock_oauth.assert_called_once()
+            assert headers["Authorization"] == "Bearer oauth-token"
+
+    def test_pat_used_when_api_key_set(self, monkeypatch):
+        """PAT is used when api_key is provided."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="dapi-test-key",
+            api_base="https://adb-123.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+        )
+
+        assert headers["Authorization"] == "Bearer dapi-test-key"
+
+
+class TestValidateEnvironmentUserAgent:
+    """Test User-Agent is correctly set in validate_environment."""
+
+    def test_default_user_agent(self, monkeypatch):
+        """Default user agent is set when no custom agent provided."""
+        monkeypatch.delenv("DATABRICKS_USER_AGENT", raising=False)
+        monkeypatch.delenv("LITELLM_USER_AGENT", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://adb-123.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+            custom_user_agent=None,
+        )
+
+        assert headers["User-Agent"].startswith("litellm/")
+        assert "_" not in headers["User-Agent"].split("/")[0]
+
+    def test_custom_user_agent_via_param(self, monkeypatch):
+        """Custom user agent is prefixed when passed as parameter."""
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://adb-123.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+            custom_user_agent="mycompany/1.0.0",
+        )
+
+        assert headers["User-Agent"].startswith("mycompany_litellm/")
+
+
+class TestSDKPartnerTelemetry:
+    """Test that SDK partner telemetry is registered."""
+
+    def test_sdk_partner_registered(self):
+        """useragent.with_partner is called when using SDK."""
+        databricks_base = DatabricksBase()
+
+        mock_workspace_client = MagicMock()
+        mock_workspace_client.config.host = "https://adb-123.net"
+        mock_workspace_client.config.authenticate.return_value = {
+            "Authorization": "Bearer token"
+        }
+
+        with patch(
+            "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client
+        ):
+            with patch("databricks.sdk.useragent.with_partner") as mock_with_partner:
+                databricks_base._get_databricks_credentials(
+                    api_key=None,
+                    api_base=None,
+                    headers=None,
+                )
+
+                mock_with_partner.assert_called_once_with("litellm")
+
+
+class TestUserAgentFromEnvironment:
+    """Test User-Agent is correctly picked up from environment variables."""
+
+    def test_user_agent_from_databricks_env_var(self, monkeypatch):
+        """DATABRICKS_USER_AGENT environment variable is used."""
+        monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner")
+        monkeypatch.delenv("LITELLM_USER_AGENT", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://adb-123.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+            custom_user_agent="envpartner",  # Simulating what transformation.py passes
+        )
+
+        assert headers["User-Agent"].startswith("envpartner_litellm/")
+
+    def test_custom_param_takes_precedence(self, monkeypatch):
+        """Custom user_agent parameter takes precedence over environment."""
+        monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner")
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://adb-123.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+            custom_user_agent="parampartner/1.0.0",
+        )
+
+        assert headers["User-Agent"].startswith("parampartner_litellm/")
+
+
+class TestLiteLLMCompletionUserAgent:
+    """Test User-Agent is correctly passed through LiteLLM completion calls."""
+
+    def test_completion_passes_user_agent_to_headers(self):
+        """litellm.completion() correctly passes user_agent to request headers."""
+        from litellm.llms.databricks.chat.transformation import DatabricksConfig
+
+        config = DatabricksConfig()
+        optional_params = {"user_agent": "testpartner/1.0.0"}
+
+        # Mock the validation to capture what headers are set
+        with patch.object(
+            config,
+            "databricks_validate_environment",
+            return_value=(
+                "https://test.net/serving-endpoints/chat/completions",
+                {
+                    "Authorization": "Bearer test",
+                    "User-Agent": "testpartner_litellm/1.0.0",
+                },
+            ),
+        ) as mock_validate:
+            result = config.validate_environment(
+                headers={},
+                model="databricks/test-model",
+                messages=[],
+                optional_params=optional_params,
+                litellm_params={},
+                api_key="test-key",
+                api_base="https://test.net/serving-endpoints",
+            )
+
+            # Verify user_agent was passed to databricks_validate_environment
+            mock_validate.assert_called_once()
+            call_kwargs = mock_validate.call_args[1]
+            assert call_kwargs.get("custom_user_agent") == "testpartner/1.0.0"
+
+    def test_user_agent_removed_from_optional_params(self):
+        """user_agent is removed from optional_params so it's not sent to API."""
+        from litellm.llms.databricks.chat.transformation import DatabricksConfig
+
+        config = DatabricksConfig()
+        optional_params = {
+            "user_agent": "testpartner/1.0.0",
+            "temperature": 0.7,
+        }
+
+        with patch.object(
+            config,
+            "databricks_validate_environment",
+            return_value=(
+                "https://test.net/chat/completions",
+                {"Authorization": "Bearer test", "User-Agent": "test"},
+            ),
+        ):
+            config.validate_environment(
+                headers={},
+                model="databricks/test-model",
+                messages=[],
+                optional_params=optional_params,
+                litellm_params={},
+                api_key="test-key",
+                api_base="https://test.net/serving-endpoints",
+            )
+
+            # user_agent should be removed from optional_params
+            assert "user_agent" not in optional_params
+            # Other params should remain
+            assert optional_params.get("temperature") == 0.7
+
+
+class TestLiteLLMEmbeddingUserAgent:
+    """Test User-Agent is correctly passed through LiteLLM embedding calls."""
+
+    def test_embedding_passes_user_agent_to_headers(self):
+        """litellm.embedding() correctly passes user_agent to request headers."""
+        from litellm.llms.databricks.embed.handler import DatabricksEmbeddingHandler
+
+        handler = DatabricksEmbeddingHandler()
+        optional_params = {"user_agent": "embedpartner/1.0.0"}
+
+        with patch.object(
+            handler,
+            "databricks_validate_environment",
+            return_value=(
+                "https://test.net/serving-endpoints/embeddings",
+                {
+                    "Authorization": "Bearer test",
+                    "User-Agent": "embedpartner_litellm/1.0.0",
+                },
+            ),
+        ) as mock_validate:
+            with patch(
+                "litellm.llms.openai_like.embedding.handler.OpenAILikeEmbeddingHandler.embedding"
+            ):
+                try:
+                    handler.embedding(
+                        model="databricks/test-model",
+                        input=["test"],
+                        timeout=30,
+                        api_key="test-key",
+                        api_base="https://test.net/serving-endpoints",
+                        optional_params=optional_params,
+                    )
+                except Exception:
+                    pass  # We just want to verify the mock was called
+
+            # Verify user_agent was passed
+            if mock_validate.called:
+                call_kwargs = mock_validate.call_args[1]
+                assert call_kwargs.get("custom_user_agent") == "embedpartner/1.0.0"
+
+
+class TestAuthenticationPriority:
+    """Test that authentication methods are used in correct priority order."""
+
+    def test_oauth_used_when_no_api_key_provided(self, monkeypatch):
+        """OAuth M2M is used when OAuth creds are set and no api_key is provided."""
+        monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id")
+        monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret")
+        monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints")
+
+        databricks_base = DatabricksBase()
+
+        with patch.object(
+            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+        ) as mock_oauth:
+            api_base, headers = databricks_base.databricks_validate_environment(
+                api_key=None,  # No PAT provided - OAuth should be used
+                api_base=None,
+                endpoint_type="chat_completions",
+                custom_endpoint=False,
+                headers=None,
+            )
+
+            # OAuth should be used
+            mock_oauth.assert_called_once()
+            assert headers["Authorization"] == "Bearer oauth-token"
+
+    def test_explicit_pat_takes_priority_over_oauth_env(self, monkeypatch):
+        """Explicit api_key takes priority over OAuth token in final headers."""
+        monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id")
+        monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret")
+        monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints")
+
+        databricks_base = DatabricksBase()
+
+        # Mock the OAuth call - it will be attempted but PAT should override
+        with patch.object(
+            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+        ):
+            api_base, headers = databricks_base.databricks_validate_environment(
+                api_key="dapi-explicit-pat",
+                api_base=None,
+                endpoint_type="chat_completions",
+                custom_endpoint=False,
+                headers=None,
+            )
+
+            # PAT should override OAuth token since api_key was explicitly provided
+            assert headers["Authorization"] == "Bearer dapi-explicit-pat"
+
+    def test_pat_used_when_no_oauth_credentials(self, monkeypatch):
+        """PAT is used when OAuth credentials are not set."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="dapi-pat-token",
+            api_base="https://test.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+        )
+
+        assert headers["Authorization"] == "Bearer dapi-pat-token"
+
+    def test_sdk_fallback_when_no_credentials(self, monkeypatch):
+        """Databricks SDK is used when no API key or OAuth credentials."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+        monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        mock_workspace_client = MagicMock()
+        mock_workspace_client.config.host = "https://adb-123.net"
+        mock_workspace_client.config.authenticate.return_value = {
+            "Authorization": "Bearer sdk-token"
+        }
+
+        with patch(
+            "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client
+        ):
+            with patch("databricks.sdk.useragent.with_partner"):
+                api_base, headers = databricks_base.databricks_validate_environment(
+                    api_key=None,
+                    api_base=None,
+                    endpoint_type="chat_completions",
+                    custom_endpoint=False,
+                    headers=None,
+                )
+
+                assert "Authorization" in headers
+
+
+class TestEndpointURLConstruction:
+    """Test that endpoint URLs are correctly constructed."""
+
+    def test_chat_completions_endpoint(self, monkeypatch):
+        """Chat completions endpoint is correctly appended."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://test.net/serving-endpoints",
+            endpoint_type="chat_completions",
+            custom_endpoint=False,
+            headers=None,
+        )
+
+        assert api_base.endswith("/chat/completions")
+
+    def test_embeddings_endpoint(self, monkeypatch):
+        """Embeddings endpoint is correctly appended."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://test.net/serving-endpoints",
+            endpoint_type="embeddings",
+            custom_endpoint=False,
+            headers=None,
+        )
+
+        assert api_base.endswith("/embeddings")
+
+    def test_custom_endpoint_not_modified(self, monkeypatch):
+        """Custom endpoints are not modified."""
+        monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+        monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+
+        databricks_base = DatabricksBase()
+
+        api_base, headers = databricks_base.databricks_validate_environment(
+            api_key="test-key",
+            api_base="https://test.net/custom/endpoint",
+            endpoint_type="chat_completions",
+            custom_endpoint=True,
+            headers=None,
+        )
+
+        assert api_base == "https://test.net/custom/endpoint"