mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
fix(databricks): use the native Anthropic route for Claude /v1/messages (#46091)
* fix(databricks): use the native Anthropic route for Claude /v1/messages Claude models on a Databricks workspace were bridged through chat completions on /v1/messages, which drops the safeguards field and never returns safeguard_results. They now call the workspace's /serving-endpoints/anthropic/v1/messages, the proxy forwards the client's anthropic-beta and anthropic-version headers to Databricks deployments, and the beta header map marks dangerous-tool-use-2026-09-03 as supported there. Any route whose platform does not take that beta drops safeguards together with it instead of answering 400. Non-Claude Databricks models keep the chat completions bridge * fix(databricks): derive the dangerous-tool-use beta from safeguards instead of forwarding client betas A Databricks Claude serving endpoint answers 400 to betas Claude Code sends on every turn, prompt-caching-scope-2026-01-05 among them, so the proxy no longer forwards the client's anthropic-beta header to Databricks deployments. The beta the classifier needs is derived from the safeguards body field, the same way the Azure AI route derives it. The Rust beta-header mirror stops rejecting the Databricks dangerous-tool-use cell so its config-cell test matches the JSON again * fix(databricks): keep the billing block, build the messages route once, and pass the partner user agent * fix(databricks): route Unity Catalog Claude names to the gateway messages route and keep env credentials Unity Catalog model-service names now reach the AI Gateway's native Anthropic route instead of a serving-endpoints path that answers 404 for them, and the native route resolves the key and base through the same order the chat route uses, so environment-only deployments keep working. The native route is limited to Databricks foundation-model Claude names, an unlisted provider forwards safeguards and the beta as sent, and the ten betas the workspace rejects are marked unsupported for Databricks in the beta map and its Rust mirror. * test(rust): expect Databricks to drop the advanced-tool-use beta * fix(databricks): keep passthrough safeguards, partner user agent, and classifier fallback Only the OpenAI-compatible JSON providers drop safeguards now, so a deployment opted into /v1/messages passthrough forwards the field and the beta as sent. When a route drops safeguards and the proxy has a classifier model, the proxy answers the verdict itself, as the chat completions bridge did. The Databricks Messages route also reads the databricks_user_agent param like the chat route. * fix(databricks): drop rejected server tools under drop_params, short-circuit web search, route UC names from a full Messages URL * fix(databricks): drop container on the native Messages route, as the chat completions bridge did * fix(databricks): keep every api_base on a route the workspace serves and gate the derived beta on the map Gate the dangerous-tool-use beta derivation on the beta map the way #43955 does, so GitHub Copilot, DeepSeek, and opt-in passthroughs keep the merge base wire. Recognize /serving-endpoints/anthropic api_bases, fall back to the merge base routing for an api_base path the workspace trim does not know, strip Bedrock invocation metrics from the streamed message_stop, forward client betas to Databricks filtered by the map, and keep tool_choice when a dropped server tool has no name * test(databricks): cover the native Messages route, derived beta and JSON provider safeguards drop on the wire * test(anthropic): expect the dangerous-tool-use beta derived from safeguards next to the client's betas * fix(databricks): send client anthropic headers only on the native Messages route and keep SDK auth Databricks chat completions, called directly or through the Messages bridge, no longer receive the client's anthropic-beta and anthropic-version, as at the merge base. The native Messages route still forwards them, filtered by the beta map, without counting them as caller credentials, so a deployment with no api_key keeps falling back to Databricks SDK auth. A pathless api_base goes to the native route only for Unity Catalog names, matching the chat route. * test(databricks): cover SDK auth with client anthropic headers and a Unity Catalog workspace root on the native route * fix(databricks): keep client anthropic headers off an opted-in passthrough api_base A databricks Claude deployment whose api_base is no workspace route and that opts into /v1/messages with supported_endpoints falls back to the OpenAI-like passthrough. That config does not filter betas through the beta map, yet it received the native-only Databricks header entry, so Claude Code's prompt-caching-scope beta reached the upstream raw and got a 400 where the merge base answered 200. Each Messages config now says whether it is the provider's native route, the OpenAI-like passthrough says no, and the handler passes that answer when it reads provider-specific headers, so native-only entries reach only native configs. * test(databricks): cover client anthropic betas on an opted-in passthrough api_base A databricks Claude deployment whose api_base is no workspace route and whose model_info.supported_endpoints opts into /v1/messages takes the OpenAI-like passthrough. The cell sends Claude Code's prompt-caching-scope and context-1m betas through the proxy and asserts the caller gets the upstream's message and the gateway receives no anthropic-beta, as on the merge base. The scripted gateway answers 400 on the prompt-caching-scope beta, as the real one does * fix(anthropic): keep beta map loads off shared threads and build the safeguards evaluator only when the proxy classifies Unfiltered configs such as first-party Anthropic never read the beta map to decide safeguards support. Filtered configs await the map on its dedicated loader before reading it, on both the native and the MCP paths. The native handler decides after the upstream call and builds the evaluator only when the proxy classifies, so malformed messages keep the upstream 400. The Databricks dispatch tests assert the dispatched route instead of the config class, and the sail and Eden AI wire rows expect safeguards dropped, as JSON providers do. * test(databricks): cover a loaded beta map that marks the safeguards beta unsupported A remote beta map that nulls dangerous-tool-use for databricks while the bundled map allows it must send neither safeguards nor the beta to the serving endpoint, with or without a client beta, and the caller gets a 200 with no verdict. * test(anthropic): annotate the MCP beta map recordings as Final * fix(anthropic): drop safeguards on routes without the classifier only under drop_params A route whose upstream does not answer safeguards (a JSON provider, or a map that marks the dangerous-tool-use beta unsupported) now forwards the field and the beta as the merge base did unless drop_params is set or the proxy runs its own classifier, matching how litellm treats any unsupported param. Databricks keeps dropping both, since its merge base bridge never forwarded the field. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
a6245c1bbb
commit
4bebcdf757
33 changed files with 2310 additions and 109 deletions
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]] = (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
194
litellm/llms/databricks/messages/transformation.py
Normal file
194
litellm/llms/databricks/messages/transformation.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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") == ()
|
||||
|
|
|
|||
0
tests/unit/llms/databricks/messages/__init__.py
Normal file
0
tests/unit/llms/databricks/messages/__init__.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue