diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 0db4de3f43c..b1509313b88 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -195,9 +195,10 @@ impl AnthropicBeta { fn rejected_by(&self) -> &'static [BetaProvider] { match self { - Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => { - &[BetaProvider::BedrockConverse] + Self::AdvancedToolUse20251120 => { + &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } + Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], Self::AdvisorTool20260301 | Self::Compact20260904 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, @@ -231,12 +232,11 @@ impl AnthropicBeta { BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::VertexAi, + BetaProvider::Databricks, ], Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[], - Self::DangerousToolUse20260903 => { - &[BetaProvider::BedrockConverse, BetaProvider::Databricks] - } + Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::Effort20251124 => &[BetaProvider::VertexAi], Self::FastMode20260201 | Self::Oauth20250420 => &[ BetaProvider::AzureAi, @@ -244,6 +244,7 @@ impl AnthropicBeta { BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::VertexAi, + BetaProvider::Databricks, ], Self::FineGrainedToolStreaming20250514 => { &[BetaProvider::AzureAi, BetaProvider::VertexAi] diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index 3c18bdcf437..a0843621043 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -76,7 +76,7 @@ mod tests { #[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))] #[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))] #[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))] - #[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))] + #[case::databricks(BetaProvider::Databricks, None)] #[case::bedrock_converse(BetaProvider::BedrockConverse, None)] fn filter_renames_per_host( #[case] provider: BetaProvider, diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 7690372a214..236b18fe428 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -220,36 +220,36 @@ }, "databricks": { "advisor-tool-2026-03-01": null, - "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20", + "advanced-tool-use-2025-11-20": null, "bash_20241022": null, "bash_20250124": null, - "code-execution-2025-08-25": "code-execution-2025-08-25", + "code-execution-2025-08-25": null, "compact-2026-01-12": "compact-2026-01-12", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", - "dangerous-tool-use-2026-09-03": null, + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", - "fast-mode-2026-02-01": "fast-mode-2026-02-01", - "files-api-2025-04-14": "files-api-2025-04-14", + "fast-mode-2026-02-01": null, + "files-api-2025-04-14": null, "structured-output-2024-03-01": null, "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", - "mcp-client-2025-11-20": "mcp-client-2025-11-20", - "mcp-client-2025-04-04": "mcp-client-2025-04-04", + "mcp-client-2025-11-20": null, + "mcp-client-2025-04-04": null, "mcp-servers-2025-12-04": null, - "oauth-2025-04-20": "oauth-2025-04-20", + "oauth-2025-04-20": null, "output-128k-2025-02-19": "output-128k-2025-02-19", "per-turn-control-2026-07-01": null, - "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", - "skills-2025-10-02": "skills-2025-10-02", + "prompt-caching-scope-2026-01-05": null, + "skills-2025-10-02": null, "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", "text_editor_20241022": null, "text_editor_20250124": null, "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", - "web-fetch-2025-09-10": "web-fetch-2025-09-10", + "web-fetch-2025-09-10": null, "web-search-2025-03-05": "web-search-2025-03-05" } } diff --git a/litellm/litellm_core_utils/get_provider_specific_headers.py b/litellm/litellm_core_utils/get_provider_specific_headers.py index 20afd7535f3..bc325c3c702 100644 --- a/litellm/litellm_core_utils/get_provider_specific_headers.py +++ b/litellm/litellm_core_utils/get_provider_specific_headers.py @@ -9,6 +9,7 @@ class ProviderSpecificHeaderUtils: def get_provider_specific_headers( provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None, custom_llm_provider: str | None, + native_anthropic_messages: bool = False, ) -> Mapping[str, str]: """ Get the provider specific headers for the given custom llm provider. @@ -17,7 +18,8 @@ class ProviderSpecificHeaderUtils: carries its own comma-separated provider list, so headers that are safe for several providers and headers that are safe for exactly one can travel on the same request without sharing a scope. Entries whose provider list does not contain - `custom_llm_provider` contribute nothing. + `custom_llm_provider` contribute nothing, and entries marked + `native_anthropic_messages_only` contribute only to a provider's native Messages route. Returns: Dict: The provider specific headers for the given custom llm provider @@ -31,6 +33,12 @@ class ProviderSpecificHeaderUtils: matched_headers: Final = {} for scoped_header in scoped_headers: + if ( + "native_anthropic_messages_only" in scoped_header + and scoped_header["native_anthropic_messages_only"] + and not native_anthropic_messages + ): + continue stored_providers = scoped_header.get("custom_llm_provider", "") provider_list = [p.strip() for p in stored_providers.split(",")] if custom_llm_provider in provider_list: diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index 88cfeaad66a..002f39c121c 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -7,13 +7,14 @@ import asyncio import contextvars -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from functools import partial from typing import Final, cast from pydantic import TypeAdapter import litellm +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( @@ -39,8 +40,7 @@ from ..responses_adapters.handler import LiteLLMMessagesToResponsesAPIHandler from ..safeguards import ( SafeguardsEvaluator, native_route_classifies, - read_classifier_model_setting, - requested_dangerous_tool_use, + proxy_serves_safeguards, with_safeguard_results, ) from .interceptors import get_messages_interceptors @@ -106,11 +106,17 @@ def _should_route_to_responses_api( async def _with_native_safeguard_results( pending: Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]], - evaluator: SafeguardsEvaluator, + proxy_classifies: Callable[[], Awaitable[bool]], + build_evaluator: Callable[[], SafeguardsEvaluator | None], ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: from .streaming_iterator import AnthropicMessagesStreamingResponse, SafeguardedAnthropicSSEStream response: Final = await pending + if not await proxy_classifies(): + return response + evaluator: Final = await asyncio.to_thread(build_evaluator) + if evaluator is None: + return response if isinstance(response, AnthropicMessagesStreamingResponse): return AnthropicMessagesStreamingResponse( completion_stream=SafeguardedAnthropicSSEStream(response, evaluator), @@ -121,6 +127,10 @@ async def _with_native_safeguard_results( return SafeguardedAnthropicSSEStream(response, evaluator) +def _native_safeguards_evaluator(kwargs: Mapping[str, object], messages: object) -> SafeguardsEvaluator | None: + return request_safeguards_evaluator(kwargs, _SAFEGUARD_MESSAGES.validate_python(messages)) + + def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: """Whether the deployment opted into forwarding /v1/messages untranslated. @@ -139,7 +149,7 @@ def _deployment_supports_cache_control_ttl(model_info: object) -> bool: def _native_messages_config( - model: str, custom_llm_provider: str | None, model_info: object + model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None ) -> BaseAnthropicMessagesConfig | None: provider_config: Final = ( ProviderConfigManager.get_provider_anthropic_messages_config( @@ -149,8 +159,11 @@ def _native_messages_config( if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders] else None ) - if provider_config is not None or not _deployment_passes_through_anthropic_messages(model_info): - return provider_config + native_config: Final = ( + provider_config if provider_config is not None and provider_config.serves_api_base(model, api_base) else None + ) + if native_config is not None or not _deployment_passes_through_anthropic_messages(model_info): + return native_config from litellm.llms.openai_like.messages.transformation import ( OpenAILikeAnthropicMessagesConfig, ) @@ -158,12 +171,26 @@ def _native_messages_config( return OpenAILikeAnthropicMessagesConfig(cache_control_ttl=_deployment_supports_cache_control_ttl(model_info)) -def proxy_classifies_tool_uses(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool: +async def _proxy_classifies(model: str, safeguards: object, native_config: BaseAnthropicMessagesConfig | None) -> bool: + if not proxy_serves_safeguards(safeguards): + return False + if native_config is None: + return True + if native_config.should_filter_anthropic_beta_headers(): + await aload_beta_headers_config() + return not native_config.supports_safeguards() or native_route_classifies(model, safeguards) + + +async def proxy_classifies_tool_uses(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool: """Whether this proxy, rather than the upstream, answers ``safeguards`` for a request routed this way.""" - safeguards: Final = kwargs.get("safeguards") - if _native_messages_config(model, custom_llm_provider, kwargs.get("model_info")) is not None: - return native_route_classifies(model, safeguards) - return requested_dangerous_tool_use(safeguards) is not None and read_classifier_model_setting() is not None + api_base: Final = kwargs.get("api_base") + return await _proxy_classifies( + model, + kwargs.get("safeguards"), + _native_messages_config( + model, custom_llm_provider, kwargs.get("model_info"), api_base if isinstance(api_base, str) else None + ), + ) ####### ENVIRONMENT VARIABLES ################### @@ -608,7 +635,7 @@ def anthropic_messages_handler( ) anthropic_messages_provider_config: Final = _native_messages_config( - model, custom_llm_provider, kwargs.get("model_info") + model, custom_llm_provider, kwargs.get("model_info"), api_base ) if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. @@ -691,11 +718,6 @@ def anthropic_messages_handler( safeguard_kwargs: Final = ( _SAFEGUARD_KWARGS.validate_python(kwargs) if is_async is True and "safeguards" in kwargs else None ) - safeguards_evaluator: Final = ( - request_safeguards_evaluator(safeguard_kwargs, _SAFEGUARD_MESSAGES.validate_python(messages)) - if safeguard_kwargs is not None and native_route_classifies(model, safeguard_kwargs.get("safeguards")) - else None - ) native_response: Final = base_llm_http_handler.anthropic_messages_handler( model=model, messages=cast( # cast-ok: anthropic messages are dicts at runtime @@ -716,6 +738,10 @@ def anthropic_messages_handler( stream=stream, kwargs=kwargs, ) - if safeguards_evaluator is None or not isinstance(native_response, Coroutine): + if safeguard_kwargs is None or not isinstance(native_response, Coroutine): return native_response - return _with_native_safeguard_results(native_response, safeguards_evaluator) + return _with_native_safeguard_results( + native_response, + partial(_proxy_classifies, model, safeguard_kwargs.get("safeguards"), anthropic_messages_provider_config), + partial(_native_safeguards_evaluator, safeguard_kwargs, messages), + ) diff --git a/litellm/llms/anthropic/pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/pass_through/messages/mcp_handler.py index 9331c90aa5d..f369e10f22c 100644 --- a/litellm/llms/anthropic/pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/pass_through/messages/mcp_handler.py @@ -136,7 +136,7 @@ async def anthropic_messages_with_mcp( ) custom_llm_provider: Final = kwargs.get("custom_llm_provider") - classified_here: Final = proxy_classifies_tool_uses( + classified_here: Final = await proxy_classifies_tool_uses( model, custom_llm_provider if isinstance(custom_llm_provider, str) else None, kwargs ) loop_call_args: Final[Mapping[str, object]] = ( diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index ad45485fa56..30d45f3cec1 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -3,10 +3,12 @@ from typing import ClassVar, Final import httpx +import litellm from litellm._logging import verbose_logger -from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config +from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_beta_header_supported_by_bundled_config from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -110,6 +112,27 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def should_filter_anthropic_beta_headers(self) -> bool: return self._resolved_provider != "anthropic" + def supports_safeguards(self) -> bool: + """ + Whether the upstream platform answers the ``safeguards`` body field itself. A platform + whose beta header map column marks the dangerous-tool-use beta unsupported does not, and + a platform the map does not list is trusted to unless its config says otherwise. + """ + return not self.should_filter_anthropic_beta_headers() or ( + ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value + not in get_unsupported_headers(self._resolved_provider) + ) + + def drops_safeguards(self, safeguards: object, drop_params: bool) -> bool: + """ + ``safeguards`` and the dangerous-tool-use beta leave together for a platform that does not + answer them, when the proxy classifies in its place or drop_params is set. Otherwise both + are forwarded as sent and the platform answers its own 400, like any unsupported param. + """ + return not self.supports_safeguards() and ( + drop_params or litellm.drop_params is True or proxy_serves_safeguards(safeguards) + ) + def _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None: """ Remove `scope` field from cache_control blocks. @@ -208,7 +231,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ids were found without tool_result blocks immediately after") while consecutive user turns merge upstream. Billing-header system blocks are stripped from the top-level ``system`` - field regardless of whether anything was hoisted. + field when ``should_strip_billing_metadata`` says so, regardless of + whether anything was hoisted. Subclasses whose upstream rejects the role opt in by calling this from their ``transform_anthropic_messages_request``; the first-party Anthropic @@ -242,7 +266,11 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ) for block in as_system_content_blocks(source) ] - filtered_system: Final = self._filter_billing_headers_from_system(system_content) + filtered_system: Final = ( + self._filter_billing_headers_from_system(system_content) + if self.should_strip_billing_metadata() + else system_content + ) if filtered_system: anthropic_messages_request["system"] = filtered_system else: @@ -592,6 +620,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): status_code=400, ) + if "safeguards" in anthropic_messages_optional_request_params and self.drops_safeguards( + anthropic_messages_optional_request_params["safeguards"], # pyright: ignore[reportUnknownArgumentType] # the request params arrive as a bare dict + litellm_params.drop_params is True, + ): + anthropic_messages_optional_request_params.pop("safeguards") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + self._translate_reasoning_effort_to_anthropic( model=model, optional_params=anthropic_messages_optional_request_params, @@ -684,9 +718,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): self, model: str, httpx_response: httpx.Response, - request_body: dict, + request_body: dict[str, object], # mutable-ok: mirrors the base signature litellm_logging_obj: LiteLLMLoggingObj, - ) -> AsyncIterator: + ) -> AsyncIterator[bytes]: """Helper function to handle Anthropic streaming responses using the existing logging handlers""" from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, @@ -703,13 +737,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): litellm_logging_obj=litellm_logging_obj, ) - @staticmethod def _update_headers_with_anthropic_beta( + self, headers: dict, optional_params: dict, custom_llm_provider: str = "anthropic", messages: Sequence[object] = (), beta_map_provider: str | None = None, + drop_params: bool = False, ) -> dict: """ Auto-inject anthropic-beta headers based on features used. @@ -727,6 +762,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): custom_llm_provider: Provider name for looking up correct tool search header messages: Request messages, scanned for per-message output_config beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider + drop_params: Whether the deployment or request set drop_params """ beta_values: Final[set[str]] = set() @@ -805,9 +841,20 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): else () ) all_beta_values: Final = beta_values.union(thinking_display_betas, tool_change_betas) + dangerous_tool_use: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value + forwarded_beta_values: Final = ( # pyright: ignore[reportUnknownVariableType] # derived from the bare beta set + all_beta_values - {dangerous_tool_use} + if dangerous_tool_use in all_beta_values + and self.drops_safeguards( + optional_params.get("safeguards"), # pyright: ignore[reportUnknownArgumentType] # the optional params arrive as a bare dict + drop_params, + ) + else all_beta_values + ) if not all_beta_values: return headers merged: Final = {key: value for key, value in headers.items() if key.lower() != "anthropic-beta"} - merged["anthropic-beta"] = ",".join(sorted(all_beta_values)) + if forwarded_beta_values: + merged["anthropic-beta"] = ",".join(sorted(forwarded_beta_values)) return merged diff --git a/litellm/llms/anthropic/pass_through/safeguards.py b/litellm/llms/anthropic/pass_through/safeguards.py index 6eef313dc92..eededde7e92 100644 --- a/litellm/llms/anthropic/pass_through/safeguards.py +++ b/litellm/llms/anthropic/pass_through/safeguards.py @@ -273,12 +273,12 @@ def read_classifier_model_setting() -> str | None: return configured if isinstance(configured, str) and configured else None +def proxy_serves_safeguards(safeguards: object) -> bool: + return requested_dangerous_tool_use(safeguards) is not None and read_classifier_model_setting() is not None + + def native_route_classifies(model: str, safeguards: object) -> bool: - return ( - "claude" not in model.lower() - and requested_dangerous_tool_use(safeguards) is not None - and read_classifier_model_setting() is not None - ) + return "claude" not in model.lower() and proxy_serves_safeguards(safeguards) def build_safeguards_evaluator( diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index 06717207dc0..bda7015ab2b 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -138,6 +138,15 @@ class BaseAnthropicMessagesConfig(ABC): """ return True + def supports_safeguards(self) -> bool: + return True + + def serves_api_base(self, model: str, api_base: str | None) -> bool: + return True + + def is_native_messages_route(self) -> bool: + return True + def handles_web_search_natively(self) -> bool: """ Whether the upstream this config routes to executes ``web_search`` tools diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index df85c72ecdc..12421aa94fd 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -98,6 +98,9 @@ class AmazonAnthropicClaudeMessagesConfig( def beta_headers_provider(self) -> str: return self.custom_llm_provider or "bedrock" + def should_strip_billing_metadata(self) -> bool: + return True + def drops_dangerous_tool_use_beta_without_safeguards(self) -> bool: return True diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 1ad84761bc9..c71721d800e 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -176,9 +176,9 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): self, model: str, httpx_response: httpx.Response, - request_body: dict, + request_body: dict[str, object], # mutable-ok: mirrors the base signature litellm_logging_obj: LiteLLMLoggingObj, - ) -> AsyncIterator: + ) -> AsyncIterator[bytes]: return AnthropicMessagesConfig.get_async_streaming_response_iterator( self, model=model, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index bbda4f57ec8..b4dae7ba8b6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2314,6 +2314,7 @@ class BaseLLMHTTPHandler: provider_specific_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( provider_specific_header=provider_specific_header, custom_llm_provider=custom_llm_provider, + native_anthropic_messages=anthropic_messages_provider_config.is_native_messages_route(), ) forwarded_headers: Final = kwargs.get("headers", None) # Also check for extra_headers in kwargs (from config or direct calls) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 74bdd2cd733..2ecbb2d9c93 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -278,7 +278,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2 + use_ai_gateway: Final = self.uses_ai_gateway(model) api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base complete_url: Final = f"{url_base}/chat/completions" diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 5e85c91a6b3..3112b2458b4 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -12,12 +12,16 @@ Authentication priority: import os import re +from collections.abc import Mapping +from types import MappingProxyType from typing import Final, Literal from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict +import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str _DEFAULT_M2M_TOKEN_LIFETIME_SECONDS: Final = 3600 @@ -26,6 +30,28 @@ class DatabricksException(BaseLLMException): pass +DATABRICKS_ANTHROPIC_MESSAGES_PATH: Final = "/anthropic/v1/messages" +DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH: Final = f"/serving-endpoints{DATABRICKS_ANTHROPIC_MESSAGES_PATH}" +DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH: Final = "/ai-gateway/anthropic" +DATABRICKS_AI_GATEWAY_MESSAGES_PATH: Final = f"{DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH}/v1/messages" +DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = ( + "/serving-endpoints", + "/serving-endpoints/anthropic", + DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH, + "/ai-gateway/mlflow/v1", + DATABRICKS_AI_GATEWAY_ANTHROPIC_PATH, + DATABRICKS_AI_GATEWAY_MESSAGES_PATH, +) +DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-") +DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType( + { + "chat_completions": "/chat/completions", + "embeddings": "/embeddings", + "anthropic_messages": DATABRICKS_ANTHROPIC_MESSAGES_PATH, + } +) + + class OAuthM2MToken(BaseModel): model_config = ConfigDict(frozen=True) @@ -144,6 +170,47 @@ class DatabricksBase: return redacted + @staticmethod + def resolve_api_key(api_key: str | None) -> str | None: + return api_key or litellm.api_key or litellm.databricks_key or get_secret_str("DATABRICKS_API_KEY") + + @staticmethod + def resolve_api_base(api_base: str | None) -> str | None: + return api_base or litellm.api_base or os.getenv("DATABRICKS_API_BASE") + + @staticmethod + def serves_native_anthropic_messages(model: str) -> bool: + return model.removeprefix("databricks/").lower().startswith(DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES) + + @staticmethod + def uses_ai_gateway(model: str) -> bool: + return model.removeprefix("databricks/").count(".") >= 2 + + @staticmethod + def recognizes_workspace_route(api_base: str, model: str) -> bool: + route: Final = urlsplit(api_base).path.rstrip("/") + return route.endswith(DATABRICKS_WORKSPACE_ROUTE_PREFIXES) or ( + route == "" and DatabricksBase.uses_ai_gateway(model) + ) + + @staticmethod + def workspace_root(api_base: str) -> str: + normalized: Final = api_base.rstrip("/") + route_prefix: Final = next( + (prefix for prefix in DATABRICKS_WORKSPACE_ROUTE_PREFIXES if normalized.endswith(prefix)), "" + ) + return normalized.removesuffix(route_prefix) + + @staticmethod + def configured_user_agent(litellm_params: Mapping[str, object]) -> str | None: + configured: Final = ( + litellm_params.get("user_agent") + or litellm_params.get("databricks_user_agent") + or os.getenv("LITELLM_USER_AGENT") + or os.getenv("DATABRICKS_USER_AGENT") + ) + return configured if isinstance(configured, str) else None + @staticmethod def _build_user_agent(custom_user_agent: str | None = None) -> str: """ @@ -327,11 +394,11 @@ class DatabricksBase: self, api_key: str | None, api_base: str | None, - endpoint_type: Literal["chat_completions", "embeddings"], + endpoint_type: Literal["chat_completions", "embeddings", "anthropic_messages"], custom_endpoint: bool | None, headers: dict | None, custom_user_agent: str | None = None, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: # mutable-ok: becomes the outbound headers """ Validate and configure the Databricks environment. @@ -343,7 +410,7 @@ class DatabricksBase: Args: api_key: Personal access token (PAT) api_base: Databricks workspace URL with /serving-endpoints - endpoint_type: Type of endpoint (chat_completions or embeddings) + endpoint_type: Type of endpoint (chat_completions, embeddings or anthropic_messages) custom_endpoint: Whether using a custom endpoint URL headers: Existing headers dict custom_user_agent: Optional custom user agent to prefix @@ -410,9 +477,5 @@ class DatabricksBase: # Debug logging with redaction (never log actual tokens) verbose_logger.debug("Databricks request headers: %s", self.redact_headers_for_logging(headers)) - if endpoint_type == "chat_completions" and custom_endpoint is not True: - api_base = f"{api_base}/chat/completions" - elif endpoint_type == "embeddings" and custom_endpoint is not True: - api_base = f"{api_base}/embeddings" - - return api_base, headers + endpoint_path: Final = "" if custom_endpoint is True else DATABRICKS_ENDPOINT_PATHS[endpoint_type] + return f"{api_base}{endpoint_path}", headers diff --git a/litellm/llms/databricks/messages/transformation.py b/litellm/llms/databricks/messages/transformation.py new file mode 100644 index 00000000000..cb8d2b652a2 --- /dev/null +++ b/litellm/llms/databricks/messages/transformation.py @@ -0,0 +1,194 @@ +import json +from collections.abc import AsyncIterator, Mapping +from typing import Final + +import httpx + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) +from litellm.llms.databricks.common_utils import ( + DATABRICKS_AI_GATEWAY_MESSAGES_PATH, + DATABRICKS_ANTHROPIC_MESSAGES_PATH, + DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH, + DatabricksBase, +) +from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS +from litellm.types.router import GenericLiteLLMParams + +DATABRICKS_UNSUPPORTED_SERVER_TOOL_PREFIXES: Final = ("web_search_", "web_fetch_", "code_execution_") +DROP_CONTAINER_WARNING: Final = ( + "Dropping `container` for model=%s. The Databricks Anthropic Messages route rejects it, " + "and the chat completions bridge it replaces never forwarded it." +) +BEDROCK_INVOCATION_METRICS_FIELD: Final = "amazon-bedrock-invocationMetrics" + + +def _is_unsupported_server_tool(tool: object) -> bool: + tool_type: Final = (as_str_mapping(tool) or {}).get("type") + return isinstance(tool_type, str) and tool_type.startswith(DATABRICKS_UNSUPPORTED_SERVER_TOOL_PREFIXES) + + +def _tool_name(tool: object) -> object: + return (as_str_mapping(tool) or {}).get("name") + + +def _without_unsupported_server_tools(params: Mapping[str, object]) -> dict[str, object]: + if not isinstance(params.get("tools"), list): + return dict(params) + sent_tools: Final = tuple(iter_items(params.get("tools"))) + kept_tools: Final = [tool for tool in sent_tools if not _is_unsupported_server_tool(tool)] + dropped_names: Final = frozenset(_tool_name(tool) for tool in sent_tools if _is_unsupported_server_tool(tool)) + forced_tool: Final = _tool_name(params.get("tool_choice")) + keeps_tool_choice: Final = ( + "tool_choice" in params and bool(kept_tools) and (forced_tool is None or forced_tool not in dropped_names) + ) + remaining: Final = {key: value for key, value in params.items() if key not in ("tools", "tool_choice")} + return ( + remaining + | ({"tools": kept_tools} if kept_tools else {}) + | ({"tool_choice": params["tool_choice"]} if keeps_tool_choice else {}) + ) + + +def _without_invocation_metrics(sse_line: bytes) -> bytes: + if not sse_line.startswith(b"data:") or BEDROCK_INVOCATION_METRICS_FIELD.encode() not in sse_line: + return sse_line + event: Final = as_str_mapping(json.loads(sse_line.removeprefix(b"data:"))) + if event is None: + return sse_line + kept_fields: Final = {key: value for key, value in event.items() if key != BEDROCK_INVOCATION_METRICS_FIELD} + return b"data: " + json.dumps(kept_fields, separators=(",", ":")).encode() + + +async def _strip_invocation_metrics(sse_stream: AsyncIterator[bytes]) -> AsyncIterator[bytes]: + partial_line: Final = bytearray() + async for chunk in sse_stream: + partial_line.extend(chunk) + if complete_end := partial_line.rfind(b"\n") + 1: + yield b"\n".join( + _without_invocation_metrics(line) for line in bytes(partial_line[:complete_end]).split(b"\n") + ) + del partial_line[:complete_end] + if partial_line: + yield _without_invocation_metrics(bytes(partial_line)) + + +class DatabricksAnthropicMessagesConfig(DatabricksBase, AnthropicMessagesConfig): + """ + Forwards /v1/messages for Claude models to the workspace's native Anthropic + route (``/serving-endpoints/anthropic/v1/messages`` for a serving endpoint, + ``/ai-gateway/anthropic/v1/messages`` for a Unity Catalog model service), so + Anthropic-only request fields and response fields (``safeguards`` and + ``safeguard_results``) survive instead of being bridged through chat completions. + """ + + @property + def custom_llm_provider(self) -> str | None: + return "databricks" + + def handles_web_search_natively(self) -> bool: + return False + + def drops_safeguards(self, safeguards: object, drop_params: bool) -> bool: + return not self.supports_safeguards() + + def serves_api_base(self, model: str, api_base: str | None) -> bool: + resolved_api_base: Final = self.resolve_api_base(api_base) + return resolved_api_base is None or self.recognizes_workspace_route(resolved_api_base, model) + + def get_async_streaming_response_iterator( + self, + model: str, + httpx_response: httpx.Response, + request_body: dict[str, object], # mutable-ok: mirrors the base signature + litellm_logging_obj: LiteLLMLoggingObj, + ) -> AsyncIterator[bytes]: + return _strip_invocation_metrics( + super().get_async_streaming_response_iterator( + model=model, + httpx_response=httpx_response, + request_body=request_body, + litellm_logging_obj=litellm_logging_obj, + ) + ) + + def validate_anthropic_messages_environment( + self, + headers: dict[str, str], # mutable-ok: mirrors the base signature + model: str, + messages: list[object], # mutable-ok: mirrors the base signature + optional_params: dict[str, object], # mutable-ok: mirrors the base signature + litellm_params: dict[str, object], # mutable-ok: mirrors the base signature + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature + client_anthropic_headers: Final = { + key: value for key, value in headers.items() if key.lower() in ANTHROPIC_API_HEADERS + } + messages_api_base, auth_headers = self.databricks_validate_environment( # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict + api_key=self.resolve_api_key(api_key), + api_base=self.resolve_api_base(api_base), + endpoint_type="anthropic_messages", + custom_endpoint=False, + headers={key: value for key, value in headers.items() if key not in client_anthropic_headers}, + custom_user_agent=self.configured_user_agent(litellm_params), + ) + workspace_api_base: Final = messages_api_base.removesuffix(DATABRICKS_ANTHROPIC_MESSAGES_PATH) + normalized_headers: Final = { + ("content-type" if key.lower() == "content-type" else key): value + for key, value in {**auth_headers, **client_anthropic_headers}.items() + } + finalized_headers: Final[dict[str, str]] = self._finalize_messages_headers( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the shared helper returns a bare dict # mutable-ok: becomes the outbound headers + headers=normalized_headers, + optional_params=optional_params, + messages=messages, + ) + return finalized_headers, workspace_api_base + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], # mutable-ok: mirrors the base signature + litellm_params: dict[str, object], # mutable-ok: mirrors the base signature + stream: bool | None = None, + ) -> str: + messages_path: Final = ( + DATABRICKS_AI_GATEWAY_MESSAGES_PATH + if self.uses_ai_gateway(model) + else DATABRICKS_SERVING_ENDPOINTS_MESSAGES_PATH + ) + return f"{self.workspace_root(self._get_api_base(api_base))}{messages_path}" + + def transform_anthropic_messages_request( + self, + model: str, + messages: list[dict[str, object]], # mutable-ok: mirrors the base signature + anthropic_messages_optional_request_params: dict[str, object], # mutable-ok: mirrors the base signature + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: mirrors the base signature + ) -> dict[str, object]: # mutable-ok: mirrors the base signature + if "container" in anthropic_messages_optional_request_params: + litellm.verbose_logger.warning(DROP_CONTAINER_WARNING, model) + workspace_params: Final = { + key: value for key, value in anthropic_messages_optional_request_params.items() if key != "container" + } + drops_unsupported_params: Final = litellm.drop_params is True or litellm_params.drop_params is True + request_params: Final = ( + _without_unsupported_server_tools(workspace_params) if drops_unsupported_params else workspace_params + ) + anthropic_messages_request: Final[dict[str, object]] = super().transform_anthropic_messages_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the shared base returns a bare dict # mutable-ok: becomes the outbound request body + model=model, + messages=messages, + anthropic_messages_optional_request_params=request_params, + litellm_params=litellm_params, + headers=headers, + ) + self._normalize_system_role_messages(anthropic_messages_request, model=model) # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict + self._remove_scope_from_cache_control(anthropic_messages_request) # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict + return anthropic_messages_request diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 94f9792da35..ef501720e2d 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -33,6 +33,9 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): super().__init__() self._cache_control_ttl: Final = cache_control_ttl + def is_native_messages_route(self) -> bool: + return False + def validate_anthropic_messages_environment( self, headers: dict[str, str], @@ -59,6 +62,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): optional_params=optional_params, messages=messages, beta_map_provider=LlmProviders.OPENAI_LIKE.value, + drop_params=litellm_params.get("drop_params") is True, # pyright: ignore[reportUnknownMemberType] # the litellm params arrive as a bare dict ) return merged, api_base @@ -132,6 +136,9 @@ class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): def should_strip_billing_metadata(self) -> bool: return True + def supports_safeguards(self) -> bool: + return False + def _resolve_api_key(self, api_key: str | None) -> str | None: return api_key or get_secret_str(self._provider.api_key_env) or litellm.api_key diff --git a/litellm/main.py b/litellm/main.py index 2ac65b738a4..38591441e01 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -239,6 +239,7 @@ from .llms.cohere.embed import handler as cohere_embed from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router +from .llms.databricks.common_utils import DatabricksBase from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha from .llms.gdc.chat.transformation import GDCGeminiConfig @@ -3361,8 +3362,8 @@ def _complete_oobabooga(ctx: CompletionDispatchContext) -> _CompletionDispatchRe def _complete_databricks(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion - api_base = ctx.api_base - api_key = ctx.api_key + api_base: Final = DatabricksBase.resolve_api_base(ctx.api_base) + api_key: Final = DatabricksBase.resolve_api_key(ctx.api_key) client: Final = _dispatch_client_http(ctx) headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3374,20 +3375,6 @@ def _complete_databricks(ctx: CompletionDispatchContext) -> _CompletionDispatchR stream: Final = ctx.stream timeout: Final = ctx.timeout - api_base = ( - api_base # for databricks we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or os.getenv("DATABRICKS_API_BASE") - ) - - # set API KEY - api_key = ( - api_key - or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there - or litellm.databricks_key - or get_secret("DATABRICKS_API_KEY") - ) - headers = headers or litellm.headers ## COMPLETION CALL @@ -6832,7 +6819,7 @@ def embedding( api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE") # set API KEY - api_key = api_key or litellm.api_key or litellm.databricks_key or get_secret("DATABRICKS_API_KEY") + api_key = DatabricksBase.resolve_api_key(api_key) ## EMBEDDING CALL response = databricks_embedding.embedding( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 29fe636f5e4..ce2e3cd30c4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4089,13 +4089,25 @@ def add_provider_specific_headers_to_request( } scoped_headers: Final = [ - ProviderSpecificHeader(custom_llm_provider=providers, extra_headers=extra_headers) - for providers, extra_headers in ( - (_ANTHROPIC_API_HEADER_PROVIDERS, anthropic_api_headers), - (_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS, anthropic_oauth_credential_headers), - (LlmProviders.TYPESAFE.value, typesafe_danger_zone_headers), + scoped_header + for scoped_header in ( + ProviderSpecificHeader( + custom_llm_provider=_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers + ), + ProviderSpecificHeader( + custom_llm_provider=LlmProviders.DATABRICKS.value, + extra_headers=anthropic_api_headers, + native_anthropic_messages_only=True, + ), + ProviderSpecificHeader( + custom_llm_provider=_ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS, + extra_headers=anthropic_oauth_credential_headers, + ), + ProviderSpecificHeader( + custom_llm_provider=LlmProviders.TYPESAFE.value, extra_headers=typesafe_danger_zone_headers + ), ) - if extra_headers + if scoped_header["extra_headers"] ] if scoped_headers: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8b2b319d1e0..419c2624111 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4538,6 +4538,7 @@ class HttpHandlerRequestFields(TypedDict, total=False): class ProviderSpecificHeader(TypedDict): custom_llm_provider: str extra_headers: dict + native_anthropic_messages_only: NotRequired[ReadOnly[bool]] class SelectTokenizerResponse(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 142a1745143..ab97b69897f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9367,6 +9367,15 @@ class ProviderConfigManager: ) return AzureAnthropicMessagesConfig() + elif litellm.LlmProviders.DATABRICKS == provider: + from litellm.llms.databricks.common_utils import DatabricksBase + + if DatabricksBase.serves_native_anthropic_messages(model): + from litellm.llms.databricks.messages.transformation import ( + DatabricksAnthropicMessagesConfig, + ) + + return DatabricksAnthropicMessagesConfig() elif litellm.LlmProviders.MINIMAX == provider: from litellm.llms.minimax.messages.transformation import ( MinimaxMessagesConfig, diff --git a/tests/integration/_support/safeguards_rig.py b/tests/integration/_support/safeguards_rig.py index 677a41e274b..8ba6f14f71c 100644 --- a/tests/integration/_support/safeguards_rig.py +++ b/tests/integration/_support/safeguards_rig.py @@ -287,6 +287,7 @@ def safeguards_rig( directory: Path, *, environment: Callable[[str], Mapping[str, str]] = no_environment, + remove_environment: tuple[str, ...] = (), workers: int = 2, classifier_model: JsonValue | Omitted = CLASSIFIER_MODEL, general_settings: Mapping[str, JsonValue] | None = None, @@ -305,7 +306,14 @@ def safeguards_rig( litellm_settings=litellm_settings, top_level=top_level, ) - with owned_proxy_process(shared, directory, dict(environment(wire.url)), config=path, workers=workers) as owned: + with owned_proxy_process( + shared, + directory, + dict(environment(wire.url)), + config=path, + remove_environment=remove_environment, + workers=workers, + ) as owned: yield Rig(wire, scripts, owned) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py index 0c41efffe73..cabde4f1e92 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -497,6 +497,7 @@ class _ProviderCase: keys: frozenset[str] = _MESSAGES_KEYS extra_headers: Mapping[str, str] = MappingProxyType({}) body_beta: list[str] | None = None + forwards_safeguards: bool = True _PROVIDERS: Final = MappingProxyType( @@ -533,11 +534,24 @@ _PROVIDERS: Final = MappingProxyType( auth=("authorization", "Bearer synthetic-sail-key"), expected_beta=None, ), + "D18_sail_with_claude_code_header_and_drop_params": _ProviderCase( + parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key", "drop_params": True}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("authorization", "Bearer synthetic-sail-key"), + expected_beta=None, + forwards_safeguards=False, + ), "D19_edenai": _ProviderCase( parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"}, auth=("authorization", "Bearer synthetic-edenai-key"), expected_beta=None, ), + "D19_edenai_with_drop_params": _ProviderCase( + parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key", "drop_params": True}, + auth=("authorization", "Bearer synthetic-edenai-key"), + expected_beta=None, + forwards_safeguards=False, + ), "D21_bedrock_claude_platform": _ProviderCase( parameters={ "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", @@ -601,7 +615,14 @@ def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_de response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta) _assert_answered(response, marker) observed: Final = _only(wire) - _expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys) + _expect_messages( + observed, + target=case.target, + auth=case.auth, + beta=case.expected_beta, + safeguards=SAFEGUARDS if case.forwards_safeguards else OMITTED, + keys=case.keys, + ) for name, value in case.extra_headers.items(): assert observed.headers.get(name) == value, observed.header_names assert observed.body.get("anthropic_beta") == case.body_beta, observed.body diff --git a/tests/integration/messages_endpoint/providers/databricks/test_databricks_native_messages_wire.py b/tests/integration/messages_endpoint/providers/databricks/test_databricks_native_messages_wire.py new file mode 100644 index 00000000000..95f8698618b --- /dev/null +++ b/tests/integration/messages_endpoint/providers/databricks/test_databricks_native_messages_wire.py @@ -0,0 +1,598 @@ +"""/v1/messages on Databricks Claude names reaches the workspace's native Anthropic route with the caller's +``safeguards``, while every other Databricks name and api_base keeps the merge base's chat completions bridge or +opted-in passthrough.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterator, Mapping, Sequence +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import list_value, object_value, string_value +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + CLAUDE_MODEL, + CLAUDE_TARGET, + DANGEROUS_TOOL_USE_BETA, + NATIVE_MODEL, + NATIVE_TARGET, + OMITTED, + OTHER_NATIVE_MODELS, + SAFEGUARDS, + Rig, + auto_mode_body, + chat_backend, + deployment, + flagged, + json_body, + messages_body, + native_backend, + native_message, + native_stream_frames, + nonce, + request_json, + safeguards_rig, + sse, + sse_events, + stream_results, + text_block, + tool_id, + tool_use_block, +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +DATABRICKS_MODEL: Final = "safeguards-databricks" +UNITY_CATALOG_MODEL: Final = "safeguards-databricks-unity-catalog" +UNITY_CATALOG_ROOT_MODEL: Final = "safeguards-databricks-unity-catalog-root" +SDK_MODEL: Final = "safeguards-databricks-sdk" +DROP_PARAMS_MODEL: Final = "safeguards-databricks-drop-params" +ENVIRONMENT_MODEL: Final = "safeguards-databricks-environment" +GPT_MODEL: Final = "safeguards-databricks-gpt" +CUSTOM_CLAUDE_MODEL: Final = "safeguards-databricks-custom-claude" +GATEWAY_MODEL: Final = "safeguards-databricks-behind-gateway" +OPTED_IN_GATEWAY_MODEL: Final = "safeguards-databricks-opted-in-gateway" +AZURE_MODEL: Final = "safeguards-azure-claude" +EDENAI_MODEL: Final = "safeguards-edenai" + +SERVING_TARGET: Final = "/dbx/serving-endpoints/anthropic/v1/messages" +UNITY_CATALOG_TARGET: Final = "/dbxuc/ai-gateway/anthropic/v1/messages" +UNITY_CATALOG_ROOT_TARGET: Final = "/ai-gateway/anthropic/v1/messages" +SDK_TARGET: Final = "/dbxsdk/serving-endpoints/anthropic/v1/messages" +ENVIRONMENT_TARGET: Final = "/dbxenv/serving-endpoints/anthropic/v1/messages" +SERVING_CHAT_TARGET: Final = "/dbx/serving-endpoints/chat/completions" +GATEWAY_CHAT_TARGET: Final = "/gw/dbx/chat/completions" +OPTED_IN_GATEWAY_TARGET: Final = "/gw/optin/v1/messages" +AZURE_TARGET: Final = "/azure/anthropic/v1/messages" +DATABRICKS_KEY: Final = "synthetic-databricks-key" +ENVIRONMENT_KEY: Final = "synthetic-databricks-environment-key" +SDK_TOKEN: Final = "synthetic-databricks-sdk-token" +PARTNER_USER_AGENT: Final = "partnerapp/2.0" +INVOCATION_METRICS: Final = "amazon-bedrock-invocationMetrics" +BILLING_BLOCK: Final[Mapping[str, JsonValue]] = { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.283.00; cc_entrypoint=sdk-cli;", +} + + +def environment(wire_url: str) -> Mapping[str, str]: + return { + "DATABRICKS_API_BASE": f"{wire_url}/dbxenv/serving-endpoints", + "DATABRICKS_API_KEY": ENVIRONMENT_KEY, + "DATABRICKS_USER_AGENT": PARTNER_USER_AGENT, + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "True", + } + + +def sdk_environment(wire_url: str) -> Mapping[str, str]: + return { + "DATABRICKS_HOST": f"{wire_url}/dbxsdk", + "DATABRICKS_TOKEN": SDK_TOKEN, + "DATABRICKS_AUTH_TYPE": "pat", + "DATABRICKS_CONFIG_FILE": "/dev/null", + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "True", + } + + +def databricks(name: str, model: str, api_base: str, **extra: JsonValue) -> Mapping[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": f"databricks/{model}", "api_base": api_base, "api_key": DATABRICKS_KEY, **extra}, + } + + +def extra_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + databricks(DATABRICKS_MODEL, "databricks-claude-sonnet-4-6", f"{wire_url}/dbx/serving-endpoints"), + databricks( + UNITY_CATALOG_MODEL, + "system.ai.claude-sonnet-4-6", + f"{wire_url}/dbxuc/serving-endpoints/anthropic/v1/messages", + ), + databricks(UNITY_CATALOG_ROOT_MODEL, "system.ai.claude-sonnet-4-6", wire_url), + databricks( + DROP_PARAMS_MODEL, "databricks-claude-opus-4-6", f"{wire_url}/dbx/serving-endpoints", drop_params=True + ), + {"model_name": ENVIRONMENT_MODEL, "litellm_params": {"model": "databricks/databricks-claude-haiku-4-5"}}, + databricks(GPT_MODEL, "databricks-gpt-oss-120b", f"{wire_url}/dbx/serving-endpoints"), + databricks(CUSTOM_CLAUDE_MODEL, "my-claude-endpoint", f"{wire_url}/dbx/serving-endpoints"), + databricks(GATEWAY_MODEL, "databricks-claude-sonnet-4-6", f"{wire_url}/gw/dbx"), + { + **databricks(OPTED_IN_GATEWAY_MODEL, "databricks-claude-sonnet-4-6", f"{wire_url}/gw/optin"), + "model_info": {"supported_endpoints": ["/v1/messages"]}, + }, + deployment(AZURE_MODEL, "azure_ai/claude-sonnet-4-6", f"{wire_url}/azure"), + ) + + +def sdk_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + { + "model_name": SDK_MODEL, + "litellm_params": { + "model": "databricks/databricks-claude-haiku-4-5", + "api_base": f"{wire_url}/dbxsdk/serving-endpoints", + }, + }, + ) + + +@pytest.fixture(scope="module") +def databricks_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with safeguards_rig( + tmp_path_factory.mktemp("databricks-native-messages"), environment=environment, extra_models=extra_models + ) as rig: + yield rig + + +@pytest.fixture(scope="module") +def sdk_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with safeguards_rig( + tmp_path_factory.mktemp("databricks-sdk-auth"), + environment=sdk_environment, + remove_environment=( + "DATABRICKS_API_KEY", + "DATABRICKS_API_BASE", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + ), + extra_models=sdk_models, + ) as rig: + yield rig + + +def beta_map_without_the_databricks_safeguards_beta() -> bytes: + beta_map: Final = object_value(json.loads(Path("litellm/anthropic_beta_headers_config.json").read_text())) + databricks_column: Final = {**object_value(beta_map["databricks"]), DANGEROUS_TOOL_USE_BETA: None} + return json.dumps({**beta_map, "databricks": databricks_column}).encode() + + +def serving_endpoint_model(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return (databricks(DATABRICKS_MODEL, "databricks-claude-sonnet-4-6", f"{wire_url}/dbx/serving-endpoints"),) + + +@pytest.fixture(scope="module") +def remote_beta_map_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Rig, Wire]]: + served_map: Final = beta_map_without_the_databricks_safeguards_beta() + with ( + wire_server(lambda _: Reply(body=served_map)) as beta_map, + safeguards_rig( + tmp_path_factory.mktemp("databricks-remote-beta-map"), + environment=lambda _: {"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map.url}/beta.json"}, + remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",), + classifier_model=OMITTED, + extra_models=serving_endpoint_model, + ) as rig, + ): + yield rig, beta_map + + +def scripted_message( + identity: str, model: str = "databricks-claude-sonnet-4-6", **extra: JsonValue +) -> Mapping[str, JsonValue]: + return native_message(identity, content=(text_block("Deleting."), tool_use_block(identity)), model=model, **extra) + + +def workspace_verdict(identity: str) -> Sequence[JsonValue]: + return [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {identity: {"type": "evaluated", "outcome": "allowed", "explanation": "workspace"}}, + }, + } + ] + + +def only_backend_call(rig: Rig, request_nonce: str) -> Request: + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(received) == 1, [request.target for request in received] + return received[0] + + +def beta_values(request: Request) -> frozenset[str]: + return frozenset(value.strip() for value in request.headers.get("anthropic-beta", "").split(",") if value.strip()) + + +def tool_types_and_names(body: Mapping[str, JsonValue]) -> tuple[tuple[JsonValue, JsonValue], ...]: + tools: Final = tuple(object_value(tool) for tool in list_value(body.get("tools", []))) + return tuple((tool.get("type"), tool.get("name")) for tool in tools) + + +def frame_data(frame: bytes) -> dict[str, JsonValue]: + line: Final = next(line for line in frame.splitlines() if line.startswith(b"data: ")) + return object_value(json.loads(line.removeprefix(b"data: "))) + + +def test_claude_serving_endpoint_reaches_the_native_route_and_keeps_the_workspace_verdict( + databricks_rig: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + message: Final = scripted_message(identity, safeguard_results=list(workspace_verdict(identity))) + databricks_rig.register(request_nonce, native_backend(message, stream=False)) + request: Final = auto_mode_body(DATABRICKS_MODEL, request_nonce) + body: Final = json_body(databricks_rig.post(request, beta=None)) + assert body == {**message, "model": DATABRICKS_MODEL}, body + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert sent.target == SERVING_TARGET, sent.target + assert sent.headers.get("authorization") == f"Bearer {DATABRICKS_KEY}", sent.headers + assert beta_values(sent) == {DANGEROUS_TOOL_USE_BETA}, sent.headers + outbound: Final = request_json(sent) + assert outbound["safeguards"] == request["safeguards"], outbound + assert outbound["messages"] == request["messages"] and outbound["tools"] == request["tools"], outbound + + +def odd_chunks(frames: Sequence[bytes]) -> tuple[bytes, ...]: + return tuple(bytes((byte,)) for byte in b"".join(frames)) + + +def whole_frames(frames: Sequence[bytes]) -> tuple[bytes, ...]: + return tuple(frames) + + +@pytest.mark.parametrize("reshape", (whole_frames, odd_chunks), ids=("whole_frames", "one_byte_chunks")) +def test_claude_stream_drops_invocation_metrics_and_keeps_every_other_frame( + databricks_rig: Rig, reshape: Callable[[Sequence[bytes]], tuple[bytes, ...]] +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + verdict: Final = workspace_verdict(identity) + metrics: Final[Mapping[str, JsonValue]] = {"inputTokenCount": 40, "outputTokenCount": 9, "invocationLatency": 812} + sent_frames: Final = ( + *native_stream_frames(scripted_message(identity), delta_extra={"safeguard_results": list(verdict)})[:-1], + sse({"type": "message_stop", INVOCATION_METRICS: dict(metrics)}), + ) + databricks_rig.register( + request_nonce, native_backend(scripted_message(identity), stream=True, chunks=reshape(sent_frames)) + ) + status, raw = databricks_rig.stream(auto_mode_body(DATABRICKS_MODEL, request_nonce)) + assert status == 200, raw + received_events: Final = sse_events(raw) + sent_events: Final = tuple(frame_data(frame) for frame in sent_frames) + assert len(received_events) == len(sent_events), received_events + assert received_events[-1] == {"type": "message_stop"}, received_events[-1] + assert received_events[1:-1] == sent_events[1:-1], received_events + opening: Final = object_value(received_events[0]["message"]) + assert opening == {**object_value(sent_events[0]["message"]), "model": DATABRICKS_MODEL}, opening + assert stream_results(raw) == verdict, raw + assert only_backend_call(databricks_rig, request_nonce).target == SERVING_TARGET + + +@pytest.mark.parametrize( + ("model", "target"), + ((UNITY_CATALOG_MODEL, UNITY_CATALOG_TARGET), (UNITY_CATALOG_ROOT_MODEL, UNITY_CATALOG_ROOT_TARGET)), + ids=("full_messages_url", "workspace_root"), +) +def test_unity_catalog_claude_name_reaches_the_ai_gateway_route(databricks_rig: Rig, model: str, target: str) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + message: Final = scripted_message(identity, safeguard_results=list(workspace_verdict(identity))) + databricks_rig.register(request_nonce, native_backend(message, stream=False)) + request: Final = auto_mode_body(model, request_nonce) + body: Final = json_body(databricks_rig.post(request)) + assert body["safeguard_results"] == workspace_verdict(identity), body + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert sent.target == target, sent.target + assert request_json(sent)["safeguards"] == request["safeguards"], sent.body + + +def test_environment_only_deployment_reaches_the_native_route_with_the_environment_key_and_partner_agent( + databricks_rig: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register(request_nonce, native_backend(scripted_message(identity), stream=False)) + json_body(databricks_rig.post(auto_mode_body(ENVIRONMENT_MODEL, request_nonce))) + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert sent.target == ENVIRONMENT_TARGET, sent.target + assert sent.headers.get("authorization") == f"Bearer {ENVIRONMENT_KEY}", sent.headers + assert sent.headers.get("user-agent", "").startswith("partnerapp_litellm/"), sent.headers + + +def test_client_anthropic_headers_leave_sdk_auth_in_charge_of_the_native_route(sdk_rig: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + sdk_rig.register( + request_nonce, native_backend(scripted_message(identity, model="databricks-claude-haiku-4-5"), stream=False) + ) + json_body( + sdk_rig.post( + auto_mode_body(SDK_MODEL, request_nonce), + beta="context-1m-2025-08-07,fast-mode-2026-02-01", + headers={"anthropic-version": "2023-06-01"}, + ) + ) + sent: Final = only_backend_call(sdk_rig, request_nonce) + assert sent.target == SDK_TARGET, sent.target + assert sent.headers.get("authorization") == f"Bearer {SDK_TOKEN}", sent.headers + assert beta_values(sent) == {"context-1m-2025-08-07", DANGEROUS_TOOL_USE_BETA}, sent.headers + assert sent.headers.get("anthropic-version") == "2023-06-01", sent.headers + + +@pytest.mark.parametrize( + ("model", "target"), + ( + (GPT_MODEL, SERVING_CHAT_TARGET), + (CUSTOM_CLAUDE_MODEL, SERVING_CHAT_TARGET), + (GATEWAY_MODEL, GATEWAY_CHAT_TARGET), + ), + ids=("gpt_oss", "custom_endpoint_named_claude", "api_base_outside_workspace_routes"), +) +def test_other_names_and_api_bases_keep_the_chat_completions_bridge( + databricks_rig: Rig, model: str, target: str +) -> None: + request_nonce: Final = nonce() + databricks_rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + body: Final = json_body( + databricks_rig.post( + auto_mode_body(model, request_nonce), + beta=f"context-1m-2025-08-07,{DANGEROUS_TOOL_USE_BETA}", + headers={"anthropic-version": "2023-06-01"}, + ) + ) + assert body["type"] == "message", body + received: Final = databricks_rig.received(request_nonce) + backend_calls: Final = databricks_rig.backend_calls(received) + assert [call.target for call in backend_calls] == [target], received + assert "safeguards" not in request_json(backend_calls[0]), backend_calls[0].body + assert "anthropic-beta" not in backend_calls[0].headers, backend_calls[0].headers + assert "anthropic-version" not in backend_calls[0].headers, backend_calls[0].headers + + +def gateway_rejecting_the_prompt_caching_scope_beta(message: Mapping[str, JsonValue]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if "prompt-caching-scope-2026-01-05" not in beta_values(request): + return Reply(body=json.dumps(message).encode()) + rejection: Final = "Unexpected value(s) `prompt-caching-scope-2026-01-05` for the `anthropic-beta` header." + return Reply(status=400, body=json.dumps({"error_code": "BAD_REQUEST", "message": rejection}).encode()) + + return respond + + +def test_opted_in_passthrough_api_base_gets_no_client_anthropic_beta(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + message: Final = native_message( + tool_id(request_nonce), + content=(text_block("Done."),), + stop_reason="end_turn", + model="databricks-claude-sonnet-4-6", + ) + databricks_rig.register(request_nonce, gateway_rejecting_the_prompt_caching_scope_beta(message)) + body: Final = json_body( + databricks_rig.post( + messages_body(OPTED_IN_GATEWAY_MODEL, request_nonce), + beta="prompt-caching-scope-2026-01-05,context-1m-2025-08-07", + headers={"anthropic-version": "2023-06-01"}, + ) + ) + assert body["id"] == message["id"] and body["content"] == message["content"], body + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert sent.target == OPTED_IN_GATEWAY_TARGET, sent.target + assert "anthropic-beta" not in sent.headers, sent.headers + + +def test_chat_completions_on_a_databricks_model_carries_no_client_anthropic_headers(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + databricks_rig.register( + request_nonce, chat_backend(tool_id(request_nonce), stream=False, text="ok", finish_reason="stop") + ) + response: Final = databricks_rig.gateway.client.post( + "/v1/chat/completions", + json={"model": GPT_MODEL, "messages": [{"role": "user", "content": f"say ok {request_nonce}"}]}, + headers={ + "Authorization": f"Bearer {databricks_rig.gateway.key}", + "anthropic-beta": f"context-1m-2025-08-07,{DANGEROUS_TOOL_USE_BETA}", + "anthropic-version": "2023-06-01", + }, + timeout=120, + ) + assert response.status_code == 200, response.text + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert sent.target == SERVING_CHAT_TARGET, sent.target + assert "anthropic-beta" not in sent.headers, sent.headers + assert "anthropic-version" not in sent.headers, sent.headers + + +def test_client_betas_reach_the_workspace_filtered_by_its_beta_map(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register(request_nonce, native_backend(scripted_message(identity), stream=False)) + client_betas: Final = ( + "context-1m-2025-08-07", + "prompt-caching-scope-2026-01-05", + "fast-mode-2026-02-01", + DANGEROUS_TOOL_USE_BETA, + ) + json_body(databricks_rig.post(auto_mode_body(DATABRICKS_MODEL, request_nonce), beta=",".join(client_betas))) + sent: Final = only_backend_call(databricks_rig, request_nonce) + assert beta_values(sent) == {"context-1m-2025-08-07", DANGEROUS_TOOL_USE_BETA}, sent.headers + + +def test_native_route_drops_container_hoists_system_turns_and_keeps_the_billing_block(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register(request_nonce, native_backend(scripted_message(identity), stream=False)) + scoped_system: Final[Mapping[str, JsonValue]] = { + "type": "text", + "text": "Synthetic agent instructions.", + "cache_control": {"type": "ephemeral", "scope": "global"}, + } + request: Final = { + **auto_mode_body(DATABRICKS_MODEL, request_nonce), + "system": [dict(BILLING_BLOCK), dict(scoped_system)], + "messages": [ + {"role": "system", "content": "leading operator note"}, + {"role": "user", "content": f"clean the build directory {request_nonce}"}, + ], + "container": "container_synthetic_0001", + } + json_body(databricks_rig.post(request)) + outbound: Final = request_json(only_backend_call(databricks_rig, request_nonce)) + assert "container" not in outbound, outbound + system: Final = tuple(object_value(block) for block in list_value(outbound["system"])) + assert system[0] == BILLING_BLOCK, system + assert {**scoped_system, "cache_control": {"type": "ephemeral"}} in system, system + assert "leading operator note" in json.dumps(system), system + roles: Final = tuple(string_value(object_value(message)["role"]) for message in list_value(outbound["messages"])) + assert roles == ("user",), outbound["messages"] + + +SERVER_TOOLS: Final[Sequence[JsonValue]] = ( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 1}, + {"type": "web_fetch_20250910", "name": "web_fetch"}, + {"type": "code_execution_20250825", "name": "code_execution"}, +) +BASH_SERVER_TOOL: Final[Mapping[str, JsonValue]] = {"type": "bash_20250124", "name": "bash"} + + +@pytest.mark.parametrize( + ("tool_choice", "kept_tool_choice"), + ( + ({"type": "tool", "name": "web_search"}, None), + ({"type": "tool", "name": "Bash"}, {"type": "tool", "name": "Bash"}), + ({"type": "auto"}, {"type": "auto"}), + ), + ids=("forced_dropped_tool", "forced_kept_tool", "auto"), +) +def test_drop_params_drops_only_the_server_tools_the_workspace_rejects( + databricks_rig: Rig, tool_choice: Mapping[str, JsonValue], kept_tool_choice: Mapping[str, JsonValue] | None +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register(request_nonce, native_backend(scripted_message(identity), stream=False)) + request: Final = auto_mode_body(DROP_PARAMS_MODEL, request_nonce) + tools: Final = [*list_value(request["tools"]), *SERVER_TOOLS, dict(BASH_SERVER_TOOL)] + json_body(databricks_rig.post({**request, "tools": tools, "tool_choice": dict(tool_choice)})) + outbound: Final = request_json(only_backend_call(databricks_rig, request_nonce)) + assert tool_types_and_names(outbound) == ((None, "Bash"), ("bash_20250124", "bash")), outbound["tools"] + assert outbound.get("tool_choice") == kept_tool_choice, outbound + + +def test_drop_params_with_only_rejected_server_tools_sends_no_tools_and_no_tool_choice(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register( + request_nonce, + native_backend(native_message(identity, content=(text_block("Done."),), stop_reason="end_turn"), stream=False), + ) + request: Final = auto_mode_body(DROP_PARAMS_MODEL, request_nonce) + json_body(databricks_rig.post({**request, "tools": list(SERVER_TOOLS), "tool_choice": {"type": "any"}})) + outbound: Final = request_json(only_backend_call(databricks_rig, request_nonce)) + assert "tools" not in outbound and "tool_choice" not in outbound, outbound + + +def workspace_rejecting_server_tools(request: Request) -> Reply: + message: Final = "tool type 'web_search_20250305' is not supported for this model" + return Reply(status=400, body=json.dumps({"error_code": "BAD_REQUEST", "message": message}).encode()) + + +def test_without_drop_params_server_tools_reach_the_workspace_as_sent(databricks_rig: Rig) -> None: + request_nonce: Final = nonce() + databricks_rig.register(request_nonce, workspace_rejecting_server_tools) + request: Final = auto_mode_body(DATABRICKS_MODEL, request_nonce) + tools: Final = [*list_value(request["tools"]), *SERVER_TOOLS] + response: Final = databricks_rig.post({**request, "tools": tools}) + assert response.status_code == 400, response.text + assert "not supported for this model" in response.text, response.text + outbound: Final = request_json(only_backend_call(databricks_rig, request_nonce)) + assert outbound["tools"] == tools, outbound["tools"] + + +DERIVES_THE_BETA: Final[Mapping[str, tuple[str, bool]]] = { + CLAUDE_MODEL: (CLAUDE_TARGET, True), + AZURE_MODEL: (AZURE_TARGET, True), + DATABRICKS_MODEL: (SERVING_TARGET, True), + NATIVE_MODEL: (NATIVE_TARGET, False), + "safeguards-tencent": (OTHER_NATIVE_MODELS["safeguards-tencent"], False), + "safeguards-minimax": (OTHER_NATIVE_MODELS["safeguards-minimax"], False), + "safeguards-openai-like": (OTHER_NATIVE_MODELS["safeguards-openai-like"], False), +} + + +@pytest.mark.parametrize(("model", "expected"), tuple(DERIVES_THE_BETA.items()), ids=tuple(DERIVES_THE_BETA)) +def test_safeguards_without_a_client_beta_derive_it_only_where_the_beta_map_allows( + databricks_rig: Rig, model: str, expected: tuple[str, bool] +) -> None: + target, derives = expected + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register( + request_nonce, + native_backend(scripted_message(identity, safeguard_results=list(workspace_verdict(identity))), stream=False), + ) + json_body(databricks_rig.post(auto_mode_body(model, request_nonce), beta=None)) + backend_calls: Final = databricks_rig.backend_calls(databricks_rig.received(request_nonce)) + assert [call.target for call in backend_calls] == [target], backend_calls + assert request_json(backend_calls[0])["safeguards"] == list(SAFEGUARDS), backend_calls[0].body + assert (DANGEROUS_TOOL_USE_BETA in beta_values(backend_calls[0])) is derives, backend_calls[0].headers + + +def workspace_rejecting_safeguards(message: Mapping[str, JsonValue]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if "safeguards" not in request_json(request) and DANGEROUS_TOOL_USE_BETA not in beta_values(request): + return Reply(body=json.dumps(message).encode()) + rejection: Final = "Unknown field `safeguards` or beta `dangerous-tool-use-2026-09-03` for this workspace." + return Reply(status=400, body=json.dumps({"error_code": "BAD_REQUEST", "message": rejection}).encode()) + + return respond + + +@pytest.mark.parametrize("beta", (DANGEROUS_TOOL_USE_BETA, None), ids=("client_beta", "no_client_beta")) +def test_a_loaded_beta_map_without_the_databricks_safeguards_beta_drops_safeguards_and_the_beta( + remote_beta_map_rig: tuple[Rig, Wire], beta: str | None +) -> None: + rig, beta_map = remote_beta_map_rig + request_nonce: Final = nonce() + message: Final = scripted_message(tool_id(request_nonce)) + rig.register(request_nonce, workspace_rejecting_safeguards(message)) + body: Final = json_body(rig.post(auto_mode_body(DATABRICKS_MODEL, request_nonce), beta=beta)) + assert body == {**message, "model": DATABRICKS_MODEL}, body + sent: Final = only_backend_call(rig, request_nonce) + assert sent.target == SERVING_TARGET, sent.target + assert "safeguards" not in request_json(sent), sent.body + assert DANGEROUS_TOOL_USE_BETA not in beta_values(sent), sent.headers + assert beta_map.received.qsize() >= 1, "the proxy never fetched the served beta map" + + +@pytest.mark.parametrize("beta", (DANGEROUS_TOOL_USE_BETA, None), ids=("client_beta", "no_client_beta")) +def test_json_provider_gets_neither_safeguards_nor_the_beta_and_the_proxy_classifies( + databricks_rig: Rig, beta: str | None +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + databricks_rig.register(request_nonce, native_backend(scripted_message(identity, model="gpt-4o"), stream=False)) + body: Final = json_body(databricks_rig.post(auto_mode_body(EDENAI_MODEL, request_nonce), beta=beta)) + assert body["safeguard_results"] == flagged((identity,)), body + received: Final = databricks_rig.received(request_nonce) + assert len(databricks_rig.classifier_calls(received)) == 1, [request.target for request in received] + backend_calls: Final = databricks_rig.backend_calls(received) + assert [call.target for call in backend_calls] == [OTHER_NATIVE_MODELS[EDENAI_MODEL]], backend_calls + assert "safeguards" not in request_json(backend_calls[0]), backend_calls[0].body + assert DANGEROUS_TOOL_USE_BETA not in beta_values(backend_calls[0]), backend_calls[0].headers diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_native.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_native.py index ff44d04e180..97139fd81d3 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_native.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_native.py @@ -229,14 +229,22 @@ def test_every_native_messages_provider_is_classified(rig: Rig, model: str, targ classifier_call, backend_calls = one_classifier_call(rig, request_nonce) assert ids_under_review(classifier_call) == (identity,), classifier_call.body assert [call.target for call in backend_calls] == [target], backend_calls - assert request_json(backend_calls[0])["safeguards"] == list(SAFEGUARDS), backend_calls[0].body + assert request_json(backend_calls[0]).get("safeguards") == forwarded_safeguards(model), backend_calls[0].body +JSON_PROVIDER_MODELS: Final = frozenset({"safeguards-edenai"}) NATIVE_TARGETS: Final[Mapping[str, str]] = {NATIVE_MODEL: NATIVE_TARGET, **OTHER_NATIVE_MODELS} +SAFEGUARDS_FORWARDING_TARGETS: Final[Mapping[str, str]] = { + model: target for model, target in NATIVE_TARGETS.items() if model not in JSON_PROVIDER_MODELS +} + + +def forwarded_safeguards(model: str) -> JsonValue: + return None if model in JSON_PROVIDER_MODELS else list(SAFEGUARDS) @pytest.mark.parametrize("stream", (False, True), ids=("non_stream", "stream")) -@pytest.mark.parametrize(("model", "target"), tuple(NATIVE_TARGETS.items())) +@pytest.mark.parametrize(("model", "target"), tuple(SAFEGUARDS_FORWARDING_TARGETS.items())) def test_native_backend_receives_the_request_safeguards_unchanged( rig: Rig, model: str, target: str, stream: bool ) -> None: diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_web_search_short_circuit.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_web_search_short_circuit.py index 1f287c64fa7..48adbcc640d 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_web_search_short_circuit.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_web_search_short_circuit.py @@ -3,7 +3,7 @@ still receives an ``available`` result with no tool uses to judge.""" from __future__ import annotations -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence from typing import Final import pytest @@ -16,6 +16,7 @@ from integration._support.safeguards_rig import ( auto_mode_body, available, json_body, + deployment, nonce, safeguards_rig, sse_events, @@ -26,11 +27,16 @@ from pydantic import JsonValue pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) WEB_SEARCH_TOOLS: Final[list[JsonValue]] = [{"type": "web_search_20250305", "name": "web_search", "max_uses": 1}] +DATABRICKS_CLAUDE_MODEL: Final = "safeguards-databricks-claude" +SEARCHING_MODELS: Final = (CHAT_MODEL, DATABRICKS_CLAUDE_MODEL) LITELLM_SETTINGS: Final[Mapping[str, JsonValue]] = { "callbacks": ["websearch_interception"], - "websearch_interception_params": {"enabled_providers": ["fireworks_ai"], "search_tool_name": "tavily-search"}, + "websearch_interception_params": { + "enabled_providers": ["fireworks_ai", "databricks"], + "search_tool_name": "tavily-search", + }, } @@ -46,10 +52,20 @@ def search_tools(wire_url: str) -> Mapping[str, JsonValue]: return {"search_tools": [dict(tavily)]} +def databricks_claude(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + deployment( + DATABRICKS_CLAUDE_MODEL, "databricks/databricks-claude-sonnet-4-6", f"{wire_url}/dbx/serving-endpoints" + ), + ) + + @pytest.fixture(scope="module") def searching(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: directory: Final = tmp_path_factory.mktemp("web-search") - with safeguards_rig(directory, litellm_settings=LITELLM_SETTINGS, top_level=search_tools) as rig: + with safeguards_rig( + directory, litellm_settings=LITELLM_SETTINGS, top_level=search_tools, extra_models=databricks_claude + ) as rig: yield rig @@ -72,20 +88,22 @@ def assert_one_search_and_nothing_else(rig: Rig, request_nonce: str) -> None: assert received[0].headers.get("authorization") == "Bearer synthetic-tavily-key", received[0].headers -def test_short_circuited_search_is_available_with_no_tool_uses(searching: Rig) -> None: +@pytest.mark.parametrize("model", SEARCHING_MODELS) +def test_short_circuited_search_is_available_with_no_tool_uses(searching: Rig, model: str) -> None: request_nonce: Final = nonce() body: Final = json_body( - searching.post(auto_mode_body(CHAT_MODEL, request_nonce, tools=WEB_SEARCH_TOOLS, prompt="search the web for")) + searching.post(auto_mode_body(model, request_nonce, tools=WEB_SEARCH_TOOLS, prompt="search the web for")) ) assert request_nonce in search_text(body["content"]), body assert body["safeguard_results"] == available({}), body assert_one_search_and_nothing_else(searching, request_nonce) -def test_short_circuited_search_stream_is_available_with_no_tool_uses(searching: Rig) -> None: +@pytest.mark.parametrize("model", SEARCHING_MODELS) +def test_short_circuited_search_stream_is_available_with_no_tool_uses(searching: Rig, model: str) -> None: request_nonce: Final = nonce() status, raw = searching.stream( - auto_mode_body(CHAT_MODEL, request_nonce, tools=WEB_SEARCH_TOOLS, prompt="search the web for") + auto_mode_body(model, request_nonce, tools=WEB_SEARCH_TOOLS, prompt="search the web for") ) assert status == 200, raw assert request_nonce in streamed_text(raw), raw diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index 4f707856fc9..4fbc78c2f4d 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -23,9 +23,17 @@ class TestTryShortCircuitSearch: """Tests for WebSearchInterceptionLogger.try_short_circuit_search""" @pytest.mark.asyncio - async def test_short_circuits_single_web_search_tool(self): + @pytest.mark.parametrize( + ("provider", "model"), + [ + ("github_copilot", "github_copilot/claude-sonnet-5.5"), + ("databricks", "databricks/databricks-claude-sonnet-5-5"), + ("databricks", "databricks/system.ai.claude-sonnet-5-5"), + ], + ) + async def test_short_circuits_single_web_search_tool(self, provider, model): """Single web_search_20250305 tool → short-circuit fires""" - logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + logger = WebSearchInterceptionLogger(enabled_providers=[provider]) with patch.object( logger, "_execute_search", new_callable=AsyncMock @@ -36,14 +44,14 @@ class TestTryShortCircuitSearch: ) result = await logger.try_short_circuit_search( - model="github_copilot/claude-sonnet-5.5", + model=model, messages=[ {"role": "user", "content": "Search for Claude Code releases"} ], tools=[ {"type": "web_search_20250305", "name": "web_search", "max_uses": 8} ], - custom_llm_provider="github_copilot", + custom_llm_provider=provider, ) assert result is not None diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 6323132db36..1308c5aa01f 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6138,6 +6138,12 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): "AmazonAnthropicClaudeConfig", True, ), + ( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation", + "AmazonAnthropicClaudeMessagesConfig", + True, + ), + ("litellm.llms.databricks.messages.transformation", "DatabricksAnthropicMessagesConfig", False), ( "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation", "VertexAIAnthropicConfig", diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index e44cdf2080b..ffe9e46bf56 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1882,6 +1882,20 @@ async def test_native_claude_messages_response_is_left_to_claude_code_when_its_b assert "safeguard_results" not in response +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_native_claude_route_that_drops_safeguards_gets_the_classifier_verdict(stream: bool): + response = await _native_messages( + "sail/claude-sonnet-4-6", _native_upstream(stream=stream, model="claude-sonnet-4-6"), stream=stream + ) + + if stream: + message_deltas = [event for event in response if event["type"] == "message_delta"] + assert [delta["delta"]["safeguard_results"] for delta in message_deltas] == [_FLAGGED_CURL] + else: + assert list(response["safeguard_results"]) == _FLAGGED_CURL + + @pytest.mark.asyncio async def test_anthropic_messages_litellm_router_latency_metadata_tracking(): """ @@ -2248,3 +2262,35 @@ async def test_anthropic_messages_bedrock_dynamic_region(): mock_get_credentials.assert_called_once() credentials_args = mock_get_credentials.call_args.kwargs assert credentials_args.get("aws_region_name") == test_region + + +@pytest.mark.asyncio +@pytest.mark.parametrize("malformed_messages", [["hi"], "hi", [{"role": "user", "content": "hi"}, 7]]) +async def test_anthropic_messages_with_safeguards_and_no_classifier_keep_the_upstream_400_for_malformed_messages( + malformed_messages: object, +) -> None: + from litellm.llms.anthropic.pass_through.messages import handler + + def upstream_rejects_the_messages(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + json={"type": "error", "error": {"type": "invalid_request_error", "message": "messages: invalid"}}, + request=request, + ) + + upstream: Final = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_rejects_the_messages)) + + with pytest.raises(litellm.BadRequestError) as rejected: + await handler.anthropic_messages( + max_tokens=16, + messages=malformed_messages, # pyright: ignore[reportArgumentType] # the malformed body under test + model="anthropic/claude-haiku-4-5", + custom_llm_provider="anthropic", + api_key="sk-test", + client=upstream, + safeguards=[{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}], + ) + + assert rejected.value.status_code == 400 + assert "invalid_request_error" in str(rejected.value) diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py index 0481f7785cd..c63bfbdd301 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_transformation.py @@ -1,5 +1,7 @@ +from collections.abc import Mapping from pathlib import Path from typing import Final +from unittest.mock import patch import pytest @@ -10,6 +12,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( from litellm.llms.anthropic.wif import get_anthropic_wif_token from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig from litellm.llms.tencent.messages.transformation import TencentAnthropicMessagesConfig +from litellm.types.router import GenericLiteLLMParams from tests.unit.llms.anthropic.test_anthropic_wif import ( ScriptedPoster, make_engine, @@ -116,3 +119,184 @@ def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent( assert config._allows_workload_identity is False assert len(poster.requests) == 1 + + +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" +AUTO_MODE_SAFEGUARDS: Final = [ + {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} +] + + +def _classifier_configs() -> tuple[AnthropicMessagesConfig, ...]: + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig + from litellm.llms.databricks.messages.transformation import DatabricksAnthropicMessagesConfig + + return (AnthropicMessagesConfig(), AzureAnthropicMessagesConfig(), DatabricksAnthropicMessagesConfig()) + + +class _BetaMapRejectsTheClassifierConfig(AnthropicMessagesConfig): + @property + def custom_llm_provider(self) -> str | None: + return "bedrock_converse" + + +def _no_classifier_configs() -> tuple[AnthropicMessagesConfig, ...]: + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + from litellm.llms.openai_like.messages.transformation import JSONProviderAnthropicMessagesConfig + + sail: Final = JSONProviderRegistry.get("sail") + assert sail is not None + return (JSONProviderAnthropicMessagesConfig(sail), _BetaMapRejectsTheClassifierConfig()) + + +def _configs_that_forward_safeguards_as_sent() -> tuple[AnthropicMessagesConfig, ...]: + from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig + from litellm.llms.github_copilot.messages.transformation import GithubCopilotAnthropicMessagesConfig + from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig + + return ( + DeepSeekAnthropicMessagesConfig(), + GithubCopilotAnthropicMessagesConfig(), + MinimaxMessagesConfig(), + OpenAILikeAnthropicMessagesConfig(), + ) + + +def _safeguards_request(config: AnthropicMessagesConfig, litellm_params: GenericLiteLLMParams | None = None) -> dict: + return config.transform_anthropic_messages_request( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello"}], + anthropic_messages_optional_request_params={"max_tokens": 64, "safeguards": list(AUTO_MODE_SAFEGUARDS)}, + litellm_params=litellm_params or GenericLiteLLMParams(), + headers={}, + ) + + +def _beta_header(config: AnthropicMessagesConfig, client_beta: str) -> dict: + return config._update_headers_with_anthropic_beta( + headers={"anthropic-beta": client_beta}, optional_params={}, messages=[] + ) + + +def _derived_beta_header(config: AnthropicMessagesConfig) -> dict: + return config._update_headers_with_anthropic_beta( + headers={}, optional_params={"safeguards": list(AUTO_MODE_SAFEGUARDS)}, messages=[] + ) + + +@pytest.mark.parametrize("config", _classifier_configs(), ids=type) +def test_routes_whose_platform_runs_the_classifier_keep_safeguards_and_the_beta( + config: AnthropicMessagesConfig, local_beta_headers_config: None +) -> None: + """Claude Code's auto mode sends `safeguards` with dangerous-tool-use-2026-09-03 and reads + `safeguard_results` back. Anthropic, Azure AI Foundry, and Databricks serving endpoints + honour the pair (probed live 2026-10-12), so both halves must reach them unchanged.""" + assert _safeguards_request(config)["safeguards"] == AUTO_MODE_SAFEGUARDS + assert _beta_header(config, DANGEROUS_TOOL_USE_BETA)["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + assert _derived_beta_header(config)["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + + +@pytest.mark.parametrize("config", _configs_that_forward_safeguards_as_sent(), ids=type) +def test_unjudged_routes_forward_safeguards_and_the_beta_as_sent( + config: AnthropicMessagesConfig, local_beta_headers_config: None +) -> None: + """A platform the beta header map has no column for, or a deployment the operator opted into + /v1/messages passthrough, was never judged, so the request reaches it the way the caller + sent it, as it did before the pair could be dropped.""" + assert _safeguards_request(config)["safeguards"] == AUTO_MODE_SAFEGUARDS + assert _beta_header(config, DANGEROUS_TOOL_USE_BETA)["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + assert _derived_beta_header(config)["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + + +_CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "classifier"} + + +@pytest.mark.parametrize("config", _no_classifier_configs(), ids=type) +@pytest.mark.parametrize( + ("litellm_params", "global_drop_params", "general_settings", "forwarded"), + ( + pytest.param(GenericLiteLLMParams(), False, {}, AUTO_MODE_SAFEGUARDS, id="as_sent"), + pytest.param(GenericLiteLLMParams(drop_params=True), False, {}, None, id="deployment_drop_params"), + pytest.param(GenericLiteLLMParams(), True, {}, None, id="global_drop_params"), + pytest.param(GenericLiteLLMParams(), False, _CLASSIFIER_SETTINGS, None, id="proxy_classifier"), + ), +) +def test_a_route_without_the_classifier_drops_safeguards_only_under_drop_params_or_the_proxy_classifier( + config: AnthropicMessagesConfig, + litellm_params: GenericLiteLLMParams, + global_drop_params: bool, + general_settings: Mapping[str, str], + forwarded: list[dict] | None, + monkeypatch: pytest.MonkeyPatch, + local_beta_headers_config: None, +) -> None: + """A platform without the classifier answers 400 to `safeguards`, which is how litellm treats any + unsupported param: dropped under drop_params, forwarded as sent without it. The proxy's own + classifier answers in the platform's place, so the field is dropped then too.""" + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + with patch("litellm.proxy.proxy_server.general_settings", general_settings): + request: Final = _safeguards_request(config, litellm_params) + assert request.get("safeguards") == forwarded + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("drop_params", "forwarded", "forwarded_beta"), + ( + pytest.param(False, AUTO_MODE_SAFEGUARDS, f"context-1m-2025-08-07,{DANGEROUS_TOOL_USE_BETA}", id="as_sent"), + pytest.param(True, None, "context-1m-2025-08-07", id="drop_params"), + ), +) +async def test_a_json_provider_gets_safeguards_and_the_beta_together_unless_drop_params_drops_both( + drop_params: bool, forwarded: list[dict] | None, forwarded_beta: str, local_beta_headers_config: None +) -> None: + """JSON-configured OpenAI-compatible providers forward every beta header untouched, so under + drop_params the request transform itself drops the beta along with the field.""" + import json + + import httpx + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake platform records what litellm sent + + def platform_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + json={ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(platform_answers)) + try: + response = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello"}], + model="sail/claude-sonnet-4-6", + api_key="test-api-key", + api_base="https://sail.example.test/v1", + safeguards=list(AUTO_MODE_SAFEGUARDS), + extra_headers={"anthropic-beta": f"context-1m-2025-08-07,{DANGEROUS_TOOL_USE_BETA}"}, + drop_params=drop_params, + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1 + sent: Final = received[0] + assert json.loads(sent.content).get("safeguards") == forwarded + assert sent.headers["anthropic-beta"] == forwarded_beta + assert "safeguard_results" not in response + assert response["content"][0]["text"] == "hello" diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_mcp_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_mcp_handler.py index b20aeb40e95..3aca6181d12 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_mcp_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_mcp_handler.py @@ -1,4 +1,5 @@ from importlib import import_module +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -402,3 +403,54 @@ async def test_anthropic_messages_with_mcp_classifies_only_the_response_the_clie assert result["safeguard_results"][0]["status"]["tool_uses"] == { "toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": "unrequested"} } + + +async def _mcp_safeguards_turn_recording_beta_map_loads(custom_llm_provider: str, model: str) -> tuple[int, ...]: + import threading + + from litellm import anthropic_beta_headers_manager + from litellm.llms.anthropic.pass_through.messages import mcp_handler + from litellm.responses.mcp.request_context import MCPRequestContext + + loader_threads: Final[list[int]] = [] # mutable-ok: records the thread each beta map load ran on + + def recording_loader(url: str) -> object: + loader_threads.append(threading.get_ident()) + return anthropic_beta_headers_manager.GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + + with ( + patch.object(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None), + patch.object(anthropic_beta_headers_manager, "get_beta_headers_config", new=recording_loader), + patch.object(MCPRequestContext, "resolve", return_value=MCPRequestContext(user_api_key_auth=None)), + patch.object( + import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + new=AsyncMock(return_value=([], {})), + ), + patch("litellm.anthropic_messages", new=AsyncMock(return_value={"stop_reason": "end_turn", "content": []})), + patch("litellm.proxy.proxy_server.general_settings", {"safeguards_classifier_model": "classifier"}), + ): + await mcp_handler.anthropic_messages_with_mcp( + max_tokens=100, + messages=[{"role": "user", "content": "hi"}], + model=model, + custom_llm_provider=custom_llm_provider, + tools=[MCP_REFERENCE], + safeguards=[{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}], + ) + return tuple(loader_threads) + + +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_loads_a_filtered_routes_beta_map_off_the_event_loop(): + import threading + + loader_threads: Final = await _mcp_safeguards_turn_recording_beta_map_loads("azure_ai", "claude-sonnet-4-6") + + assert len(loader_threads) == 1 + assert threading.get_ident() not in loader_threads + + +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_never_loads_the_beta_map_for_first_party_anthropic(): + assert await _mcp_safeguards_turn_recording_beta_map_loads("anthropic", "claude-sonnet-4-6") == () diff --git a/tests/unit/llms/databricks/messages/__init__.py b/tests/unit/llms/databricks/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/messages/test_databricks_messages_transformation.py b/tests/unit/llms/databricks/messages/test_databricks_messages_transformation.py new file mode 100644 index 00000000000..3c5f4fd72f2 --- /dev/null +++ b/tests/unit/llms/databricks/messages/test_databricks_messages_transformation.py @@ -0,0 +1,752 @@ +import json +import sys +from collections.abc import AsyncIterator +from types import ModuleType, SimpleNamespace +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.databricks.chat.transformation import DatabricksConfig +from litellm.llms.databricks.messages.transformation import DatabricksAnthropicMessagesConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager + +DATABRICKS_API_BASE: Final = "https://my.workspace.cloud.databricks.com/serving-endpoints" +DATABRICKS_API_KEY: Final = "dapimykey" +CLAUDE_MODEL: Final = "databricks-claude-sonnet-4-6" +UNITY_CATALOG_CLAUDE_MODEL: Final = "system.ai.claude-sonnet-4-6" +WORKSPACE_URL: Final = "https://my.workspace.cloud.databricks.com" +MESSAGES_URL: Final = f"{DATABRICKS_API_BASE}/anthropic/v1/messages" +ANTHROPIC_ROUTE_API_BASE: Final = f"{DATABRICKS_API_BASE}/anthropic" +OPTED_IN_MODEL_INFO: Final = {"supported_endpoints": ["/v1/messages"]} +AI_GATEWAY_MESSAGES_URL: Final = f"{WORKSPACE_URL}/ai-gateway/anthropic/v1/messages" +DANGEROUS_TOOL_USE_BETA: Final = "dangerous-tool-use-2026-09-03" +AUTO_MODE_SAFEGUARDS: Final = [ + {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} +] +SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": {"type": "available", "tool_uses": {"toolu_bdrk_01": {"type": "evaluated", "outcome": "not_flagged"}}}, + } +] +BASH_TOOL: Final = { + "name": "Bash", + "description": "Run a shell command", + "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}, +} +USER_TURN: Final = [{"role": "user", "content": "Use the Bash tool to run: echo hello"}] + + +@pytest.mark.parametrize( + ("model", "messages_url"), + [(CLAUDE_MODEL, MESSAGES_URL), (UNITY_CATALOG_CLAUDE_MODEL, AI_GATEWAY_MESSAGES_URL)], +) +def test_foundation_claude_names_dispatch_to_a_native_messages_route(model: str, messages_url: str) -> None: + config: Final = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=litellm.LlmProviders.DATABRICKS + ) + + assert config is not None + assert ( + config.get_complete_url( + api_base=DATABRICKS_API_BASE, api_key=None, model=model, optional_params={}, litellm_params={} + ) + == messages_url + ) + + +@pytest.mark.parametrize("model", ["databricks-gpt-oss-120b", "my-claude-finetune", "main.default.claude-proxy"]) +def test_names_outside_the_foundation_claude_models_keep_the_chat_completions_bridge(model: str) -> None: + config: Final = ProviderConfigManager.get_provider_anthropic_messages_config( + model=model, provider=litellm.LlmProviders.DATABRICKS + ) + + assert config is None + + +@pytest.mark.parametrize( + "api_base", + [DATABRICKS_API_BASE, f"{DATABRICKS_API_BASE}/", ANTHROPIC_ROUTE_API_BASE, MESSAGES_URL], +) +def test_complete_url_targets_the_workspace_anthropic_messages_route(api_base: str) -> None: + url: Final = DatabricksAnthropicMessagesConfig().get_complete_url( + api_base=api_base, api_key=None, model=CLAUDE_MODEL, optional_params={}, litellm_params={} + ) + + assert url == MESSAGES_URL + + +def test_validate_environment_signs_with_the_pat_and_keeps_the_workspace_base() -> None: + headers, api_base = DatabricksAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model=CLAUDE_MODEL, + messages=USER_TURN, + optional_params={}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + ) + + assert api_base == DATABRICKS_API_BASE + assert headers["Authorization"] == f"Bearer {DATABRICKS_API_KEY}" + assert headers["anthropic-version"] == "2023-06-01" + assert headers["content-type"] == "application/json" + assert headers["User-Agent"].startswith("litellm/") + + +def test_validate_environment_reads_the_workspace_credentials_from_the_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DATABRICKS_API_KEY", DATABRICKS_API_KEY) + monkeypatch.setenv("DATABRICKS_API_BASE", DATABRICKS_API_BASE) + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + config: Final = DatabricksAnthropicMessagesConfig() + + headers, api_base = config.validate_anthropic_messages_environment( + headers={}, + model=CLAUDE_MODEL, + messages=USER_TURN, + optional_params={}, + litellm_params={}, + ) + url: Final = config.get_complete_url( + api_base=api_base, api_key=None, model=CLAUDE_MODEL, optional_params={}, litellm_params={} + ) + + assert headers["Authorization"] == f"Bearer {DATABRICKS_API_KEY}" + assert url == MESSAGES_URL + + +@pytest.mark.parametrize( + "api_base", + [DATABRICKS_API_BASE, f"{DATABRICKS_API_BASE}/", ANTHROPIC_ROUTE_API_BASE, MESSAGES_URL], +) +def test_validation_then_url_building_appends_the_messages_route_once(api_base: str) -> None: + config: Final = DatabricksAnthropicMessagesConfig() + _, validated_api_base = config.validate_anthropic_messages_environment( + headers={}, + model=CLAUDE_MODEL, + messages=USER_TURN, + optional_params={}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=api_base, + ) + + url: Final = config.get_complete_url( + api_base=validated_api_base, api_key=DATABRICKS_API_KEY, model=CLAUDE_MODEL, optional_params={}, litellm_params={} + ) + + assert url == MESSAGES_URL + + +@pytest.mark.parametrize( + "api_base", + [ + DATABRICKS_API_BASE, + WORKSPACE_URL, + f"{WORKSPACE_URL}/", + f"{WORKSPACE_URL}/ai-gateway/mlflow/v1", + f"{WORKSPACE_URL}/ai-gateway/anthropic", + AI_GATEWAY_MESSAGES_URL, + ANTHROPIC_ROUTE_API_BASE, + MESSAGES_URL, + f"{MESSAGES_URL}/", + ], +) +def test_unity_catalog_claude_names_route_to_the_ai_gateway_messages_route(api_base: str) -> None: + config: Final = DatabricksAnthropicMessagesConfig() + _, validated_api_base = config.validate_anthropic_messages_environment( + headers={}, + model=UNITY_CATALOG_CLAUDE_MODEL, + messages=USER_TURN, + optional_params={}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=api_base, + ) + + url: Final = config.get_complete_url( + api_base=validated_api_base, + api_key=DATABRICKS_API_KEY, + model=UNITY_CATALOG_CLAUDE_MODEL, + optional_params={}, + litellm_params={}, + ) + + assert url == AI_GATEWAY_MESSAGES_URL + + +@pytest.mark.parametrize( + ("model", "chat_url", "messages_url"), + [ + (CLAUDE_MODEL, f"{DATABRICKS_API_BASE}/chat/completions", MESSAGES_URL), + (UNITY_CATALOG_CLAUDE_MODEL, f"{WORKSPACE_URL}/ai-gateway/mlflow/v1/chat/completions", AI_GATEWAY_MESSAGES_URL), + ], +) +def test_messages_and_chat_routes_agree_on_which_names_use_the_ai_gateway( + model: str, chat_url: str, messages_url: str +) -> None: + chat: Final = DatabricksConfig().get_complete_url( + api_base=DATABRICKS_API_BASE, api_key=DATABRICKS_API_KEY, model=model, optional_params={}, litellm_params={} + ) + messages: Final = DatabricksAnthropicMessagesConfig().get_complete_url( + api_base=DATABRICKS_API_BASE, api_key=DATABRICKS_API_KEY, model=model, optional_params={}, litellm_params={} + ) + + assert (chat, messages) == (chat_url, messages_url) + + +def test_partner_user_agent_matches_the_chat_route(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DATABRICKS_USER_AGENT", "partner/1.0") + messages_headers, _ = DatabricksAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model=CLAUDE_MODEL, + messages=USER_TURN, + optional_params={}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + ) + chat_headers: Final = DatabricksConfig().validate_environment( + headers={}, + model=CLAUDE_MODEL, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + ) + + assert messages_headers["User-Agent"].startswith("partner_litellm/") + assert messages_headers["User-Agent"] == chat_headers["User-Agent"] + + +def test_request_transform_normalizes_system_turns_and_strips_cache_scope() -> None: + request: Final = DatabricksAnthropicMessagesConfig().transform_anthropic_messages_request( + model=CLAUDE_MODEL, + messages=[ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "reply"}, + {"role": "system", "content": "mid-turn instruction"}, + {"role": "user", "content": "second"}, + ], + anthropic_messages_optional_request_params={ + "max_tokens": 64, + "system": [{"type": "text", "text": "be brief", "cache_control": {"type": "ephemeral", "scope": "global"}}], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert [message["role"] for message in request["messages"]] == ["user", "assistant", "user", "user"] + assert "mid-turn instruction" in json.dumps(request["messages"][2]["content"]) + assert request["system"][0]["cache_control"] == {"type": "ephemeral"} + + +def test_request_transform_keeps_the_billing_header_block() -> None: + billing_block: Final = {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.296; cc_entrypoint=cli"} + request: Final = DatabricksAnthropicMessagesConfig().transform_anthropic_messages_request( + model=CLAUDE_MODEL, + messages=[{"role": "user", "content": "first"}], + anthropic_messages_optional_request_params={ + "max_tokens": 64, + "system": [billing_block, {"type": "text", "text": "be brief"}], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["system"] == [billing_block, {"type": "text", "text": "be brief"}] + + +def _reply(request: httpx.Request, body: dict[str, object]) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + json=body, + request=request, + ) + + +SERVER_TOOLS: Final = [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 8}, + {"type": "web_fetch_20250910", "name": "web_fetch", "max_uses": 1}, + {"type": "code_execution_20250825", "name": "code_execution"}, +] + + +async def _workspace_request_body(**request_params: object) -> dict[str, object]: + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what litellm sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply( + request, + { + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + client=upstream, + **request_params, + ) + finally: + await upstream.close() + assert len(received) == 1 + return json.loads(received[0].content) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "tool_choice", "expected_tools", "expected_tool_choice"), + [ + ([*SERVER_TOOLS, BASH_TOOL], {"type": "any"}, [BASH_TOOL], {"type": "any"}), + ([*SERVER_TOOLS, BASH_TOOL], {"type": "tool", "name": "web_search"}, [BASH_TOOL], None), + (SERVER_TOOLS[:1], {"type": "auto"}, None, None), + ([{"type": "web_search_20250305", "max_uses": 8}, BASH_TOOL], {"type": "any"}, [BASH_TOOL], {"type": "any"}), + ], +) +async def test_drop_params_drops_the_server_tools_the_workspace_rejects( + tools: list[dict[str, object]], + tool_choice: dict[str, object], + expected_tools: list[dict[str, object]] | None, + expected_tool_choice: dict[str, object] | None, +) -> None: + body: Final = await _workspace_request_body(tools=tools, tool_choice=tool_choice, drop_params=True) + + assert body.get("tools") == expected_tools + assert body.get("tool_choice") == expected_tool_choice + + +@pytest.mark.asyncio +async def test_server_tools_reach_the_workspace_as_sent_without_drop_params() -> None: + body: Final = await _workspace_request_body(tools=[*SERVER_TOOLS, BASH_TOOL], tool_choice={"type": "any"}) + + assert body["tools"] == [*SERVER_TOOLS, BASH_TOOL] + assert body["tool_choice"] == {"type": "any"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [True, False]) +@pytest.mark.parametrize("container", ["cntr_abc", {"id": "cntr_abc"}]) +async def test_container_never_reaches_the_workspace(drop_params: bool, container: object) -> None: + body: Final = await _workspace_request_body( + container=container, metadata={"user_id": "user-1"}, drop_params=drop_params + ) + + assert "container" not in body + assert body["metadata"] == {"user_id": "user-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_agent_param", ["databricks_user_agent", "user_agent"]) +async def test_per_call_partner_user_agent_reaches_the_workspace_like_the_chat_route(user_agent_param: str) -> None: + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what litellm sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply( + request, + { + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + client=upstream, + **{user_agent_param: "partnerx/1.0"}, + ) + finally: + await upstream.close() + chat_headers: Final = DatabricksConfig().validate_environment( + headers={}, + model=CLAUDE_MODEL, + messages=[{"role": "user", "content": "hi"}], + optional_params={user_agent_param: "partnerx/1.0"}, + litellm_params={}, + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + ) + + assert len(received) == 1 + assert received[0].headers["user-agent"].startswith("partnerx_litellm/") + assert received[0].headers["user-agent"] == chat_headers["User-Agent"] + assert user_agent_param not in json.loads(received[0].content) + + +@pytest.mark.asyncio +async def test_safeguards_reach_the_workspace_and_safeguard_results_come_back( + local_beta_headers_config: None, +) -> None: + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what the proxy sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply( + request, + { + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [ + {"type": "tool_use", "id": "toolu_bdrk_01", "name": "Bash", "input": {"command": "echo hello"}} + ], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + "safeguard_results": SAFEGUARD_RESULTS, + }, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + response = await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + tools=[BASH_TOOL], + safeguards=AUTO_MODE_SAFEGUARDS, + client=upstream, + ) + finally: + await upstream.close() + + assert [str(request.url) for request in received] == [MESSAGES_URL] + sent: Final = received[0] + sent_body: Final = json.loads(sent.content) + assert sent.headers["authorization"] == f"Bearer {DATABRICKS_API_KEY}" + assert sent.headers["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + assert sent_body["safeguards"] == AUTO_MODE_SAFEGUARDS + assert sent_body["model"] == CLAUDE_MODEL + assert response["safeguard_results"] == SAFEGUARD_RESULTS + assert response["content"][0]["id"] == "toolu_bdrk_01" + + +@pytest.mark.asyncio +async def test_streaming_message_delta_keeps_safeguard_results(local_beta_headers_config: None) -> None: + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None, "safeguard_results": SAFEGUARD_RESULTS}, + "usage": {"output_tokens": 6}, + } + events: Final = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_bdrk_01", "name": "Bash", "input": {}}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"command": "echo hello"}'}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", message_delta), + ("message_stop", {"type": "message_stop"}), + ) + sse_body: Final = "".join(f"event: {name}\ndata: {json.dumps(payload)}\n\n" for name, payload in events).encode() + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what the proxy sent + + def workspace_streams(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + content=sse_body, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_streams)) + try: + stream = await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + tools=[BASH_TOOL], + safeguards=AUTO_MODE_SAFEGUARDS, + stream=True, + client=upstream, + ) + relayed: Final = b"".join([chunk async for chunk in stream]) + finally: + await upstream.close() + + assert [str(request.url) for request in received] == [MESSAGES_URL] + assert received[0].headers["anthropic-beta"] == DANGEROUS_TOOL_USE_BETA + sent_body: Final = json.loads(received[0].content) + assert sent_body["safeguards"] == AUTO_MODE_SAFEGUARDS + assert sent_body["stream"] is True + relayed_deltas: Final = [ + json.loads(line.removeprefix(b"data: ")) + for line in relayed.splitlines() + if line.startswith(b'data: {"type": "message_delta"') + ] + assert len(relayed_deltas) == 1, relayed.decode() + assert relayed_deltas[0]["delta"]["safeguard_results"] == SAFEGUARD_RESULTS + assert b'"id": "toolu_bdrk_01"' in relayed + + +CHAT_COMPLETION: Final = { + "id": "chatcmpl_01", + "object": "chat.completion", + "created": 1792227137, + "model": "claude-sonnet-4-6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 12, "completion_tokens": 6, "total_tokens": 18}, +} +MESSAGE: Final = { + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "api_base", "request_params", "expected_url"), + [ + (UNITY_CATALOG_CLAUDE_MODEL, "http://gateway.internal/dbx", {}, "http://gateway.internal/dbx/chat/completions"), + (CLAUDE_MODEL, "http://gateway.internal/dbx", {}, "http://gateway.internal/dbx/chat/completions"), + ( + CLAUDE_MODEL, + "http://gateway.internal/dbx-anthropic", + {"model_info": OPTED_IN_MODEL_INFO}, + "http://gateway.internal/dbx-anthropic/v1/messages", + ), + (CLAUDE_MODEL, ANTHROPIC_ROUTE_API_BASE, {"model_info": OPTED_IN_MODEL_INFO}, MESSAGES_URL), + (CLAUDE_MODEL, "http://gateway.internal", {}, "http://gateway.internal/chat/completions"), + (UNITY_CATALOG_CLAUDE_MODEL, WORKSPACE_URL, {}, AI_GATEWAY_MESSAGES_URL), + (UNITY_CATALOG_CLAUDE_MODEL, f"{WORKSPACE_URL}/", {}, AI_GATEWAY_MESSAGES_URL), + ], +) +async def test_each_api_base_reaches_a_route_the_workspace_serves( + model: str, api_base: str, request_params: dict[str, object], expected_url: str +) -> None: + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake upstream records what litellm sent + + def upstream_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply(request, CHAT_COMPLETION if request.url.path.endswith("/chat/completions") else MESSAGE) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(upstream_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{model}", + api_key=DATABRICKS_API_KEY, + api_base=api_base, + client=upstream, + **request_params, + ) + finally: + await upstream.close() + + assert [str(request.url) for request in received] == [expected_url] + + +SDK_TOKEN: Final = "sdk-test-token" + + +@pytest.mark.asyncio +async def test_client_anthropic_headers_keep_the_sdk_credentials_on_the_native_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for name in ("DATABRICKS_API_BASE", "DATABRICKS_API_KEY", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + sdk_config: Final = SimpleNamespace(host=WORKSPACE_URL, authenticate=lambda: {"Authorization": f"Bearer {SDK_TOKEN}"}) + sdk_module: Final = ModuleType("databricks.sdk") + setattr(sdk_module, "WorkspaceClient", lambda: SimpleNamespace(config=sdk_config)) + setattr(sdk_module, "useragent", SimpleNamespace(with_partner=lambda _: None)) + package_module: Final = ModuleType("databricks") + package_module.__path__ = [] + setattr(package_module, "sdk", sdk_module) + monkeypatch.setitem(sys.modules, "databricks", package_module) + monkeypatch.setitem(sys.modules, "databricks.sdk", sdk_module) + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what litellm sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply(request, MESSAGE) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_base=DATABRICKS_API_BASE, + client=upstream, + provider_specific_header={ + "custom_llm_provider": "databricks", + "extra_headers": {"anthropic-beta": "context-1m-2025-08-07", "anthropic-version": "2023-06-01"}, + "native_anthropic_messages_only": True, + }, + ) + finally: + await upstream.close() + + assert [str(request.url) for request in received] == [MESSAGES_URL] + assert received[0].headers["Authorization"] == f"Bearer {SDK_TOKEN}" + assert received[0].headers["anthropic-beta"] == "context-1m-2025-08-07" + + +@pytest.mark.asyncio +async def test_client_anthropic_headers_skip_an_opted_in_passthrough_api_base() -> None: + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake upstream records what litellm sent + + def upstream_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _reply(request, MESSAGE) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(upstream_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base="http://gateway.internal/dbx-anthropic", + client=upstream, + model_info=OPTED_IN_MODEL_INFO, + provider_specific_header={ + "custom_llm_provider": "databricks", + "extra_headers": {"anthropic-beta": "prompt-caching-scope-2026-01-05", "anthropic-version": "2023-01-01"}, + "native_anthropic_messages_only": True, + }, + ) + finally: + await upstream.close() + + assert [str(request.url) for request in received] == ["http://gateway.internal/dbx-anthropic/v1/messages"] + assert "anthropic-beta" not in received[0].headers + assert received[0].headers["anthropic-version"] != "2023-01-01" + + +class _ChunkedStream(httpx.AsyncByteStream): + def __init__(self, body: bytes, chunk_size: int) -> None: + self._chunks: Final = tuple(body[start : start + chunk_size] for start in range(0, len(body), chunk_size)) + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk_size", [7, 4096]) +async def test_streamed_message_stop_drops_the_bedrock_invocation_metrics(chunk_size: int) -> None: + events: Final = ( + ("message_start", {"type": "message_start", "message": {**MESSAGE, "content": [], "stop_reason": None}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}), + ( + "message_stop", + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": 18, + "outputTokenCount": 4, + "invocationLatency": 496, + "firstByteLatency": 487, + }, + }, + ), + ) + sse_body: Final = "".join(f"event: {name}\ndata: {json.dumps(payload)}\n\n" for name, payload in events).encode() + + def workspace_streams(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "text/event-stream"}, + stream=_ChunkedStream(sse_body, chunk_size), + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_streams)) + try: + stream = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=USER_TURN, + model=f"databricks/{CLAUDE_MODEL}", + api_key=DATABRICKS_API_KEY, + api_base=DATABRICKS_API_BASE, + stream=True, + client=upstream, + ) + relayed: Final = b"".join([chunk async for chunk in stream]) + finally: + await upstream.close() + + relayed_events: Final = [ + json.loads(line.removeprefix(b"data:")) for line in relayed.splitlines() if line.startswith(b"data:") + ] + assert relayed_events == [*(payload for _, payload in events[:-1]), {"type": "message_stop"}], relayed.decode() diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 20a13bdc566..2fc1b72c932 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -2007,13 +2007,145 @@ def test_provider_specific_header(): data=data, headers=headers, ) - # Verify multi-provider support: anthropic headers work across multiple providers - assert data["provider_specific_header"] == { - "custom_llm_provider": "anthropic,bedrock,bedrock_mantle,vertex_ai", - "extra_headers": { - "anthropic-beta": "prompt-caching-2024-07-31", + from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils + + for provider in ("anthropic", "bedrock", "bedrock_mantle", "vertex_ai"): + assert ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=data["provider_specific_header"], custom_llm_provider=provider + ) == {"anthropic-beta": "prompt-caching-2024-07-31"} + assert ( + ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=data["provider_specific_header"], custom_llm_provider="databricks" + ) + == {} + ) + assert ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=data["provider_specific_header"], + custom_llm_provider="databricks", + native_anthropic_messages=True, + ) == {"anthropic-beta": "prompt-caching-2024-07-31"} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_client_anthropic_betas_reach_databricks_claude_filtered_to_the_ones_the_workspace_accepts(): + import httpx + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy.litellm_pre_call_utils import add_provider_specific_headers_to_request + + data: Final[dict] = {"model": "databricks-claude-sonnet-4-6"} # mutable-ok: the function under test writes into it + add_provider_specific_headers_to_request( + data=data, + headers={ + "content-type": "application/json", + "anthropic-beta": "context-1m-2025-08-07,prompt-caching-scope-2026-01-05", + "anthropic-version": "2023-06-01", }, - } + ) + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what the proxy sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_bdrk_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="databricks/databricks-claude-sonnet-4-6", + api_key="dapimykey", + api_base="https://my.workspace.cloud.databricks.com/serving-endpoints", + client=upstream, + provider_specific_header=data["provider_specific_header"], + ) + finally: + await upstream.close() + + assert received[0].headers["anthropic-beta"] == "context-1m-2025-08-07" + + +_CLIENT_ANTHROPIC_HEADERS: Final = { + "content-type": "application/json", + "anthropic-beta": "context-1m-2025-08-07", + "anthropic-version": "2023-06-01", +} +_DATABRICKS_CHAT_COMPLETION: Final = { + "id": "chatcmpl-01", + "object": "chat.completion", + "created": 1760000000, + "model": "databricks-gpt-oss-120b", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 12, "completion_tokens": 6, "total_tokens": 18}, +} + + +async def _databricks_chat_completion(provider_specific_header: object, upstream: object) -> None: + await litellm.acompletion( + model="databricks/databricks-gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + api_key="dapimykey", + api_base="https://my.workspace.cloud.databricks.com/serving-endpoints", + client=upstream, + provider_specific_header=provider_specific_header, + ) + + +async def _databricks_messages_bridge(provider_specific_header: object, upstream: object) -> None: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="databricks/databricks-gpt-oss-120b", + api_key="dapimykey", + api_base="https://my.workspace.cloud.databricks.com/serving-endpoints", + client=upstream, + provider_specific_header=provider_specific_header, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_databricks", + [_databricks_chat_completion, _databricks_messages_bridge], + ids=["chat_completions", "messages_bridge"], +) +async def test_client_anthropic_headers_never_reach_databricks_chat_completions(call_databricks): + import httpx + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy.litellm_pre_call_utils import add_provider_specific_headers_to_request + + data: Final[dict] = {"model": "databricks-gpt-oss-120b"} # mutable-ok: the function under test writes into it + add_provider_specific_headers_to_request(data=data, headers=_CLIENT_ANTHROPIC_HEADERS) + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake workspace records what the proxy sent + + def workspace_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response(status_code=200, json=_DATABRICKS_CHAT_COMPLETION, request=request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(workspace_answers)) + try: + await call_databricks(data["provider_specific_header"], upstream) + finally: + await upstream.close() + + assert [request.url.path for request in received] == ["/serving-endpoints/chat/completions"] + assert "anthropic-beta" not in received[0].headers + assert "anthropic-version" not in received[0].headers def test_provider_specific_header_multi_provider(): @@ -2080,10 +2212,10 @@ def test_provider_specific_header_multi_provider(): # Verify that provider_specific_header contains comma-separated providers assert "provider_specific_header" in data assert ( - data["provider_specific_header"]["custom_llm_provider"] + data["provider_specific_header"][0]["custom_llm_provider"] == "anthropic,bedrock,bedrock_mantle,vertex_ai" ) - assert data["provider_specific_header"]["extra_headers"] == { + assert data["provider_specific_header"][0]["extra_headers"] == { "anthropic-beta": "context-1m-2025-08-07", "anthropic-version": "2023-06-01", }