mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
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.
This commit is contained in:
mateo-berri
parent
86962c66a2
commit
d4828713ba
7 changed files with 139 additions and 31 deletions
|
|
@ -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 (
|
||||
|
|
@ -99,11 +100,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),
|
||||
|
|
@ -114,6 +121,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.
|
||||
|
||||
|
|
@ -154,16 +165,20 @@ def _native_messages_config(
|
|||
return OpenAILikeAnthropicMessagesConfig(cache_control_ttl=_deployment_supports_cache_control_ttl(model_info))
|
||||
|
||||
|
||||
def _proxy_classifies(model: str, safeguards: object, native_config: BaseAnthropicMessagesConfig | None) -> bool:
|
||||
if native_config is not None and native_config.supports_safeguards():
|
||||
return native_route_classifies(model, safeguards)
|
||||
return requested_dangerous_tool_use(safeguards) is not None and read_classifier_model_setting() is not None
|
||||
async def _proxy_classifies(model: str, safeguards: object, native_config: BaseAnthropicMessagesConfig | None) -> bool:
|
||||
if requested_dangerous_tool_use(safeguards) is None or read_classifier_model_setting() is None:
|
||||
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)
|
||||
|
||||
|
||||
def proxy_classifies_tool_uses(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
|
||||
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."""
|
||||
api_base: Final = kwargs.get("api_base")
|
||||
return _proxy_classifies(
|
||||
return await _proxy_classifies(
|
||||
model,
|
||||
kwargs.get("safeguards"),
|
||||
_native_messages_config(
|
||||
|
|
@ -697,12 +712,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 _proxy_classifies(model, safeguard_kwargs.get("safeguards"), anthropic_messages_provider_config)
|
||||
else None
|
||||
)
|
||||
native_response: Final = base_llm_http_handler.anthropic_messages_handler(
|
||||
model=model,
|
||||
messages=cast( # cast-ok: anthropic messages are dicts at runtime
|
||||
|
|
@ -723,6 +732,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]] = (
|
||||
|
|
|
|||
|
|
@ -118,8 +118,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
map does not list keeps the body field, and the beta goes with it only where that
|
||||
platform's config forwards betas unfiltered.
|
||||
"""
|
||||
return ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value not in get_unsupported_headers(
|
||||
self._resolved_provider
|
||||
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 _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
@ -532,11 +533,13 @@ _PROVIDERS: Final = MappingProxyType(
|
|||
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,
|
||||
forwards_safeguards=False,
|
||||
),
|
||||
"D21_bedrock_claude_platform": _ProviderCase(
|
||||
parameters={
|
||||
|
|
@ -601,7 +604,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
|
||||
|
|
|
|||
|
|
@ -2221,3 +2221,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)
|
||||
|
|
|
|||
|
|
@ -402,3 +402,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: 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 = 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") == ()
|
||||
|
|
|
|||
|
|
@ -41,22 +41,23 @@ BASH_TOOL: Final = {
|
|||
USER_TURN: Final = [{"role": "user", "content": "Use the Bash tool to run: echo hello"}]
|
||||
|
||||
|
||||
def test_claude_serving_endpoints_dispatch_to_the_native_messages_config() -> None:
|
||||
@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=CLAUDE_MODEL, provider=litellm.LlmProviders.DATABRICKS
|
||||
model=model, provider=litellm.LlmProviders.DATABRICKS
|
||||
)
|
||||
|
||||
assert isinstance(config, DatabricksAnthropicMessagesConfig)
|
||||
assert config.custom_llm_provider == "databricks"
|
||||
|
||||
|
||||
def test_unity_catalog_claude_names_dispatch_to_the_native_messages_config() -> None:
|
||||
config: Final = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model=UNITY_CATALOG_CLAUDE_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
|
||||
)
|
||||
|
||||
assert isinstance(config, DatabricksAnthropicMessagesConfig)
|
||||
|
||||
|
||||
@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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue