From 3a1baae45cb58a7050cd28a6f553010ed1c39768 Mon Sep 17 00:00:00 2001 From: prasadkona Date: Mon, 22 Dec 2025 11:59:20 -0800 Subject: [PATCH] feat(databricks): Add enhanced authentication, security features, and custom user-agent support - Add OAuth M2M (Machine-to-Machine) authentication via DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET - Add Databricks SDK auto-auth with automatic credential discovery - Add sensitive data redaction for secure logging (tokens, API keys, secrets) - Add custom user_agent parameter for partner attribution in Databricks telemetry - Support user_agent in LiteLLM Proxy via config.yaml litellm_params - Add 49 mocked unit tests for all new functionality - Add 13 E2E tests for real-world validation (skipped in CI) - Update documentation with new features and examples --- docs/my-website/docs/providers/databricks.md | 94 ++ .../llms/databricks/chat/transformation.py | 45 +- litellm/llms/databricks/common_utils.py | 311 ++++- litellm/llms/databricks/embed/handler.py | 12 + poetry.lock | 10 +- .../test_databricks_chat_transformation.py | 6 +- .../databricks/databricks_config.template.txt | 78 ++ .../llms/databricks/test_databricks_e2e.py | 1029 +++++++++++++++++ .../test_databricks_partner_integration.py | 662 +++++++++++ 9 files changed, 2218 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/llms/databricks/databricks_config.template.txt create mode 100644 tests/test_litellm/llms/databricks/test_databricks_e2e.py create mode 100644 tests/test_litellm/llms/databricks/test_databricks_partner_integration.py diff --git a/docs/my-website/docs/providers/databricks.md b/docs/my-website/docs/providers/databricks.md index 921b06a17b7..2791d55dff1 100644 --- a/docs/my-website/docs/providers/databricks.md +++ b/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 @@ -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 ``` diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index ac3be0c3518..2b7f5dd5995 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/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 diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 1353b5b13f6..608f29a03a7 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/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 diff --git a/litellm/llms/databricks/embed/handler.py b/litellm/llms/databricks/embed/handler.py index 2eabcdbc866..227824f72d0 100644 --- a/litellm/llms/databricks/embed/handler.py +++ b/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, diff --git a/poetry.lock b/poetry.lock index 5313167def4..4ae5cf01079 100644 --- a/poetry.lock +++ b/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" @@ -3138,15 +3138,15 @@ openai = ["openai (>=0.27.8)"] [[package]] name = "litellm-enterprise" -version = "0.1.25" +version = "0.1.27" description = "Package for LiteLLM Enterprise features" optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_enterprise-0.1.25-py3-none-any.whl", hash = "sha256:80c8f1996846453ad309e74cd6d2659d9508320370df5d462d34326b06401c4d"}, - {file = "litellm_enterprise-0.1.25.tar.gz", hash = "sha256:1c82178b8e2c85f47b31910fd103a322b46d6caea44cd7a8c80b00fdcfeacd22"}, + {file = "litellm_enterprise-0.1.27-py3-none-any.whl", hash = "sha256:41b9d41d04123f492060a742091006dc1d182b54ce3a1c0e18ee75d623c63e91"}, + {file = "litellm_enterprise-0.1.27.tar.gz", hash = "sha256:aa40c87f7c8df64beb79e75f71e1b5c0a458350efa68527e3491e6f27f2cbd57"}, ] [[package]] @@ -8051,4 +8051,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "a102d24777f1c438dcf15055abeef385722a9603cb3c3d3643c86190b7534c47" +content-hash = "996152bfbb1d7870a4b6f1837d7ff556320d6e817198ecec4198ed99539e848b" diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index a14683fac17..f437b8405f7 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/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 = [ diff --git a/tests/test_litellm/llms/databricks/databricks_config.template.txt b/tests/test_litellm/llms/databricks/databricks_config.template.txt new file mode 100644 index 00000000000..7352fdbc773 --- /dev/null +++ b/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 + diff --git a/tests/test_litellm/llms/databricks/test_databricks_e2e.py b/tests/test_litellm/llms/databricks/test_databricks_e2e.py new file mode 100644 index 00000000000..669f9e94639 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_e2e.py @@ -0,0 +1,1029 @@ +""" +End-to-End Tests for Databricks LiteLLM Integration +==================================================== + +⚠️ WARNING: These tests require REAL Databricks credentials and make ACTUAL API calls. + They are NOT suitable for automated CI/CD pipelines. + +For unit tests that use mocks and don't require credentials, see: + test_databricks_partner_integration.py + +Purpose: + - Validate actual API connectivity with Databricks + - Test all authentication methods (OAuth M2M, PAT, SDK) + - Verify User-Agent strings appear correctly in Databricks audit logs + - Test chat completions and embeddings with real models + - Test different SDK integration methods with custom user agents + +LiteLLM Integration Tests: + This test file includes tests for different ways of calling Databricks via LiteLLM: + + 1. LiteLLM SDK Direct - Using litellm.completion() with user_agent parameter + 2. LangChain + LiteLLM - Using ChatLiteLLM wrapper (requires langchain-community) + 3. LiteLLM Async - Using litellm.acompletion() async API + 4. LiteLLM Streaming - Using litellm.completion() with stream=True + 5. LiteLLM Embedding - Using litellm.embedding() with user_agent parameter + + All tests use the CUSTOM_USER_AGENT value from the config file and call + Databricks endpoints through LiteLLM's unified interface. + +Prerequisites: + - Valid Databricks workspace access + - Configured credentials (OAuth Service Principal, PAT, or Databricks CLI) + - Access to serving endpoints (e.g., databricks-gpt-oss-120b) + +Optional Dependencies (for LiteLLM integration tests): + - pip install langchain-litellm # For LangChain tests (recommended) + +Setup: + 1. Copy the template to create your config file: + cp databricks_config.template.txt ~/.databricks_litellm_config.txt + + 2. Edit the config file with your Databricks credentials: + - DATABRICKS_API_BASE (required) + - DATABRICKS_HOST (required for Databricks SDK tests) + - DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET (for OAuth) + - DATABRICKS_API_KEY (for PAT) + - CUSTOM_USER_AGENT (for partner attribution tests) + + 3. Optionally set a custom config path: + export DATABRICKS_TEST_CONFIG=/path/to/your/config.txt + +Run with: + cd /path/to/litellm + python tests/test_litellm/llms/databricks/test_databricks_e2e.py + +Config Options: + TEST_AUTH_METHOD=oauth # Test OAuth M2M authentication + TEST_AUTH_METHOD=pat # Test Personal Access Token + TEST_AUTH_METHOD=sdk # Test Databricks SDK (~/.databrickscfg) + TEST_AUTH_METHOD=all # Test all three methods sequentially +""" + +import os +import sys + +import pytest + +# Skip all tests in this module during unit test runs (make test-unit) +# These are E2E tests that require real Databricks credentials +pytestmark = pytest.mark.skip( + reason="E2E tests require real Databricks credentials. Run directly with: " + "python tests/test_litellm/llms/databricks/test_databricks_e2e.py" +) + +# Add the litellm package to path +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +) + +# Config file path - can be overridden with DATABRICKS_TEST_CONFIG env var +DEFAULT_CONFIG_PATH = os.path.expanduser("~/.databricks_litellm_config.txt") +CONFIG_FILE = os.environ.get("DATABRICKS_TEST_CONFIG", DEFAULT_CONFIG_PATH) + + +def load_config(config_file: str) -> dict: + """Load configuration from file.""" + config = {} + + template_path = os.path.join( + os.path.dirname(__file__), "databricks_config.template.txt" + ) + + if not os.path.exists(config_file): + raise FileNotFoundError( + f"Config file not found: {config_file}\n\n" + f"To set up:\n" + f" 1. Copy the template:\n" + f" cp {template_path} {config_file}\n\n" + f" 2. Edit {config_file} with your Databricks credentials\n\n" + f" 3. Or set a custom path:\n" + f" export DATABRICKS_TEST_CONFIG=/your/path/config.txt" + ) + + with open(config_file, "r") as f: + for line in f: + line = line.strip() + # Skip comments and empty lines + if not line or line.startswith("#"): + continue + + # Parse KEY=VALUE + if "=" in line: + key, value = line.split("=", 1) + key = key.strip() + value = value.strip() + if value: # Only set if value is not empty + config[key] = value + + return config + + +def setup_environment(config: dict, auth_method: str): + """Set up environment variables based on auth method.""" + # Clear any existing Databricks env vars (including SDK-specific ones) + for var in [ + "DATABRICKS_API_KEY", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_API_BASE", + "DATABRICKS_USER_AGENT", + "LITELLM_USER_AGENT", + "DATABRICKS_TOKEN", + "DATABRICKS_HOST", + ]: # Added SDK env vars + os.environ.pop(var, None) + + # Set auth based on method + if auth_method == "oauth": + if ( + "DATABRICKS_CLIENT_ID" not in config + or "DATABRICKS_CLIENT_SECRET" not in config + ): + raise ValueError( + "OAuth auth requires DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET" + ) + # For OAuth, set the API base + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + os.environ["DATABRICKS_CLIENT_ID"] = config["DATABRICKS_CLIENT_ID"] + os.environ["DATABRICKS_CLIENT_SECRET"] = config["DATABRICKS_CLIENT_SECRET"] + print(" Auth method: OAuth M2M (Service Principal)") + + elif auth_method == "pat": + if "DATABRICKS_API_KEY" not in config: + raise ValueError("PAT auth requires DATABRICKS_API_KEY") + # For PAT, set the API base + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + os.environ["DATABRICKS_API_KEY"] = config["DATABRICKS_API_KEY"] + print(" Auth method: Personal Access Token (PAT)") + + elif auth_method == "sdk": + # For SDK mode, don't set any env vars - let SDK use ~/.databrickscfg + # But we still need to pass api_base to litellm, so set it if provided + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + print(" Auth method: Databricks SDK (automatic from ~/.databrickscfg)") + + else: + raise ValueError(f"Unknown auth method: {auth_method}") + + # Set custom user agent if provided + if "CUSTOM_USER_AGENT" in config: + os.environ["DATABRICKS_USER_AGENT"] = config["CUSTOM_USER_AGENT"] + print(f" Custom User-Agent: {config['CUSTOM_USER_AGENT']}") + + +def test_user_agent_building(): + """Test User-Agent string building.""" + print("\n" + "=" * 60) + print("TEST: User-Agent Building") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + # Test 1: Default + ua = DatabricksBase._build_user_agent(None) + print(f" Default: {ua}") + assert ua.startswith("litellm/"), f"Expected litellm/, got {ua}" + print(" ✓ Default user agent works") + + # Test 2: With partner + ua = DatabricksBase._build_user_agent("mycompany/1.0.0") + print(f" With partner: {ua}") + assert ua.startswith("mycompany_litellm/"), f"Expected mycompany_litellm/, got {ua}" + print(" ✓ Partner prefixing works") + + # Test 3: Partner without version + ua = DatabricksBase._build_user_agent("acme") + print(f" Without version: {ua}") + assert ua.startswith("acme_litellm/"), f"Expected acme_litellm/, got {ua}" + print(" ✓ Partner without version works") + + print(" ✓ All user agent tests passed!") + + +def test_token_redaction(): + """Test sensitive data redaction.""" + print("\n" + "=" * 60) + print("TEST: Token Redaction") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + # Test header redaction + headers = { + "Authorization": "Bearer dapi123456789abcdef", + "Content-Type": "application/json", + } + redacted = DatabricksBase.redact_headers_for_logging(headers) + print(f" Original: Authorization: Bearer dapi123456789abcdef") + print(f" Redacted: Authorization: {redacted['Authorization']}") + assert "[REDACTED]" in redacted["Authorization"] + assert redacted["Content-Type"] == "application/json" + print(" ✓ Header redaction works") + + # Test dict redaction + data = {"api_key": "secret123", "model": "dbrx"} + redacted = DatabricksBase.redact_sensitive_data(data) + assert redacted["api_key"] == "[REDACTED]" + assert redacted["model"] == "dbrx" + print(" ✓ Dict redaction works") + + # Test PAT redaction + text = "Token: dapi_fake_test_token_for_testing" + redacted = DatabricksBase.redact_sensitive_data(text) + assert "dapi_fake_test" not in redacted + print(" ✓ PAT string redaction works") + + print(" ✓ All redaction tests passed!") + + +def test_chat_completion(config: dict): + """Test chat completion with Databricks.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion") + print("=" * 60) + + import litellm + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" API Base: {os.environ.get('DATABRICKS_API_BASE', 'Not set')}") + + try: + response = litellm.completion( + model=full_model, + messages=[ + { + "role": "user", + "content": "Say 'Hello, LiteLLM test!' in exactly those words.", + } + ], + max_tokens=50, + temperature=0.1, + ) + + content = response.choices[0].message.content + print(f" Response: {content[:100]}...") + print(f" Model returned: {response.model}") + print(f" Usage: {response.usage}") + print(" ✓ Chat completion test passed!") + return True + + except Exception as e: + print(f" ✗ Chat completion failed: {e}") + return False + + +def test_chat_completion_default_user_agent(config: dict): + """Test chat completion with default user agent (no custom agent).""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with DEFAULT User-Agent") + print("=" * 60) + + import litellm + + # Clear any custom user agent from environment + saved_user_agent = os.environ.pop("DATABRICKS_USER_AGENT", None) + saved_litellm_ua = os.environ.pop("LITELLM_USER_AGENT", None) + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" Expected User-Agent: litellm/{version}") + print(f" (No custom user agent set)") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'default' only."}], + max_tokens=10, + # Note: NOT passing user_agent parameter + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Default user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is 'litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Default user-agent test failed: {e}") + return False + + finally: + # Restore environment variables + if saved_user_agent: + os.environ["DATABRICKS_USER_AGENT"] = saved_user_agent + if saved_litellm_ua: + os.environ["LITELLM_USER_AGENT"] = saved_litellm_ua + + +def test_chat_completion_with_custom_user_agent(config: dict): + """Test chat completion with custom user agent passed as parameter.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with Custom User-Agent (parameter)") + print("=" * 60) + + import litellm + + # Clear any env user agent to ensure parameter takes precedence + saved_user_agent = os.environ.pop("DATABRICKS_USER_AGENT", None) + saved_litellm_ua = os.environ.pop("LITELLM_USER_AGENT", None) + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" Custom User-Agent param: testpartner/2.0.0") + print(f" Expected User-Agent: testpartner_litellm/{version}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'test' only."}], + max_tokens=10, + user_agent="testpartner/2.0.0", # This should result in testpartner_litellm/{version} + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Custom user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is 'testpartner_litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Custom user-agent test failed: {e}") + return False + + finally: + # Restore environment variables + if saved_user_agent: + os.environ["DATABRICKS_USER_AGENT"] = saved_user_agent + if saved_litellm_ua: + os.environ["LITELLM_USER_AGENT"] = saved_litellm_ua + + +def test_chat_completion_with_env_user_agent(config: dict): + """Test chat completion with user agent set via environment variable.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with User-Agent from ENV VAR") + print("=" * 60) + + import litellm + + # Set a specific user agent via environment + test_partner = "envpartner" + os.environ["DATABRICKS_USER_AGENT"] = test_partner + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" DATABRICKS_USER_AGENT env var: {test_partner}") + print(f" Expected User-Agent: {test_partner}_litellm/{version}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'env' only."}], + max_tokens=10, + # Note: NOT passing user_agent parameter - should use env var + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Env var user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is '{test_partner}_litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Env var user-agent test failed: {e}") + return False + + finally: + # Clean up + os.environ.pop("DATABRICKS_USER_AGENT", None) + + +def test_embedding(config: dict): + """Test embeddings with Databricks.""" + print("\n" + "=" * 60) + print("TEST: Embeddings") + print("=" * 60) + + import litellm + + model = config.get("TEST_EMBEDDING_MODEL", "databricks-bge-large-en") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + + try: + response = litellm.embedding( + model=full_model, + input=["Hello, world!"], + ) + + # Handle both object and dict response formats + if hasattr(response, "data"): + data = response.data + else: + data = response.get("data", []) + + if data: + first_item = data[0] + if hasattr(first_item, "embedding"): + embedding = first_item.embedding + else: + embedding = first_item.get("embedding", []) + + print(f" Embedding dimensions: {len(embedding)}") + print(f" First 5 values: {embedding[:5]}") + print(" ✓ Embedding test passed!") + return True + else: + print(" ✗ Embedding test failed: No data in response") + return False + + except Exception as e: + print(f" ✗ Embedding test failed: {e}") + print(" (This is expected if embedding model is not available)") + return False + + +def test_oauth_token_retrieval(config: dict): + """Test OAuth M2M token retrieval.""" + print("\n" + "=" * 60) + print("TEST: OAuth M2M Token Retrieval") + print("=" * 60) + + if "DATABRICKS_CLIENT_ID" not in config or "DATABRICKS_CLIENT_SECRET" not in config: + print(" Skipped: OAuth credentials not configured") + return None + + from litellm.llms.databricks.common_utils import DatabricksBase + + try: + db = DatabricksBase() + token = db._get_oauth_m2m_token( + api_base=config["DATABRICKS_API_BASE"], + client_id=config["DATABRICKS_CLIENT_ID"], + client_secret=config["DATABRICKS_CLIENT_SECRET"], + ) + + # Redact token for display + redacted_token = ( + f"{token[:10]}...[REDACTED]" if len(token) > 10 else "[REDACTED]" + ) + print(f" Token obtained: {redacted_token}") + print(" ✓ OAuth M2M token retrieval passed!") + return True + + except Exception as e: + print(f" ✗ OAuth token retrieval failed: {e}") + return False + + +# ============================================================================== +# SDK INTEGRATION TESTS - Different ways of calling Databricks via LiteLLM +# ============================================================================== + + +def test_litellm_sdk_with_config_user_agent(config: dict): + """ + Test 1: LiteLLM SDK with custom user agent from config file. + + This test uses the LiteLLM SDK directly with the CUSTOM_USER_AGENT + specified in the databricks config file. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM SDK with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'LiteLLM SDK test' only."}], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, # Use config user agent + ) + + content = response.choices[0].message.content + print(f" Response: {content}") + print(" ✓ LiteLLM SDK with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM SDK test failed: {e}") + return False + + +def test_langchain_litellm_with_user_agent(config: dict): + """ + Test 2: LangChain with LiteLLM integration. + + This test uses LangChain's ChatLiteLLM wrapper to call Databricks + with custom user agent from config. + + Requires: pip install langchain-litellm (recommended) + or: pip install langchain langchain-community (deprecated) + """ + print("\n" + "=" * 60) + print("TEST: LangChain + LiteLLM with Config User-Agent") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + # Try the new langchain-litellm package first, fall back to deprecated import + ChatLiteLLM = None + HumanMessage = None + + try: + from langchain_litellm import ChatLiteLLM + from langchain_core.messages import HumanMessage + + print(" Using: langchain-litellm package (recommended)") + except ImportError: + try: + # Fall back to deprecated import + import warnings + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=DeprecationWarning) + from langchain_community.chat_models import ChatLiteLLM + from langchain_core.messages import HumanMessage + print( + " Using: langchain-community (deprecated, consider: pip install langchain-litellm)" + ) + except ImportError: + print(" Skipped: langchain-litellm not installed") + print(" Install with: pip install langchain-litellm") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + # Set user agent via environment for LangChain integration + os.environ["DATABRICKS_USER_AGENT"] = custom_ua + + chat = ChatLiteLLM( + model=full_model, + max_tokens=20, + temperature=0.1, + ) + + messages = [HumanMessage(content="Say 'LangChain test' only.")] + response = chat.invoke(messages) + + content = response.content + print(f" Response: {content}") + print(" ✓ LangChain + LiteLLM with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LangChain + LiteLLM test failed: {e}") + import traceback + + traceback.print_exc() + return False + + finally: + # Clean up env var + os.environ.pop("DATABRICKS_USER_AGENT", None) + + +def test_litellm_async_completion(config: dict): + """ + Test 3: LiteLLM Async Completion API with custom User-Agent. + + This test uses LiteLLM's async completion API (acompletion) to call + Databricks with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Async Completion with Config User-Agent") + print("=" * 60) + + import asyncio + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + async def run_async_completion(): + response = await litellm.acompletion( + model=full_model, + messages=[{"role": "user", "content": "Say 'LiteLLM async test' only."}], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, + ) + return response + + try: + response = asyncio.run(run_async_completion()) + + content = response.choices[0].message.content + print(f" Response: {content}") + print(" ✓ LiteLLM async completion with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM async completion test failed: {e}") + import traceback + + traceback.print_exc() + return False + + +def test_litellm_streaming_completion(config: dict): + """ + Test 4: LiteLLM Streaming Completion with custom User-Agent. + + This test uses LiteLLM's streaming completion API to call + Databricks with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Streaming Completion with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + # Use streaming completion + response = litellm.completion( + model=full_model, + messages=[ + {"role": "user", "content": "Say 'LiteLLM streaming test' only."} + ], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, + stream=True, + ) + + # Collect streamed content + collected_content = "" + for chunk in response: + if chunk.choices and chunk.choices[0].delta.content: + collected_content += chunk.choices[0].delta.content + + print(f" Response (streamed): {collected_content}") + print(" ✓ LiteLLM streaming completion with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM streaming completion test failed: {e}") + import traceback + + traceback.print_exc() + return False + + +def test_litellm_embedding_with_user_agent(config: dict): + """ + Test 5: LiteLLM Embedding API with custom User-Agent. + + This test uses LiteLLM's embedding API to call Databricks + with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Embedding with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_EMBEDDING_MODEL", "databricks-bge-large-en") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + response = litellm.embedding( + model=full_model, + input=["Hello, this is a LiteLLM embedding test with custom user agent!"], + user_agent=custom_ua, + ) + + # Handle both object and dict response formats + if hasattr(response, "data"): + data = response.data + else: + data = response.get("data", []) + + if data: + first_item = data[0] + if hasattr(first_item, "embedding"): + embedding = first_item.embedding + else: + embedding = first_item.get("embedding", []) + + print(f" Embedding dimensions: {len(embedding)}") + print(f" First 3 values: {embedding[:3]}") + print(" ✓ LiteLLM embedding with config user-agent test passed!") + return True + else: + print(" ✗ LiteLLM embedding test failed: No data in response") + return False + + except Exception as e: + print(f" ✗ LiteLLM embedding test failed: {e}") + print(" (This may fail if embedding model is not available)") + import traceback + + traceback.print_exc() + return False + + +def run_integration_tests_for_auth_method(config: dict, auth_method: str) -> list: + """Run integration tests for a specific auth method. Returns list of (name, result) tuples.""" + results = [] + + print("\n" + "=" * 60) + print(f"INTEGRATION TESTS - {auth_method.upper()} Authentication") + print("=" * 60) + + # Setup environment for this auth method + try: + setup_environment(config, auth_method) + except ValueError as e: + print(f" ✗ Setup failed: {e}") + return [(f"[{auth_method.upper()}] Setup", False)] + + # Test OAuth token retrieval (only for oauth method) + if auth_method == "oauth": + results.append( + ( + f"[{auth_method.upper()}] OAuth Token Retrieval", + test_oauth_token_retrieval(config), + ) + ) + + # Test chat completion + results.append( + (f"[{auth_method.upper()}] Chat Completion", test_chat_completion(config)) + ) + + # Test embeddings + results.append((f"[{auth_method.upper()}] Embeddings", test_embedding(config))) + + return results + + +def main(): + print("=" * 60) + print("DATABRICKS LITELLM INTEGRATION TESTS") + print("=" * 60) + + # Load config + print(f"\nLoading config from: {CONFIG_FILE}") + try: + config = load_config(CONFIG_FILE) + print(f" Loaded {len(config)} configuration values") + except FileNotFoundError as e: + print(f"\nERROR: {e}") + return 1 + + # Validate required config + if "DATABRICKS_API_BASE" not in config: + print("\nERROR: DATABRICKS_API_BASE is required in config file") + return 1 + + auth_method = config.get("TEST_AUTH_METHOD", "pat").lower() + print(f"\nTest Configuration:") + print(f" API Base: {config['DATABRICKS_API_BASE']}") + print(f" Auth Method: {auth_method}") + + # Run unit tests (no credentials needed) + print("\n" + "=" * 60) + print("UNIT TESTS (No credentials needed)") + print("=" * 60) + + test_user_agent_building() + test_token_redaction() + + all_results = [] + + # Determine which auth methods to test + if auth_method == "all": + auth_methods_to_test = ["oauth", "pat", "sdk"] + print("\n" + "#" * 60) + print("# TESTING ALL AUTHENTICATION METHODS") + print("#" * 60) + else: + auth_methods_to_test = [auth_method] + + # Run integration tests for each auth method + for method in auth_methods_to_test: + results = run_integration_tests_for_auth_method(config, method) + all_results.extend(results) + + # Run User-Agent tests (only once, using the last auth method or 'pat' for 'all') + print("\n" + "-" * 60) + print("USER-AGENT INTEGRATION TESTS") + print("-" * 60) + + # Setup environment for user-agent tests (use 'pat' as it's simplest) + if auth_method == "all": + setup_environment(config, "pat") + + # Test 1: Default user agent (no custom agent set) + all_results.append( + ( + "Chat with DEFAULT User-Agent", + test_chat_completion_default_user_agent(config), + ) + ) + + # Test 2: Custom user agent passed as parameter + all_results.append( + ( + "Chat with Custom User-Agent (param)", + test_chat_completion_with_custom_user_agent(config), + ) + ) + + # Test 3: User agent from environment variable + all_results.append( + ( + "Chat with User-Agent from ENV", + test_chat_completion_with_env_user_agent(config), + ) + ) + + # Run SDK Integration Tests with different calling methods + print("\n" + "#" * 60) + print("# SDK INTEGRATION TESTS - DIFFERENT CALLING METHODS") + print("# Using CUSTOM_USER_AGENT from config file") + print("#" * 60) + + # Setup environment for SDK tests (use 'pat' as it's most compatible) + setup_environment(config, "pat") + + # Test 1: LiteLLM SDK with config user agent + all_results.append( + ( + "LiteLLM SDK with Config User-Agent", + test_litellm_sdk_with_config_user_agent(config), + ) + ) + + # Test 2: LangChain + LiteLLM with config user agent + all_results.append( + ( + "LangChain + LiteLLM with Config User-Agent", + test_langchain_litellm_with_user_agent(config), + ) + ) + + # Test 3: LiteLLM Async Completion with config user agent + all_results.append( + ( + "LiteLLM Async Completion with Config User-Agent", + test_litellm_async_completion(config), + ) + ) + + # Test 4: LiteLLM Streaming Completion with config user agent + all_results.append( + ( + "LiteLLM Streaming Completion with Config User-Agent", + test_litellm_streaming_completion(config), + ) + ) + + # Test 5: LiteLLM Embedding with config user agent + all_results.append( + ( + "LiteLLM Embedding with Config User-Agent", + test_litellm_embedding_with_user_agent(config), + ) + ) + + # Summary + print("\n" + "=" * 60) + print("TEST SUMMARY") + print("=" * 60) + + passed = sum(1 for _, r in all_results if r is True) + failed = sum(1 for _, r in all_results if r is False) + skipped = sum(1 for _, r in all_results if r is None) + + for name, result in all_results: + status = ( + "✓ PASSED" + if result is True + else ("✗ FAILED" if result is False else "○ SKIPPED") + ) + print(f" {status}: {name}") + + print(f"\n Total: {passed} passed, {failed} failed, {skipped} skipped") + + if auth_method == "all": + print(f"\n Auth methods tested: {', '.join(auth_methods_to_test)}") + + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py new file mode 100644 index 00000000000..b4dc6c68bb0 --- /dev/null +++ b/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"