From 59538ff1d99282ba973872c7689086ba5d54f188 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 11 Nov 2026 16:07:20 -0800 Subject: [PATCH] feat(messages): serve anthropic messages on rust only (#51119) * feat(messages): apply modify_params rewrites on rust * feat(messages): mint anthropic workload identity tokens on rust * feat(messages): apply prompt caching injection on rust * feat(messages): serve the easy python-only detectors on rust * feat(messages): retry rejected thinking signatures on rust * feat(messages): serve response cache reads on rust * feat(messages): retry rejected strict tools on rust * refactor(messages): delete the python path for anthropic messages Provider anthropic on /v1/messages is now RUST_REQUIRED, so it runs on Rust without LITELLM_RUST or litellm.rust(True) and never falls back to Python. A call the rules would send to Python (a MESSAGES_PYTHON_FEATURES feature or a detector that raises) gets UnsupportedParamsError naming the feature, raised from one place in litellm/messages/dispatch.py. Other providers keep the Python fallback. A call whose arguments the handler signature rejects still gets Python's TypeError. Deleted Python code that only served anthropic Messages: - workload identity in AnthropicMessagesConfig.validate/avalidate (subclasses were never eligible; the Rust host mints the token itself) - the anthropic Foundry beta-row and affinity-base branch in async_anthropic_messages_handler, and the "forward" mode of the beta filter that only that branch used The proxy MCP tools detector now uses the light tool precheck so a function tool never imports the gateway handler, since a raising detector now refuses. * feat(messages): forward message threads on rust The Rust path for anthropic /v1/messages dropped the thread body field, so a create or continue turn reached Anthropic without its thread or the message-threads-2026-08-12 beta. It now matches the Python handler from #50976: - the host forwards thread with the other body params - a non-null thread adds the message-threads beta - a well-formed continue turn keeps advisor blocks in history and derives no compaction, per-turn-control or tool-change beta from its messages - the recorded httpx_response carries the request that was sent, for JSON and streamed calls, so the proxy's thread ownership hook resends the create turn's betas on a continue as it does for Python Upstream error.details already reached the exception message unchanged, so thread_not_found relays and keeps the deployment out of cooldown on both paths. Parity cases cover create, continue, streaming, sync, retries, thread_not_found and the ownership hook round trip. * feat(messages): serve mock responses on rust * feat(messages): send caller clients and unsupported tls through python on rust * feat(messages): run pre-request hooks on rust * feat(messages): run agentic loop hooks on rust * feat(messages): run the safeguards classifier on rust * feat(messages): run registered interceptors on rust * fix(messages): match python stream cache gate, logging and interceptor order on rust Store a streamed Messages response in the cache only when Python would: a message_stop event line, no provider error frame and at least one frame with output. Leave dict results untouched by mark_rust_response so hidden params never reach result dicts, logging objects or hook digests. Ask interceptors whether they handle the call after the cache lookup, not at decode time, so a cache hit never calls can_handle. Emit the track_llm_api_timing service event for the provider exchange. Bill and log a stream superseded by a router fallback as the failure of its error frame, with the partial usage warning. Hand litellm_logging_obj to deployment pre-call hooks as @client does, so the autorouter baseline cache sees retried requests. * fix(messages): forward python-only json values to anthropic as python does Accept any max_tokens, sampling scalar and messages collection that Python main accepts, and carry values serde_json cannot hold (ints beyond 64 bits, NaN and infinities, strings with lone surrogates) as validated raw JSON leaves written back verbatim on the wire, so the upstream body matches Python's json.dumps * fix(messages): match python's body reads, error bodies, encoded heads and tls warnings on rust * fix(messages): log python's plain provider header names on rust streams * fix(messages): match python's request params, logged params and validation on rust Read a non-bool stream by Python truthiness and send it as given, log the caller's stream flag as Python's handler does, raise the raw Python errors main raises for malformed messages, tools, model_info and hook-set litellm_params or providers, read the tool custom-field drop from drops_unsupported_params, log the caller's nested fields that additional_drop_params strips and the Python values of raw JSON leaves, keep a missing stop_reason absent, and map provider-config rejections (missing max_tokens, unknown reasoning effort) as Python's HTTP handler maps them, with Python's text, effort order, model and provider * fix(messages): match python's tls curve names, redirects, auth token, header encoding and timeouts on rust * fix(messages): match python on hooked streams, empty cache writes and mcp cache hits * fix(messages): join already buffered stream reads and trim the env api key as python does on rust * fix(messages): deliver stream bytes read before a cut, keep wide reply numbers and follow python's redirect limit on rust A stream cut by a transport error now hands the caller the bytes already buffered before raising, as aiohttp does for a caller that reads them before the event loop sees the cut, so the frames and their usage reach the caller and the failure log. Provider replies keep integers wider than 64 bits exactly and read numbers that overflow a float as infinity, through the response, the response cache and the logged result. The provider client follows up to 20 redirects as httpx does, and running out raises Python's text. A provider_specific_header value the ascii codec cannot encode fails before the send like a caller extra header. The deployment success hook gets the stream as AnthropicMessagesStreamingResponse, as Python's handler hands it. * feat(messages): serve proxy mcp tools on rust * test(messages): cover hook, agentic and interceptor edge cases in the parity suite * fix(messages): keep provider raw-json marked strings, drop non-string tool_result texts and key clamped streams as python does on rust * fix(messages): match python's api key encoding, malformed api_base and string provider header errors on rust A header value Python's client cannot encode as ASCII now fails before the send whichever header carries it, so a caller api_key with a non-ASCII character raises Python's InternalServerError instead of a raw UnicodeEncodeError after the reply. A caller api_base without an http(s) scheme, or with whitespace, control or non-ASCII characters, is sent through Python's client at the URL Python builds, so its relative-URL and invalid-URL errors match, and a base urlsplit cannot split raises Python's bare ValueError before the send. A provider_specific_header of the wrong shape is scoped through Python's ProviderSpecificHeaderUtils first, so a plain string raises Python's AttributeError. * fix(messages): keep integer temperature and top_p digits a float cannot hold and raise python's typeerror for a string max_tokens while capping the thinking budget on rust temperature and top_p now deserialize through recognized::exact_float, which keeps an integer where float(n) != n (beyond 2**53) as an unrecognized value with its exact digits, so it is sent upstream and logged as the integer python forwards, and comparisons still go through scalar.rs cap_thinking_budget_to_max_tokens compares max_tokens <= 1024 first, so a max_tokens python cannot order against an int (a string, list or dict) raises TypeError "'<=' not supported between instances of 'str' and 'int'" in both the reasoning effort translation and the adaptive effort reshape. rust now returns ErrorDetail::UnorderableTypes there, the bridge raises it as a python TypeError marked as a late rejection, and map_failure hands it to the caller unmapped * test(messages): format the non-string tool_result parity helper * fix(messages): flag and log the call before python's pre-send raises on rust The bridge now hands settings() and map_failure() whether the call is async. A provider_specific_header of the wrong shape sets the aanthropic_messages flag before re-raising, as Python's handler has already flagged the call. A base urlsplit cannot split also logs the resolved model, provider and optional params first, as Python raises it only after its logger is updated. The TypeError raised while capping the thinking budget for a string max_tokens sets the flag too * feat(messages): serve the web search short-circuit override on rust through the intercept host call * fix(messages): give the web search entry kwargs its mutable-ok reason * test(messages): drop the body hook gap the rust short-circuit closes * fix(messages): name the partially unknown handler types on the short-circuit host call --- litellm-rust/Cargo.lock | 7 + .../crates/cache-response/src/service.rs | 4 + .../python_contract.json | 28 +- .../callbacks-legacy-python/src/adapter.rs | 392 +- .../callbacks-legacy-python/src/callbacks.rs | 22 +- .../crates/callbacks-legacy-python/src/lib.rs | 2 +- .../callbacks-legacy-python/src/logger.rs | 6 +- .../callbacks-legacy-python/src/mapping.rs | 3 + .../callbacks-legacy-python/src/python.rs | 6 + .../src/test_support.rs | 14 +- .../crates/gateway-inference/tests/caching.rs | 7 +- litellm-rust/crates/host-native/src/driver.rs | 7 + litellm-rust/crates/host-python/Cargo.toml | 2 +- .../crates/host-python/src/binding.rs | 4 +- litellm-rust/crates/host-python/src/driver.rs | 12 +- .../host-python/src/hooks/chain/dispatch.rs | 9 + litellm-rust/crates/host-python/src/lib.rs | 3 +- .../crates/host-python/src/marshal.rs | 182 +- .../crates/host-python/tests/support/mod.rs | 3 + litellm-rust/crates/host/Cargo.toml | 1 + litellm-rust/crates/host/src/interceptors.rs | 28 + litellm-rust/crates/host/src/lib.rs | 1 + litellm-rust/crates/host/src/lifecycle.rs | 15 +- .../crates/host/src/machine/context.rs | 10 +- litellm-rust/crates/host/src/protocol.rs | 6 +- litellm-rust/crates/host/src/raw_json.rs | 362 ++ litellm-rust/crates/http/src/config.rs | 20 +- litellm-rust/crates/http/src/error.rs | 4 + litellm-rust/crates/http/src/outbound.rs | 15 +- litellm-rust/crates/http/src/pool.rs | 28 +- litellm-rust/crates/http/src/tls.rs | 8 +- .../crates/inference-chat/tests/caching.rs | 1 + .../crates/inference-messages/Cargo.toml | 4 + .../inference-messages/src/constants.rs | 3 + .../crates/inference-messages/src/encoding.rs | 437 ++ .../crates/inference-messages/src/handler.rs | 968 ++++- .../inference-messages/src/host_answer.rs | 11 + .../crates/inference-messages/src/lib.rs | 141 +- .../inference-messages/src/pre_request.rs | 7 + .../crates/inference-messages/src/reads.rs | 394 ++ .../crates/inference-messages/src/route.rs | 99 + .../inference-messages/src/safeguards.rs | 248 ++ .../inference-messages/src/transport.rs | 59 + .../crates/inference-messages/src/types.rs | 43 +- .../src/workload_identity.rs | 11 + .../inference-messages/tests/caching.rs | 132 + .../tests/messages/gateway.rs | 99 + .../inference-messages/tests/messages/host.rs | 199 +- .../inference-messages/tests/messages/main.rs | 4 + .../tests/messages/pre_request.rs | 141 + .../tests/messages/response.rs | 119 +- .../tests/messages/retry.rs | 364 ++ .../tests/messages/stream.rs | 4 +- .../tests/messages/transport.rs | 256 ++ .../inference-messages/tests/support/mod.rs | 10 + .../inference-ocr/tests/ocr/lifecycle.rs | 1 + .../crates/inference-ocr/tests/ocr/machine.rs | 6 + litellm-rust/crates/inference/AGENTS.md | 1 + litellm-rust/crates/inference/Cargo.toml | 2 +- litellm-rust/crates/inference/src/caching.rs | 17 +- litellm-rust/crates/inference/src/context.rs | 13 +- litellm-rust/crates/inference/src/error.rs | 19 +- litellm-rust/crates/inference/src/outbound.rs | 17 +- .../src/formats/messages/request.rs | 64 +- .../src/providers/anthropic/beta.rs | 5 + .../crates/llms-types/src/recognized.rs | 149 + litellm-rust/crates/llms/Cargo.toml | 2 + .../crates/llms/src/anthropic/beta_headers.rs | 12 +- .../crates/llms/src/anthropic/common_utils.rs | 375 +- .../llms/src/anthropic/messages/handler.rs | 44 +- .../crates/llms/src/anthropic/messages/mod.rs | 1 + .../llms/src/anthropic/messages/scalar.rs | 113 + .../llms/src/anthropic/messages/thinking.rs | 210 +- .../src/anthropic/messages/transformation.rs | 660 ++- .../src/azure_ai/messages/transformation.rs | 25 +- .../llms/src/base_llm/messages/context.rs | 8 + .../crates/llms/src/base_llm/messages/mod.rs | 1 + .../src/base_llm/messages/normalization.rs | 26 +- .../llms/src/base_llm/messages/tools.rs | 194 + .../src/base_llm/messages/transformation.rs | 71 +- .../messages/transformation.rs | 27 +- .../anthropic_claude_transformation.rs | 18 +- .../bedrock_mantle/messages/transformation.rs | 31 +- .../src/deepseek/messages/transformation.rs | 27 +- .../src/edenai/messages/transformation.rs | 18 + litellm-rust/crates/llms/src/error.rs | 12 + .../src/minimax/messages/transformation.rs | 28 +- .../src/openrouter/messages/transformation.rs | 13 + .../src/tencent/messages/transformation.rs | 28 +- .../src/vertex_ai/messages/transformation.rs | 26 +- .../anthropic_messages_http_error_retry.json | 623 +++ .../llms/tests/messages_http_error_retry.rs | 138 + .../crates/python-bridge/src/cache/AGENTS.md | 2 +- .../crates/python-bridge/src/cache/mod.rs | 2 +- .../python-bridge/src/cache/python/AGENTS.md | 2 +- .../python-bridge/src/cache/python/host.rs | 77 +- .../python-bridge/src/cache/python/service.rs | 23 +- .../python-bridge/src/cache/selection.rs | 46 +- .../crates/python-bridge/src/errors.rs | 117 +- litellm-rust/crates/python-bridge/src/http.rs | 47 +- litellm-rust/crates/python-bridge/src/lib.rs | 3 + .../crates/python-bridge/src/marshal.rs | 15 +- .../src/routes/messages/agentic.rs | 172 + .../src/routes/messages/answer.rs | 29 + .../python-bridge/src/routes/messages/host.rs | 758 +++- .../python-bridge/src/routes/messages/mod.rs | 80 +- .../src/routes/messages/modify_params.rs | 152 + .../src/routes/messages/pre_request.rs | 35 + .../src/routes/messages/safeguards.rs | 183 + .../src/routes/messages/transport.rs | 275 ++ .../src/routes/messages/workload_identity.rs | 74 + .../crates/python-bridge/src/routes/mod.rs | 35 +- litellm-rust/crates/python-compat/src/lib.rs | 2 + .../crates/python-compat/src/reprlib.rs | 179 + litellm-rust/crates/router/src/lib.rs | 25 +- litellm/__init__.py | 4 +- litellm/anthropic_beta_headers_manager.py | 12 +- litellm/caching/caching_handler.py | 265 +- litellm/litellm_core_utils/core_helpers.py | 14 +- .../llms/anthropic/messages/native/handler.py | 52 +- .../messages/native/transformation.py | 23 +- litellm/llms/custom_httpx/http_handler.py | 4 + litellm/llms/custom_httpx/llm_http_handler.py | 16 +- .../llms/dashscope/messages/transformation.py | 4 +- litellm/messages/dispatch.py | 108 +- litellm/rust_bridge/AGENTS.md | 5 +- litellm/rust_bridge/_native.pyi | 2 + litellm/rust_bridge/catalog.py | 43 +- litellm/rust_bridge/detectors.py | 32 +- litellm/rust_bridge/execution.py | 6 +- litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/cache.py | 200 +- .../host/callbacks_legacy_python.py | 200 +- litellm/rust_bridge/host/http_transport.py | 25 + litellm/rust_bridge/host/routes/messages.py | 1255 +++++- litellm/rust_bridge/messages/detectors.py | 114 - litellm/rust_bridge/messages/entrypoints.py | 35 +- litellm/rust_bridge/response_metadata.py | 7 +- litellm/rust_bridge/runtime.py | 2 +- litellm/utils.py | 99 +- .../sdk/messages/test_native_parity.py | 3523 ++++++++++++++++- .../cache/test_python_cache.py | 140 +- tests/test_litellm_rust/messages/data.py | 4 - .../messages/test_body_hooks.py | 36 +- .../messages/test_dispatch.py | 635 ++- .../messages/test_workload_identity.py | 219 + .../support/recording_server.py | 40 +- tests/test_litellm_rust/support/served.py | 10 +- tests/test_litellm_rust/test_dispatch.py | 2 +- .../litellm_core_utils/test_core_helpers.py | 27 + .../test_litellm_logging.py | 3 +- .../test_anthropic_chat_transformation.py | 3 +- .../native/test_anthropic_messages_handler.py | 9 +- .../native/test_anthropic_messages_router.py | 16 +- .../test_anthropic_messages_transformation.py | 51 +- .../messages/native/test_response_cache.py | 50 +- .../anthropic/test_anthropic_common_utils.py | 38 - .../test_anthropic_prompt_cache_prediction.py | 2 + .../test_messages_http_error_retry_cases.py | 188 + .../test_anthropic_claude_transformation.py | 3 + .../custom_httpx/test_llm_http_handler.py | 33 +- .../hooks/test_autorouter_baseline_cache.py | 7 + .../unit/responses/mcp/test_tool_precheck.py | 21 +- .../test_context_compaction.py | 16 + .../router_strategy/test_complexity_router.py | 15 + .../test_fallback_event_handlers.py | 16 + .../rust_bridge/host/routes/test_messages.py | 1314 +++++- tests/unit/rust_bridge/host/test_cache.py | 119 +- .../host/test_callbacks_legacy_python.py | 155 +- tests/unit/rust_bridge/test_catalog.py | 25 +- tests/unit/rust_bridge/test_detectors.py | 161 +- tests/unit/rust_bridge/test_runtime.py | 56 +- tests/unit/test_private_usage_aliases.py | 2130 +++++----- tests/unit/test_router/test_router.py | 2 + .../test_router/test_router_helper_utils.py | 2 + 175 files changed, 19176 insertions(+), 2747 deletions(-) create mode 100644 litellm-rust/crates/host/src/raw_json.rs create mode 100644 litellm-rust/crates/inference-messages/src/encoding.rs create mode 100644 litellm-rust/crates/inference-messages/src/host_answer.rs create mode 100644 litellm-rust/crates/inference-messages/src/pre_request.rs create mode 100644 litellm-rust/crates/inference-messages/src/reads.rs create mode 100644 litellm-rust/crates/inference-messages/src/safeguards.rs create mode 100644 litellm-rust/crates/inference-messages/src/transport.rs create mode 100644 litellm-rust/crates/inference-messages/src/workload_identity.rs create mode 100644 litellm-rust/crates/inference-messages/tests/messages/gateway.rs create mode 100644 litellm-rust/crates/inference-messages/tests/messages/pre_request.rs create mode 100644 litellm-rust/crates/inference-messages/tests/messages/retry.rs create mode 100644 litellm-rust/crates/inference-messages/tests/messages/transport.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/messages/scalar.rs create mode 100644 litellm-rust/crates/llms/src/base_llm/messages/tools.rs create mode 100644 litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json create mode 100644 litellm-rust/crates/llms/tests/messages_http_error_retry.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/agentic.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/answer.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/pre_request.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/safeguards.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/transport.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs create mode 100644 litellm-rust/crates/python-compat/src/reprlib.rs delete mode 100644 litellm/rust_bridge/messages/detectors.py create mode 100644 tests/test_litellm_rust/messages/test_workload_identity.py create mode 100644 tests/unit/llms/anthropic/test_messages_http_error_retry_cases.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed658e074ee..2cdc9ee91c8 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4013,6 +4013,7 @@ dependencies = [ "litellm-coroutine", "litellm-secret-value", "rstest", + "serde", "serde_json", "thiserror 2.0.19", "tokio", @@ -4155,9 +4156,12 @@ name = "litellm-inference-messages" version = "0.1.0" dependencies = [ "bytes", + "flate2", "futures-util", + "http 1.4.2", "litellm-auth", "litellm-auth-aws", + "litellm-cache", "litellm-cache-memory", "litellm-cache-response", "litellm-core-utils", @@ -4178,6 +4182,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "thiserror 2.0.19", "tokio", "tracing", "wiremock", @@ -4292,6 +4297,7 @@ dependencies = [ "base64 0.22.1", "bytes", "data-url", + "fancy-regex 0.19.2", "futures-util", "litellm-auth", "litellm-auth-aws", @@ -4307,6 +4313,7 @@ dependencies = [ "litellm-router-types", "litellm-secret-value", "litellm-secrets", + "litellm-tracing", "reqwest", "rstest", "serde", diff --git a/litellm-rust/crates/cache-response/src/service.rs b/litellm-rust/crates/cache-response/src/service.rs index 0bdf948ec48..ed4dccdc355 100644 --- a/litellm-rust/crates/cache-response/src/service.rs +++ b/litellm-rust/crates/cache-response/src/service.rs @@ -27,6 +27,10 @@ impl Default for ResponseCacheConfig { pub trait ResponseCacheService: Send + Sync { fn config(&self) -> &ResponseCacheConfig; + fn keyed_by_caller(&self) -> bool { + false + } + fn lookup<'a>( &'a self, request: &'a ResponseCacheRequest, diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index 8a212223523..a96e9f04bdc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -11,7 +11,8 @@ "logger", "kwargs", "start_time", - "end_time" + "end_time", + "cache_hit" ], "update_logging": [ "logger", @@ -41,13 +42,24 @@ "logger", "status", "headers", - "body" + "body", + "sent" + ], + "retried_request": [ + "logger", + "logged_body", + "body", + "reasons" ], "record_api_window": [ "logger", "start", "end" ], + "times_provider_exchange": [ + "client", + "timeout" + ], "defers_async_logging": [ "logger" ], @@ -89,6 +101,12 @@ "start", "end" ], + "cache_hit_success": [ + "logger", + "response", + "start", + "end" + ], "enqueue_logging": [ "coroutine" ], @@ -113,14 +131,16 @@ ], "stream_opened": [ "logger", - "head" + "head", + "keeps_requested" ], "stamp_stream_head": [ "head", "logger", "kwargs", "start_time", - "end_time" + "end_time", + "cache_hit" ], "stream_success": [ "logger", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index f68541c5c7b..71ca89b0790 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -6,10 +6,10 @@ use crate::LoggingOperation; use litellm_host_python::PythonOwned; use litellm_host::{ - interceptors::{RawResponse, RequestContext, WireRequest}, + interceptors::{RawResponse, RequestContext, RetriedRequest, RetryReason, WireRequest}, lifecycle::{FailureOrigin, Timing, epoch_seconds}, }; -use litellm_host_python::{Error as HostError, HookStep, from_py, to_py}; +use litellm_host_python::{Error as HostError, HookStep, from_py, json_from_py, json_to_py, to_py}; use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, @@ -19,7 +19,7 @@ use pyo3::{ use serde_json::Value; use crate::{ - DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, + DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, SentRequest, deferred::{PendingLogging, PendingSuccess}, finalize, is_internal_call, python::{Logging, Streaming}, @@ -41,6 +41,7 @@ struct DeliveredStream { struct LoggedRequest { body: Py, headers: Py, + url: String, context: RequestContext, } @@ -59,6 +60,16 @@ pub struct LegacyLogging { cache_key: Option, /// When the provider request went out, until its window is recorded (Messages only). api_started: Option, + host_answered: bool, +} + +fn retry_reason(reason: &RetryReason) -> Value { + match reason { + RetryReason::ThinkingSignature => serde_json::json!({"kind": "thinking_signature"}), + RetryReason::StrictTools { dropped } => { + serde_json::json!({"kind": "strict_tools", "dropped": dropped}) + } + } } fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult> { @@ -90,6 +101,7 @@ impl LegacyLogging { internal: false, cache_key: None, api_started: None, + host_answered: false, } } @@ -171,6 +183,14 @@ impl LegacyLogging { Ok(context) } + fn times_provider_exchange(&self, py: Python<'_>) -> PyResult { + let client = self.call.lookup(py, "client")?; + let timeout = self.call.lookup(py, "timeout")?; + Logging::TimesProviderExchange + .call(py, (client, timeout))? + .extract() + } + /// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync /// client is untimed. fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> { @@ -264,6 +284,7 @@ impl LegacyLogging { self.call.kwargs(), &self.start, &self.end, + self.messages_cache_hit(), )?; self.response .as_ref() @@ -271,6 +292,60 @@ impl LegacyLogging { .ok_or_else(|| HostError::MissingState.into()) } + fn messages_cache_hit(&self) -> bool { + self.operation == LoggingOperation::Messages && self.cache_key.is_some() + } + + fn flag_unsent_messages_call(&self, py: Python<'_>) -> PyResult<()> { + if self.operation != LoggingOperation::Messages + || self.request.is_some() + || self.cache_key.is_some() + || self.host_answered + { + return Ok(()); + } + let details = self.logger()?.object(py).getattr("model_call_details")?; + details + .call_method1("setdefault", ("litellm_params", PyDict::new(py)))? + .set_item("aanthropic_messages", self.asynchronous) + } + + fn unstreamed_stream_request(&self, py: Python<'_>) -> PyResult { + if self.operation != LoggingOperation::Messages + || self.stream.is_some() + || self.messages_cache_hit() + { + return Ok(false); + } + let Some(stream) = self.call.kwargs().bind(py).get_item("stream")? else { + return Ok(false); + }; + if !stream.is_exact_instance_of::() || !stream.is_truthy()? { + return Ok(false); + } + let logger = self.logger()?.object(py); + logger.setattr("stream", true)?; + logger + .getattr("model_call_details")? + .set_item("stream", true)?; + Ok(true) + } + + fn dispatch_cache_hit(&self, py: Python<'_>) -> PyResult<()> { + let logger = self.logger()?; + let logged = Logging::CacheHitSuccess.call( + py, + (logger.object(py), &self.response, &self.start, &self.end), + ); + match logged { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(logger.object(py))); + Ok(()) + } + result => result.map(|_| ()), + } + } + fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> { match self.try_dispatch_success(py) { Err(error) if error.is_instance_of::(py) => { @@ -489,6 +564,10 @@ impl LegacyLogging { self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); if self.runs_deployment_hooks() { + self.call + .kwargs() + .bind(py) + .set_item("litellm_logging_obj", self.logger()?.object(py))?; return Ok(HookStep::Await( DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type())?, Self::resume_begin, @@ -512,9 +591,13 @@ impl LegacyLogging { context, self.logged_request_params(context).as_ref(), )?; - let body = to_py(py, &wire.body)? - .into_bound(py) - .cast_into::()?; + let lossless = self.operation == LoggingOperation::Messages; + let body = if lossless { + json_to_py(py, &wire.body)? + } else { + to_py(py, &wire.body)?.into_bound(py) + } + .cast_into::()?; for (name, sent) in wire.body.as_object().into_iter().flatten() { if let Some(value) = self.caller_value(py, name, sent)? { body.set_item(name, value)?; @@ -534,6 +617,7 @@ impl LegacyLogging { self.request = Some(LoggedRequest { body: body.clone().unbind(), headers: headers.clone().unbind(), + url: wire.url.clone(), context: context.clone(), }); let input = self.logged_input(py, &body)?; @@ -545,7 +629,10 @@ impl LegacyLogging { &headers, &wire.url, )?; - if self.operation == LoggingOperation::Messages && self.asynchronous { + if self.operation == LoggingOperation::Messages + && self.asynchronous + && self.times_provider_exchange(py)? + { self.api_started = Some(epoch_seconds()); } let headers = headers @@ -553,7 +640,11 @@ impl LegacyLogging { .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; Ok(HookStep::Ready(Box::new(WireRequest { - body: from_py(&body)?, + body: if lossless { + json_from_py(body.as_any())? + } else { + from_py(&body)? + }, headers, ..*wire }))) @@ -567,7 +658,11 @@ impl LegacyLogging { ) -> PyResult>> { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.runs_deployment_hooks() { + self.flag_unsent_messages_call(py)?; + if self.unstreamed_stream_request(py)? { + return self.finalize(py); + } + if self.runs_deployment_hooks() && !self.messages_cache_hit() { return Ok(HookStep::Await( DeploymentHooks::after_success( py, @@ -588,7 +683,21 @@ impl LegacyLogging { ) -> PyResult> { use litellm_host::interceptors::{Cost, ResultSource}; + self.cache_key = match &facts.source { + ResultSource::Provider => None, + ResultSource::Cache { key } => Some(key.clone()), + ResultSource::Host => { + self.host_answered = true; + return Ok(HookStep::Ready(())); + } + }; let logger = self.logger()?.object(py); + if self.messages_cache_hit() { + logger + .getattr("model_call_details")? + .set_item("cache_hit", true)?; + return Ok(HookStep::Ready(())); + } let params = logger .getattr("litellm_params")? .cast_into::()? @@ -616,10 +725,6 @@ impl LegacyLogging { }; logger.setattr("_provider_reported_cost", reported_cost)?; details.call_method1("pop", ("response_cost", py.None()))?; - self.cache_key = match &facts.source { - ResultSource::Provider => None, - ResultSource::Cache { key } => Some(key.clone()), - }; if self.cache_key.is_some() { details.set_item("cache_hit", true)?; } else { @@ -644,8 +749,13 @@ impl LegacyLogging { // response, and logging replays it through `AnthropicConfig.transform_response`. self.close_api_window(py, epoch_seconds())?; if let Some(head) = &raw.head { + let sent = self.request.as_ref().map(|request| SentRequest { + url: &request.url, + headers: &request.headers, + body: &request.body, + }); self.logger()? - .record_response(py, head.status, &head.headers, &raw.body)?; + .record_response(py, head.status, &head.headers, &raw.body, sent)?; } return Ok(HookStep::Ready(())); } @@ -664,6 +774,40 @@ impl LegacyLogging { Ok(HookStep::Ready(())) } + pub(crate) fn request_retried( + &mut self, + py: Python<'_>, + retry: &RetriedRequest, + ) -> PyResult> { + let Some(request) = &self.request else { + return Ok(HookStep::Ready(())); + }; + let reasons: Vec = retry.reasons.iter().map(retry_reason).collect(); + let retried = Logging::RetriedRequest.call( + py, + ( + self.logger()?.object(py), + request.body.bind(py), + to_py(py, &retry.body)?, + to_py(py, &reasons)?, + ), + )?; + let now = epoch_seconds(); + if self.api_started.is_some() { + self.close_api_window(py, now)?; + self.api_started = Some(now); + } + Ok(HookStep::Await(retried.unbind(), Self::resume_retried)) + } + + fn resume_retried( + &mut self, + _py: Python<'_>, + result: PyResult>, + ) -> PyResult> { + result.map(|_| HookStep::Ready(())) + } + pub(crate) fn succeeded( &mut self, py: Python<'_>, @@ -674,6 +818,8 @@ impl LegacyLogging { self.response = Some(response.clone_ref(py)); match &self.stream { Some(stream) => self.stream_success(py, stream, false)?, + None if self.messages_cache_hit() => self.dispatch_cache_hit(py)?, + None if self.unstreamed_stream_request(py)? => {} None => self.dispatch_success(py)?, } Ok(HookStep::Ready(())) @@ -708,7 +854,11 @@ impl LegacyLogging { if self.stream_billing().is_none() { return Err(HostError::MissingState.into()); } - if let Some(key) = &self.cache_key { + if let Some(key) = self + .cache_key + .as_ref() + .filter(|_| !self.messages_cache_hit()) + { head.bind(py).set_item("cache_key", key)?; head.bind(py).set_item("cache_hit", true)?; } @@ -723,10 +873,18 @@ impl LegacyLogging { self.call.kwargs(), &self.start, datetime(py, epoch_seconds())?, + self.messages_cache_hit(), ), )?; } - Streaming::Opened.call(py, (self.logger()?.object(py), head.bind(py)))?; + Streaming::Opened.call( + py, + ( + self.logger()?.object(py), + head.bind(py), + self.operation == LoggingOperation::Messages, + ), + )?; self.stream = Some(DeliveredStream { chunks: PyList::empty(py).unbind(), first_chunk: None, @@ -747,7 +905,7 @@ impl LegacyLogging { // A caller that goes away mid-stream cancels the call before it ends: Python's Messages // stream iterator still bills the chunks it read, straight away rather than deferred. fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> { - if !self.asynchronous || self.end.is_some() { + if !self.asynchronous || self.end.is_some() || self.messages_cache_hit() { return Ok(()); } let Some(stream) = self.stream.take() else { @@ -918,6 +1076,24 @@ if asynchronous: }); } + #[rstest::rstest] + fn the_deployment_pre_call_hook_receives_the_logger_as_client_injects_it() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +kwargs = {'logger': logger} +logger.hooks['pre'] = lambda value: seen.append(value.get('litellm_logging_obj')) or 'awaitable' +seen = [] +", + ); + let (_, step) = begin(py, &locals, true); + assert!(awaits_deployment_hook(&step)); + run(py, &locals, c"assert seen == [logger]"); + }); + } + #[rstest::rstest] fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { Python::initialize(); @@ -1038,6 +1214,82 @@ assert finalized is replacement }); } + #[rstest] + #[case::unary(false)] + #[case::stream(true)] + fn a_host_answer_keeps_the_logger_as_set_up_and_a_streamed_one_is_finalized_unstreamed( + #[case] stream: bool, + ) { + use litellm_host::interceptors::{Cost, ExecutionFacts, ProviderIdentity, ResultSource}; + + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +kwargs = {'logger': logger} +response = object() +logger.hooks = {'pre': lambda kwargs: kwargs} +", + ); + local(&locals, "kwargs").set_item("stream", stream).unwrap(); + let (mut logging, step) = begin(py, &locals, true); + logging.operation = crate::LoggingOperation::Messages; + resume( + &mut logging, + step, + py, + Ok(local(&locals, "kwargs").unbind()), + ) + .unwrap(); + let facts = ExecutionFacts { + provider: ProviderIdentity { + model: "claude-sonnet-4-5".into(), + provider: "anthropic".into(), + }, + cost: Cost::Deferred, + source: ResultSource::Host, + }; + assert!(matches!( + logging.result_ready(py, &facts).unwrap(), + HookStep::Ready(()) + )); + let response = local(&locals, "response").unbind(); + let step = logging + .transform_response(py, response.clone_ref(py), TIMING) + .unwrap(); + assert_eq!(awaits_deployment_hook(&step), !stream); + if !stream { + return; + } + let HookStep::Ready(returned) = step else { + panic!("a host-answered stream is returned as is"); + }; + assert!(returned.bind(py).is(response.bind(py))); + logging + .on_event( + py, + PythonCallEvent::Succeeded { + timing: TIMING, + response: &response, + }, + ) + .unwrap(); + run( + py, + &locals, + c" +names = logger.names() +assert ('finalize', response) in logger.calls, names +assert 'update_logging' not in names, names +assert not [name for name in names if 'success' in name], names +assert logger.stream is True +assert logger.model_call_details['stream'] is True +", + ); + }); + } + #[rstest] #[case::pre_call(false)] #[case::post_call(true)] @@ -1129,7 +1381,9 @@ mod payload_tests { use std::ffi::CStr; use litellm_host::hooks::CallHooks; - use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; + use litellm_host::interceptors::{ + RawResponse, RequestContext, RetriedRequest, RetryReason, WireRequest, + }; use litellm_host::lifecycle::ExecutionEvent; use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; use litellm_secret_value::SecretValue; @@ -1296,6 +1550,61 @@ assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, json!({"type": "document_url", "document_url": source}) } + #[test] + fn a_retried_request_hands_the_logged_body_the_retried_body_and_the_reasons_to_logging() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let context = RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params: json!({}), + secret_fields: Vec::new(), + api_key: None, + }; + let wire = WireRequest { + url: "https://provider.invalid/v1/messages".into(), + headers: Vec::new(), + body: json!({"messages": [{"role": "user", "content": "hi"}], "thinking": {"type": "enabled"}}), + }; + logging + .before_provider_request(py, Box::new(wire.clone()), &context) + .unwrap(); + let retried = RetriedRequest { + reasons: vec![ + RetryReason::ThinkingSignature, + RetryReason::StrictTools { dropped: 2 }, + ], + body: json!({"messages": [{"role": "user", "content": "hi"}]}), + }; + + let step = logging + .on_event( + py, + PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { + retry: retried, + }), + ) + .unwrap(); + + assert!(matches!(step, HookStep::Await(_, _))); + run( + py, + &locals, + c" +[(logged, body, reasons)] = [value for name, value in logger.calls if name == 'retried_request'] +assert logged is logger.pre['complete_input_dict'], logged +assert body == {'messages': [{'role': 'user', 'content': 'hi'}]}, body +assert reasons == [{'kind': 'thinking_signature'}, {'kind': 'strict_tools', 'dropped': 2}], reasons +", + ); + }); + } + fn before_provider_request(script: &CStr, body: Value) -> WireRequest { before_send_with_secrets(script, json!({}), body, &[]) } @@ -2018,6 +2327,53 @@ mod terminal_tests { .unwrap() } + #[rstest] + #[case::requested_stream( + c"kwargs = {'logger': logger, 'stream': True}", + true, + &["finalize"] + )] + #[case::unstreamed(c"kwargs = {'logger': logger, 'stream': False}", false, &["success_hook"])] + #[case::stream_flag_not_true(c"kwargs = {'logger': logger, 'stream': 'true'}", false, &["success_hook"])] + fn a_whole_messages_result_for_a_requested_stream_is_finalized_without_success_logging( + #[case] script: &CStr, + #[case] handed_back: bool, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let mut logging = LegacyLogging { + operation: crate::LoggingOperation::Messages, + ..logged(py, &locals, true) + }; + let step = logging + .transform_public_response(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + assert_eq!(matches!(step, HookStep::Ready(_)), handed_back); + if handed_back { + assert!(matches!( + succeed(py, &locals, &mut logging), + HookStep::Ready(()) + )); + } + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + if handed_back { + run( + py, + &locals, + c"assert logger.stream is True and logger.model_call_details['stream'] is True", + ); + } + }); + } + #[rstest] #[case::sync_listened(false, c"", &["submit"])] #[case::async_listened( diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 093e86c43dc..9eed008910c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -3,13 +3,19 @@ //! duplication. All of it expires with the legacy callback contract. use litellm_host::interceptors::{RequestContext, WireRequest}; -use litellm_host_python::to_py; +use litellm_host_python::json_to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; use serde_json::Value; use crate::logger::PythonLogger; use crate::python::{Logging, Wrapper}; +pub struct SentRequest<'a> { + pub url: &'a str, + pub headers: &'a Py, + pub body: &'a Py, +} + pub trait LegacyCallbacks { /// `Logging.update_from_kwargs`: what the logger is told about the request it is /// about to see, with consumed credentials redacted. @@ -41,6 +47,7 @@ pub trait LegacyCallbacks { status: u16, headers: &[(String, String)], body: &str, + request: Option>, ) -> PyResult<()>; /// Records one provider HTTP exchange's window, as Python's timed `AsyncHTTPHandler.post` @@ -117,9 +124,7 @@ impl LegacyCallbacks for PythonLogger { let redacted_kwargs = redact(py, kwargs, &secret_fields)?; let optional_params = redact( py, - &to_py(py, &context.optional_params)? - .into_bound(py) - .cast_into::()?, + &json_to_py(py, &context.optional_params)?.cast_into::()?, &secret_fields, )?; let params = PyDict::new(py); @@ -138,7 +143,7 @@ impl LegacyCallbacks for PythonLogger { } } for (name, value) in request_params.into_iter().flatten() { - params.set_item(name, to_py(py, value)?)?; + params.set_item(name, json_to_py(py, value)?)?; } Logging::Update.call( py, @@ -177,8 +182,13 @@ impl LegacyCallbacks for PythonLogger { status: u16, headers: &[(String, String)], body: &str, + request: Option>, ) -> PyResult<()> { - Logging::RecordResponse.call(py, (self.object(py), status, headers.to_vec(), body))?; + let request = request.map(|sent| (sent.url, sent.headers.bind(py), sent.body.bind(py))); + Logging::RecordResponse.call( + py, + (self.object(py), status, headers.to_vec(), body, request), + )?; Ok(()) } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index f453237b469..3a8c15ae54c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -17,7 +17,7 @@ mod mapping; mod python; pub use adapter::LegacyLogging; pub use call::PublicCall; -pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; +pub(crate) use callbacks::{LegacyCallbacks, SentRequest, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub use mapping::{CallBoundary, CallbackMapping, Dispatch, callback_mappings}; diff --git a/litellm-rust/crates/callbacks-legacy-python/src/logger.rs b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs index 38f3bf29828..3019d87a20b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/logger.rs @@ -77,8 +77,12 @@ pub fn finalize( kwargs: &Py, start: &Py, end: &Option>, + cache_hit: bool, ) -> PyResult<()> { - Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?; + Wrapper::Finalize.call( + py, + (response, logger.object(py), kwargs, start, end, cache_hit), + )?; Ok(()) } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs index cfa9e53aa7b..e321b93b03d 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/mapping.rs @@ -179,6 +179,9 @@ impl CallHooks for LegacyLogging { PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { (AFTER.invoke)(self, py, raw) } + PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => { + self.request_retried(py, &retry) + } PythonCallEvent::Succeeded { timing, response } => { (SUCCESS.invoke)(self, py, timing, response) } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 5b6faa3dbc2..58ddf7c708a 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -50,6 +50,10 @@ pub(crate) enum Logging { RecordResponse, #[strum(serialize = "record_api_window")] RecordApiWindow, + #[strum(serialize = "times_provider_exchange")] + TimesProviderExchange, + #[strum(serialize = "retried_request")] + RetriedRequest, #[strum(serialize = "defers_async_logging")] DefersAsync, #[strum(serialize = "claim_async_success_log")] @@ -62,6 +66,8 @@ pub(crate) enum Logging { SubmitSuccess, #[strum(serialize = "async_success_handler")] AsyncSuccessHandler, + #[strum(serialize = "cache_hit_success")] + CacheHitSuccess, #[strum(serialize = "mark_async_success_enqueued")] MarkAsyncSuccessEnqueued, #[strum(serialize = "enqueue_logging")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index ee16e0c0059..5b8b2bc4f96 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -53,7 +53,7 @@ def setup(call_type, args, kwargs, start, asynchronous): FAKES = { 'setup': setup, - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'finalize': lambda response, logger, kwargs, start, end, cache_hit: logger.record('finalize', response), 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( kwargs=kwargs, model=model, @@ -66,12 +66,17 @@ FAKES = { 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args ), - 'record_response': lambda logger, status, headers, body: logger.model_call_details.update( + 'record_response': lambda logger, status, headers, body, sent: logger.model_call_details.update( response_headers=dict(headers) ), + 'retried_request': lambda logger, logged_body, body, reasons: ( + logger.record('retried_request', (logged_body, body, reasons)), + 'awaitable', + )[1], 'record_api_window': lambda logger, start, end: logger.model_call_details.update( llm_api_duration_ms=(end - start) * 1000 ), + 'times_provider_exchange': lambda client, timeout: client is None, 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), 'claim_async_success_log': lambda logger: logger.claim_async_success_log(), 'defer_success': lambda logger, pending, response, start, end: ( @@ -86,6 +91,7 @@ FAKES = { )(error, ''.join(traceback.format_exception(error)), start, end), 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'cache_hit_success': lambda logger, response, start, end: logger.record('cache_hit_success', response), 'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None), 'enqueue_logging': lambda coroutine: coroutine.enqueue(), 'restore_context': lambda logger: logger.record('restore', None), @@ -96,8 +102,8 @@ FAKES = { 'success', response, call_type ), 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger, head: logger.record('stream_opened', head), - 'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head), + 'stream_opened': lambda logger, head, keeps_requested: logger.record('stream_opened', head), + 'stamp_stream_head': lambda head, logger, kwargs, start, end, cache_hit: logger.record('stamp_stream_head', head), 'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record( 'stream_success', list(chunks) ), diff --git a/litellm-rust/crates/gateway-inference/tests/caching.rs b/litellm-rust/crates/gateway-inference/tests/caching.rs index 143bc1ba6db..030c8a66ee3 100644 --- a/litellm-rust/crates/gateway-inference/tests/caching.rs +++ b/litellm-rust/crates/gateway-inference/tests/caching.rs @@ -36,7 +36,12 @@ async fn all_inference_endpoints_share_native_cache( } else { "message_stop" }; - let events = format!("event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n"); + let output = if is_responses { + "" + } else { + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n" + }; + let events = format!("{output}event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n"); let template = if stream { ResponseTemplate::new(200).set_body_raw(events.clone(), "text/event-stream") } else { diff --git a/litellm-rust/crates/host-native/src/driver.rs b/litellm-rust/crates/host-native/src/driver.rs index af3248848a2..bb1a7025e10 100644 --- a/litellm-rust/crates/host-native/src/driver.rs +++ b/litellm-rust/crates/host-native/src/driver.rs @@ -84,6 +84,13 @@ where reply.send(self.interceptors.after_provider_response(raw).await); Ok(()) } + HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { + retry, + reply, + }) => { + reply.send(self.interceptors.provider_request_retried(retry).await); + Ok(()) + } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.demand = Some(reply); return Ok(Boundary::Open(head)); diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index 0b0107d0548..1fa5fb38830 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -14,6 +14,7 @@ bytes.workspace = true futures-util.workspace = true parking_lot.workspace = true serde.workspace = true +serde_json.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } @@ -23,4 +24,3 @@ pythonize = "0.29.0" [dev-dependencies] rstest.workspace = true -serde_json.workspace = true diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index fe6173dd100..97d6be88767 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -29,8 +29,8 @@ pub trait PythonBinding: PythonOwned { /// What the stream carries at hand-off, as the caller's stream receives it. /// The unary response after the deployment transform, before success hooks run. - fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> { - Ok(()) + fn deliver_response(&mut self, _py: Python<'_>, response: Py) -> PyResult> { + Ok(response) } fn encode_stream_head( diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 052156cbc5d..ba0c174c875 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -310,6 +310,14 @@ where reply.send(Ok(())); Ok(()) } + HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { retry, reply }) => { + self.notify(PythonCallEvent::Execution( + ExecutionEvent::ProviderRequestRetried { retry }, + )) + .await?; + reply.send(Ok(())); + Ok(()) + } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.streaming = true; let head = { @@ -405,10 +413,10 @@ where }; let delivered = { let mut parts = self.parts(); - Python::attach(|py| parts.binding.deliver_response(py, response.bind(py))) + Python::attach(|py| parts.binding.deliver_response(py, response)) }; match delivered { - Ok(()) => self.succeeded(response).await, + Ok(response) => self.succeeded(response).await, Err(error) => self.failure(error, FailureOrigin::Host).await, } } diff --git a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs index 4939995cefa..26d5a5692d2 100644 --- a/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs +++ b/litellm-rust/crates/host-python/src/hooks/chain/dispatch.rs @@ -241,6 +241,9 @@ fn retain_event(py: Python<'_>, event: PythonCallEvent<'_>) -> OwnedEvent { CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw: raw.clone() }) } + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => { + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) + } CallEvent::Succeeded { timing, response } => CallEvent::Succeeded { timing, response: response.clone_ref(py), @@ -280,6 +283,12 @@ fn dispatch( py, CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }), ), + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => hooks.on_event( + py, + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { + retry: retry.clone(), + }), + ), CallEvent::Succeeded { timing, response } => hooks.on_event( py, CallEvent::Succeeded { diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index a9ebae427d9..1cd4d65e29e 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -31,7 +31,8 @@ pub use gil::{PythonContext, attach_blocking, release_count, release_gil}; pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension}; pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime}; pub use marshal::{ - Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py, + Pythonized, from_py, from_py_argument, json_from_py, json_loads, json_object_field, json_to_py, + panic_to_pyerr, to_py, }; pub use owned::PythonOwned; pub use runtime::{ diff --git a/litellm-rust/crates/host-python/src/marshal.rs b/litellm-rust/crates/host-python/src/marshal.rs index 465521cf9d2..a6468d38d83 100644 --- a/litellm-rust/crates/host-python/src/marshal.rs +++ b/litellm-rust/crates/host-python/src/marshal.rs @@ -4,9 +4,10 @@ use std::panic::{AssertUnwindSafe, catch_unwind}; use pyo3::exceptions::PyValueError; use pyo3::panic::PanicException; use pyo3::prelude::*; -use pyo3::types::PyBytes; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}; use serde::Serialize; use serde::de::DeserializeOwned; +use serde_json::{Map, Number, Value}; /// Converts a `#[pyo3(from_py_with = ...)]` argument, reporting failures as `ValueError` /// so a bad argument reads as a bad argument rather than as whatever the conversion hit. @@ -71,6 +72,119 @@ pub fn panic_to_pyerr(payload: Box) -> PyErr { PanicException::new_err(message.to_string()) } +pub fn json_from_py(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_none() { + return Ok(Value::Null); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if let Ok(int) = value.cast::() { + if let Ok(number) = int.extract::() { + return Ok(Value::from(number)); + } + if let Ok(number) = int.extract::() { + return Ok(Value::from(number)); + } + return raw_leaf(value); + } + if let Ok(float) = value.cast::() { + return match Number::from_f64(float.value()) { + Some(number) => Ok(Value::Number(number)), + None => raw_leaf(value), + }; + } + if let Ok(text) = value.cast::() { + return json_string(text); + } + if let Ok(dict) = value.cast::() { + if !dict + .keys() + .iter() + .all(|key| key.is_instance_of::()) + { + return from_py(value); + } + return dict + .iter() + .map(|(key, item)| Ok((json_key(key.cast::()?)?, json_from_py(&item)?))) + .collect::>>() + .map(Value::Object); + } + if let Ok(list) = value.cast::() { + return list + .iter() + .map(|item| json_from_py(&item)) + .collect::>() + .map(Value::Array); + } + if let Ok(tuple) = value.cast::() { + return tuple + .iter() + .map(|item| json_from_py(&item)) + .collect::>() + .map(Value::Array); + } + from_py(value) +} + +pub fn json_to_py<'py>(py: Python<'py>, value: &Value) -> PyResult> { + Ok(match value { + Value::Null => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, *flag).to_owned().into_any(), + Value::Number(number) => match (number.as_i64(), number.as_u64()) { + (Some(int), _) => int.into_pyobject(py)?.into_any(), + (None, Some(int)) => int.into_pyobject(py)?.into_any(), + (None, None) => PyFloat::new(py, number.as_f64().unwrap_or(f64::NAN)).into_any(), + }, + Value::String(text) => match litellm_host::raw_json::decode(text) { + Some(raw) => py.import("json")?.call_method1("loads", (raw,))?, + None => PyString::new(py, text).into_any(), + }, + Value::Array(items) => PyList::new( + py, + items + .iter() + .map(|item| json_to_py(py, item)) + .collect::>>()?, + )? + .into_any(), + Value::Object(fields) => { + let dict = PyDict::new(py); + for (key, item) in fields { + dict.set_item( + json_to_py(py, &Value::String(key.clone()))?, + json_to_py(py, item)?, + )?; + } + dict.into_any() + } + }) +} + +fn json_key(text: &Bound<'_, PyString>) -> PyResult { + match json_string(text)? { + Value::String(key) => Ok(key), + _ => unreachable!("a Python str converts to a JSON string"), + } +} + +fn json_string(text: &Bound<'_, PyString>) -> PyResult { + match text.to_str() { + Ok(text) if !litellm_host::raw_json::is_marked(text) => Ok(Value::String(text.to_owned())), + _ => raw_leaf(text.as_any()), + } +} + +fn raw_leaf(value: &Bound<'_, PyAny>) -> PyResult { + let raw: String = value + .py() + .import("json")? + .call_method1("dumps", (value,))? + .extract()?; + Ok(Value::String(litellm_host::raw_json::encode(&raw))) +} + #[cfg(test)] mod tests { use serde::Serializer; @@ -131,6 +245,72 @@ mod tests { }); } + #[rstest::rstest] + #[case::plain( + c"{'a': [1, 2.5, None, True, 'x'], 'b': (1,)}", + r#"{"a":[1,2.5,null,true,"x"],"b":[1]}"# + )] + #[case::nan(c"{'t': float('nan')}", r#"{"t":NaN}"#)] + #[case::infinities(c"[float('inf'), -float('inf')]", "[Infinity,-Infinity]")] + #[case::wide_integers( + c"[2**64, -(2**70), 2**64 - 1, -(2**63)]", + "[18446744073709551616,-1180591620717411303424,18446744073709551615,-9223372036854775808]" + )] + #[case::lone_surrogates(c"{'k\\udc00': 'a\\ud800b \\u00e9'}", r#"{"k\udc00":"a\ud800b é"}"#)] + #[case::marked_text(c"['\\uf8fflitellm-raw-json:1']", r#"["litellm-raw-json:1"]"#)] + fn json_values_keep_what_serde_json_cannot_hold( + #[case] source: &std::ffi::CStr, + #[case] wire: &str, + ) { + crate::initialize_python(); + Python::attach(|py| { + let value = py.eval(source, None, None).unwrap(); + let converted = json_from_py(&value).unwrap(); + let json = py.import("json").unwrap(); + let loads = |text: &str| json.call_method1("loads", (text,)).unwrap(); + let wire_bytes = litellm_host::raw_json::to_vec(&converted).unwrap(); + assert!( + loads(std::str::from_utf8(&wire_bytes).unwrap()) + .eq(loads(wire)) + .unwrap() + || wire.contains("NaN") + ); + assert_eq!( + json.call_method1("dumps", (loads(std::str::from_utf8(&wire_bytes).unwrap()),)) + .unwrap() + .extract::() + .unwrap(), + json.call_method1("dumps", (loads(wire),)) + .unwrap() + .extract::() + .unwrap() + ); + let back = json_to_py(py, &converted).unwrap(); + assert_eq!( + json.call_method1("dumps", (back,)) + .unwrap() + .extract::() + .unwrap(), + json.call_method1("dumps", (loads(wire),)) + .unwrap() + .extract::() + .unwrap() + ); + }); + } + + #[test] + fn values_json_cannot_hold_keep_their_old_conversion() { + crate::initialize_python(); + Python::attach(|py| { + let value = py.eval(c"{1: 'a'}", None, None).unwrap(); + assert_eq!( + json_from_py(&value).unwrap_err().to_string(), + from_py::(&value).unwrap_err().to_string() + ); + }); + } + #[test] fn depythonize_preserves_python_exception_identity_and_traceback() { crate::initialize_python(); diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 15e864c0718..f4f0903cc45 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -422,6 +422,9 @@ impl CallHooks for SyntheticHooks { PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { format!("response:{}", raw.body) } + PythonCallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => { + format!("retried:{}", retry.body) + } PythonCallEvent::Succeeded { response, .. } => { format!("succeeded:{}", response.bind(py)) } diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index e0ccab4e0ad..2c113bc644c 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true litellm-secret-value.workspace = true futures-util.workspace = true litellm-coroutine.workspace = true +serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 9398ea4bad6..8146448cd98 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -10,6 +10,18 @@ pub struct WireRequest { pub body: Value, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RetryReason { + ThinkingSignature, + StrictTools { dropped: usize }, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RetriedRequest { + pub reasons: Vec, + pub body: Value, +} + /// What the route knows about the request it is sending, for a host that logs it. The /// route owns these facts; a host reads them beside the wire request and never rewrites /// them. @@ -62,6 +74,7 @@ pub struct ProviderIdentity { pub enum ResultSource { Provider, Cache { key: String }, + Host, } /// What the call cost, settled once by the route. A host records it; it never re-derives it. @@ -79,6 +92,7 @@ impl Cost { pub fn settle(source: &ResultSource, reported: Option) -> Self { match (source, reported) { (ResultSource::Cache { .. }, _) => Self::CacheHit, + (ResultSource::Host, _) => Self::Deferred, (ResultSource::Provider, Some(amount)) => Self::Reported { amount }, (ResultSource::Provider, None) => Self::Deferred, } @@ -107,6 +121,13 @@ pub trait Interceptors: Send + Sync { &self, raw: RawResponse, ) -> impl Future> + Send; + + fn provider_request_retried( + &self, + _retry: RetriedRequest, + ) -> impl Future> + Send { + async { Ok(()) } + } } impl + ?Sized> Interceptors for &T { @@ -128,6 +149,13 @@ impl + ?Sized> Interceptors for &T { ) -> impl Future> + Send { (**self).after_provider_response(raw) } + + fn provider_request_retried( + &self, + retry: RetriedRequest, + ) -> impl Future> + Send { + (**self).provider_request_retried(retry) + } } impl Interceptors for () { diff --git a/litellm-rust/crates/host/src/lib.rs b/litellm-rust/crates/host/src/lib.rs index 76fce582fb3..26c53c89f97 100644 --- a/litellm-rust/crates/host/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -13,3 +13,4 @@ pub mod lifecycle; pub mod machine; pub mod observation; pub mod protocol; +pub mod raw_json; diff --git a/litellm-rust/crates/host/src/lifecycle.rs b/litellm-rust/crates/host/src/lifecycle.rs index e6b31381494..16456ea52d0 100644 --- a/litellm-rust/crates/host/src/lifecycle.rs +++ b/litellm-rust/crates/host/src/lifecycle.rs @@ -6,7 +6,10 @@ use std::{ use futures_util::TryStreamExt; -use crate::{call::CallOutput, interceptors::RawResponse}; +use crate::{ + call::CallOutput, + interceptors::{RawResponse, RetriedRequest}, +}; /// Seconds since the Unix epoch, on one clock for every host. pub fn epoch_seconds() -> f64 { @@ -50,7 +53,7 @@ pub enum CallEvent { }, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq)] pub enum ExecutionEvent { ResultReady { facts: crate::interceptors::ExecutionFacts, @@ -58,6 +61,9 @@ pub enum ExecutionEvent { ProviderResponseReceived { raw: Raw, }, + ProviderRequestRetried { + retry: RetriedRequest, + }, } impl> CallEvent { @@ -76,6 +82,11 @@ impl> CallEvent { + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { + retry: retry.clone(), + }) + } Self::Succeeded { timing, .. } => CallEvent::Succeeded { timing: *timing, response: (), diff --git a/litellm-rust/crates/host/src/machine/context.rs b/litellm-rust/crates/host/src/machine/context.rs index 1e53605ed7b..dacb470e97f 100644 --- a/litellm-rust/crates/host/src/machine/context.rs +++ b/litellm-rust/crates/host/src/machine/context.rs @@ -5,7 +5,7 @@ use litellm_coroutine::Co; use super::coroutine::MachineFault; use crate::{ - interceptors::{RawResponse, RequestContext, WireRequest}, + interceptors::{RawResponse, RequestContext, RetriedRequest, WireRequest}, protocol::{HostRequest, InterceptRequest, Protocol, Reply, StreamDelivery}, }; @@ -119,6 +119,14 @@ where }) .await? } + + async fn provider_request_retried(&self, retry: RetriedRequest) -> Result<(), P::Error> { + self.0 + .request_reply(|reply| { + HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { retry, reply }) + }) + .await? + } } pub struct StreamSender(Channel

); diff --git a/litellm-rust/crates/host/src/protocol.rs b/litellm-rust/crates/host/src/protocol.rs index c24557b0348..5ec566b6a65 100644 --- a/litellm-rust/crates/host/src/protocol.rs +++ b/litellm-rust/crates/host/src/protocol.rs @@ -2,7 +2,7 @@ use std::ops::ControlFlow; pub use litellm_coroutine::{Abandoned, Answer, Reply, reply}; -use crate::interceptors::{RawResponse, RequestContext, WireRequest}; +use crate::interceptors::{RawResponse, RequestContext, RetriedRequest, WireRequest}; pub trait Protocol: Send + Sync + 'static { type Request: Send + 'static; @@ -33,6 +33,10 @@ pub enum InterceptRequest { raw: RawResponse, reply: Reply>, }, + ProviderRequestRetried { + retry: RetriedRequest, + reply: Reply>, + }, } pub enum StreamDelivery { diff --git a/litellm-rust/crates/host/src/raw_json.rs b/litellm-rust/crates/host/src/raw_json.rs new file mode 100644 index 00000000000..fbf9bbc6ecb --- /dev/null +++ b/litellm-rust/crates/host/src/raw_json.rs @@ -0,0 +1,362 @@ +use std::{borrow::Cow, io}; + +use serde::Serialize; +use serde_json::ser::{CharEscape, CompactFormatter, Formatter}; + +const MARK: char = '\u{F8FF}'; +const PREFIX: &str = "\u{F8FF}litellm-raw-json:"; + +pub fn encode(raw: &str) -> String { + format!("{PREFIX}{raw}") +} + +pub fn decode(text: &str) -> Option<&str> { + text.strip_prefix(PREFIX).filter(|raw| is_scalar(raw)) +} + +pub fn is_marked(text: &str) -> bool { + text.starts_with(PREFIX) +} + +pub fn to_vec(value: &T) -> serde_json::Result> { + let mut out = Vec::with_capacity(128); + let mut serializer = serde_json::Serializer::with_formatter(&mut out, RawLeaves::default()); + value.serialize(&mut serializer)?; + Ok(out) +} + +pub fn mark_wide_numbers(text: &str) -> Cow<'_, str> { + let bytes = text.as_bytes(); + let mut out = String::new(); + let mut copied = 0; + let mut index = 0; + while index < bytes.len() { + match bytes[index] { + b'"' => { + let end = string_end(bytes, index + 1); + if let Some(escaped) = escaped_marked_string(&text[index..end]) { + out.push_str(&text[copied..index]); + out.push_str(&escaped); + copied = end; + } + index = end; + } + b'-' | b'0'..=b'9' => { + let end = bytes[index..] + .iter() + .position(|byte| { + !matches!(byte, b'0'..=b'9' | b'-' | b'+' | b'.' | b'e' | b'E') + }) + .map_or(bytes.len(), |length| index + length); + if let Some(raw) = wide_number(&text[index..end]) { + out.push_str(&text[copied..index]); + out.push_str(&serde_json::Value::String(encode(raw)).to_string()); + copied = end; + } + index = end; + } + _ => index += 1, + } + } + if copied == 0 { + return Cow::Borrowed(text); + } + out.push_str(&text[copied..]); + Cow::Owned(out) +} + +fn escaped_marked_string(token: &str) -> Option { + let inner = token.strip_prefix('"')?; + let escaped_mark = inner + .get(..6) + .is_some_and(|head| head.eq_ignore_ascii_case("\\uf8ff")); + if !inner.starts_with(MARK) && !escaped_mark { + return None; + } + let text = serde_json::from_str::(token).ok()?; + is_marked(&text) + .then(|| serde_json::Value::String(encode(&ascii_string_literal(&text))).to_string()) +} + +fn ascii_string_literal(text: &str) -> String { + let mut literal = String::with_capacity(text.len() + 2); + literal.push('"'); + for character in text.chars() { + match character { + '"' => literal.push_str("\\\""), + '\\' => literal.push_str("\\\\"), + '\n' => literal.push_str("\\n"), + '\r' => literal.push_str("\\r"), + '\t' => literal.push_str("\\t"), + ' '..='~' => literal.push(character), + _ => { + let mut units = [0; 2]; + for unit in character.encode_utf16(&mut units) { + literal.push_str(&format!("\\u{unit:04x}")); + } + } + } + } + literal.push('"'); + literal +} + +fn string_end(bytes: &[u8], mut index: usize) -> usize { + while index < bytes.len() { + match bytes[index] { + b'\\' => index += 2, + b'"' => return index + 1, + _ => index += 1, + } + } + index +} + +fn wide_number(raw: &str) -> Option<&str> { + if is_integer(raw) { + return (raw.parse::().is_err() && raw.parse::().is_err()).then_some(raw); + } + match raw.parse::() { + Ok(number) if number == f64::INFINITY => Some("Infinity"), + Ok(number) if number == f64::NEG_INFINITY => Some("-Infinity"), + _ => None, + } +} + +fn is_scalar(raw: &str) -> bool { + matches!(raw, "NaN" | "Infinity" | "-Infinity") || is_integer(raw) || is_string_literal(raw) +} + +fn is_integer(raw: &str) -> bool { + let digits = raw.strip_prefix('-').unwrap_or(raw); + match digits.as_bytes() { + [b'0'] => true, + [b'1'..=b'9', rest @ ..] => rest.iter().all(u8::is_ascii_digit), + _ => false, + } +} + +fn is_string_literal(raw: &str) -> bool { + let Some(inner) = raw + .strip_prefix('"') + .and_then(|rest| rest.strip_suffix('"')) + else { + return false; + }; + let mut bytes = inner.bytes(); + while let Some(byte) = bytes.next() { + match byte { + b'\\' => match bytes.next() { + Some(b'"' | b'\\' | b'/' | b'b' | b'f' | b'n' | b'r' | b't') => {} + Some(b'u') => { + if !(0..4).all(|_| bytes.next().is_some_and(|hex| hex.is_ascii_hexdigit())) { + return false; + } + } + _ => return false, + }, + b'"' => return false, + 0x20..=0x7e => {} + _ => return false, + } + } + true +} + +fn unescaped(escape: CharEscape) -> char { + match escape { + CharEscape::Quote => '"', + CharEscape::ReverseSolidus => '\\', + CharEscape::Solidus => '/', + CharEscape::Backspace => '\u{8}', + CharEscape::FormFeed => '\u{c}', + CharEscape::LineFeed => '\n', + CharEscape::CarriageReturn => '\r', + CharEscape::Tab => '\t', + CharEscape::AsciiControl(byte) => char::from(byte), + } +} + +#[derive(Default)] +enum Text { + #[default] + Closed, + Opened, + Plain, + Marked(String), +} + +#[derive(Default)] +struct RawLeaves { + key: bool, + text: Text, +} + +impl Formatter for RawLeaves { + fn begin_object_key( + &mut self, + writer: &mut W, + first: bool, + ) -> io::Result<()> { + self.key = true; + CompactFormatter.begin_object_key(writer, first) + } + + fn end_object_key(&mut self, writer: &mut W) -> io::Result<()> { + self.key = false; + CompactFormatter.end_object_key(writer) + } + + fn begin_string(&mut self, _writer: &mut W) -> io::Result<()> { + self.text = Text::Opened; + Ok(()) + } + + fn write_string_fragment( + &mut self, + writer: &mut W, + fragment: &str, + ) -> io::Result<()> { + match &mut self.text { + Text::Opened if fragment.starts_with(MARK) => { + self.text = Text::Marked(fragment.to_owned()); + Ok(()) + } + Text::Opened => { + self.text = Text::Plain; + writer.write_all(b"\"")?; + writer.write_all(fragment.as_bytes()) + } + Text::Marked(held) => { + held.push_str(fragment); + Ok(()) + } + Text::Plain | Text::Closed => writer.write_all(fragment.as_bytes()), + } + } + + fn write_char_escape( + &mut self, + writer: &mut W, + escape: CharEscape, + ) -> io::Result<()> { + match &mut self.text { + Text::Opened => { + self.text = Text::Plain; + writer.write_all(b"\"")?; + CompactFormatter.write_char_escape(writer, escape) + } + Text::Marked(held) => { + held.push(unescaped(escape)); + Ok(()) + } + Text::Plain | Text::Closed => CompactFormatter.write_char_escape(writer, escape), + } + } + + fn end_string(&mut self, writer: &mut W) -> io::Result<()> { + match std::mem::take(&mut self.text) { + Text::Opened => writer.write_all(b"\"\""), + Text::Plain | Text::Closed => writer.write_all(b"\""), + Text::Marked(held) => match decode(&held) { + Some(raw) if !self.key || raw.starts_with('"') => writer.write_all(raw.as_bytes()), + _ => serde_json::to_writer(writer, &held).map_err(io::Error::other), + }, + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::*; + + #[rstest] + #[case::nan("NaN")] + #[case::infinity("Infinity")] + #[case::negative_infinity("-Infinity")] + #[case::wide_integer("18446744073709551616")] + #[case::negative_wide_integer("-1180591620717411303424")] + #[case::zero("0")] + #[case::surrogate_string(r#""sys \ud800 lone""#)] + #[case::escaped_string(r#""a\"b\\c\n\u00e9""#)] + fn scalar_texts_are_written_back_verbatim(#[case] raw: &str) { + let body = json!({"leaf": encode(raw), "list": [encode(raw)]}); + assert_eq!( + String::from_utf8(to_vec(&body).unwrap()).unwrap(), + format!(r#"{{"leaf":{raw},"list":[{raw}]}}"#) + ); + } + + #[rstest] + #[case::structure("1, \"injected\": true")] + #[case::object("{}")] + #[case::leading_zero("012")] + #[case::float("1.5")] + #[case::unquoted_quote(r#""a"b""#)] + #[case::non_ascii("\"\u{e9}\"")] + #[case::bad_escape(r#""\x""#)] + #[case::short_unicode_escape(r#""\u12""#)] + #[case::empty("")] + fn other_marked_texts_stay_strings(#[case] raw: &str) { + let marked = encode(raw); + assert_eq!(decode(&marked), None); + assert_eq!( + to_vec(&json!([marked])).unwrap(), + serde_json::to_vec(&json!([marked])).unwrap() + ); + } + + #[test] + fn keys_take_only_string_texts() { + let body = Value::Object( + [ + (encode(r#""k\udc00""#), json!(1)), + (encode("NaN"), json!(2)), + ] + .into_iter() + .collect(), + ); + assert_eq!( + String::from_utf8(to_vec(&body).unwrap()).unwrap(), + format!( + r#"{{"k\udc00":1,{}:2}}"#, + serde_json::to_string(&encode("NaN")).unwrap() + ) + ); + } + + #[rstest] + #[case::wide_integer(r#"{"a":18446744073709551616}"#, json!({"a": encode("18446744073709551616")}))] + #[case::negative_wide_integer(r#"[-12345678901234567890123]"#, json!([encode("-12345678901234567890123")]))] + #[case::overflowing_float(r#"{"a":[1e400,-1E+400]}"#, json!({"a": [encode("Infinity"), encode("-Infinity")]}))] + #[case::numbers_that_fit(r#"{"a":[18446744073709551615,-9223372036854775808,-0.0,5e-324,1E2,0.30000000000000004]}"#, json!({"a": [18_446_744_073_709_551_615_u64, i64::MIN, -0.0, 5e-324, 100.0, 0.300_000_000_000_000_04]}))] + #[case::provider_marked_text("{\"a\":\"\u{F8FF}litellm-raw-json:42\"}", json!({"a": encode(r#""\uf8fflitellm-raw-json:42""#)}))] + #[case::provider_escaped_mark(r#"{"\uF8FFlitellm-raw-json:1":["\uf8fflitellm-raw-json:\"x\"\n\u00e9"]}"#, json!({encode(r#""\uf8fflitellm-raw-json:1""#): [encode(r#""\uf8fflitellm-raw-json:\"x\"\n\u00e9""#)]}))] + #[case::mark_not_leading(r#"["x\uf8fflitellm-raw-json:NaN"]"#, json!(["x\u{F8FF}litellm-raw-json:NaN"]))] + #[case::numbers_inside_strings(r#"{"18446744073709551616":"1e400 \" 18446744073709551616"}"#, json!({"18446744073709551616": "1e400 \" 18446744073709551616"}))] + fn wide_numbers_are_read_as_marked_texts(#[case] text: &str, #[case] read: Value) { + assert_eq!( + serde_json::from_str::(&mark_wide_numbers(text)).unwrap(), + read + ); + } + + #[test] + fn a_text_without_wide_numbers_is_borrowed() { + assert!(matches!( + mark_wide_numbers(r#"{"a":[1,2.5,"18446744073709551616"]}"#), + Cow::Borrowed(_) + )); + } + + #[rstest] + #[case::plain(json!({"a": "text", "b": ["", "\"quoted\"\n", "\u{F8FF}"], "c": 1.5, "d": null}))] + #[case::mark_inside(json!(["x\u{F8FF}litellm-raw-json:NaN"]))] + #[case::control(json!(["\u{1}\t"]))] + fn unmarked_documents_serialize_as_serde_json_does(#[case] body: Value) { + assert_eq!(to_vec(&body).unwrap(), serde_json::to_vec(&body).unwrap()); + } +} diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 5e5f2d66fa2..8e9a7000e69 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -94,6 +94,8 @@ impl From<&HttpSettings> for Resolution { } } +const MAX_REDIRECTS: usize = 20; + impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { type Error = Error; @@ -101,7 +103,9 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { let base = reqwest::Client::builder() .tls_backend_preconfigured(rustls::ClientConfig::try_from(config)?) .connect_timeout(config.connect_timeout) - .pool_idle_timeout(config.pool_idle_timeout); + .pool_idle_timeout(config.pool_idle_timeout) + .referer(false) + .redirect(reqwest::redirect::Policy::limited(MAX_REDIRECTS)); let with_keepalive = match config.tcp_keepalive { None => base, Some(keepalive) => base @@ -186,17 +190,23 @@ mod tests { assert_eq!(resolution.unsupported, []); } - #[test] - fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() { + #[rstest] + #[case::openssl_lacks_in_rustls("secp521r1")] + #[case::nist_name("P-256")] + #[case::nist_p384("P-384")] + #[case::sec_name("secp256r1")] + #[case::lowercase("x25519")] + #[case::padded(" X25519 ")] + fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported(#[case] curve: &str) { let settings = HttpSettings { - ssl_ecdh_curve: Some("secp521r1".into()), + ssl_ecdh_curve: Some(curve.into()), ..HttpSettings::default() }; let resolution = Resolution::from(&settings); assert_eq!(resolution.config.key_exchange_group, None); assert_eq!( resolution.unsupported, - [Unsupported::EcdhCurve("secp521r1".into())] + [Unsupported::EcdhCurve(curve.into())] ); } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 9505aacf44b..d9fdfb5486f 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -251,6 +251,10 @@ impl Error { Repr::ComputedHeader(name.into()).into() } + pub fn request_body(error: serde_json::Error) -> Self { + Repr::RequestBody(error).into() + } + pub fn unsignable_body() -> Self { Repr::UnsignableBody.into() } diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index 20b0b1be972..b11041151af 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -27,7 +27,20 @@ pub fn json( body: &impl Serialize, timeout: Option, ) -> Result { - let body = serde_json::to_vec(body).map_err(Repr::RequestBody)?; + json_bytes( + url, + headers, + serde_json::to_vec(body).map_err(Repr::RequestBody)?, + timeout, + ) +} + +pub fn json_bytes( + url: &str, + headers: &[(String, String)], + body: Vec, + timeout: Option, +) -> Result { let url = Url::parse(url).map_err(Repr::ParseUrl)?; let with_headers = headers.iter().fold( http::Request::post(url.as_str()), diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 3b20bbda607..c56c5de4a11 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -16,6 +16,7 @@ use crate::{ #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { Provider, + Undecoded, NoRedirect, Media, UnpinnedMedia, @@ -103,6 +104,7 @@ impl HttpClientPool { ) -> reqwest::ClientBuilder { match variant { ClientVariant::Provider => builder, + ClientVariant::Undecoded => builder.no_gzip().no_deflate().no_brotli().no_zstd(), ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => { builder.redirect(reqwest::redirect::Policy::none()) } @@ -132,7 +134,9 @@ fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientCon }, ..config.clone() }, - ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), + ClientVariant::Provider | ClientVariant::Undecoded | ClientVariant::NoRedirect => { + config.clone() + } } } @@ -367,6 +371,28 @@ mod tests { assert!(request.contains("user-agent: litellm-test/9"), "{request}"); } + #[tokio::test] + async fn provider_client_follows_twenty_redirects_without_a_referer() { + let (address, _, requests) = serve("HTTP/1.1 307 Temporary Redirect").await; + let followed = pool() + .client(&config("a"), ClientVariant::Provider) + .unwrap() + .get(format!("http://{address}/start")) + .timeout(Duration::from_secs(5)) + .send() + .await; + assert!(followed.is_err()); + let requests = requests.lock().unwrap().clone(); + assert_eq!(requests.len(), 21, "{requests:?}"); + assert!(requests[1].starts_with("GET /elsewhere "), "{requests:?}"); + assert!( + requests + .iter() + .all(|request| !request.to_ascii_lowercase().contains("referer:")), + "{requests:?}" + ); + } + #[tokio::test] async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() { let (address, _, _) = serve("HTTP/1.1 302 Found").await; diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index a83b0018b2b..47a1429f1a4 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -23,10 +23,10 @@ impl FromStr for KeyExchangeGroup { type Err = Unsupported; fn from_str(name: &str) -> Result { - match name.trim().to_ascii_lowercase().as_str() { - "x25519" => Ok(Self::X25519), - "prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1), - "secp384r1" | "p-384" => Ok(Self::Secp384r1), + match name { + "X25519" => Ok(Self::X25519), + "prime256v1" => Ok(Self::Secp256r1), + "secp384r1" => Ok(Self::Secp384r1), _ => Err(Unsupported::EcdhCurve(name.to_owned())), } } diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index dd229d3d657..d8d3a8956c8 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -117,6 +117,7 @@ async fn the_same_route_entrypoint_reports_facts_with_or_without_caching( assert!(caching); assert!(!key.is_empty()); } + ResultSource::Host => panic!("the host never answers a chat call"), } let summaries = traces.summaries("litellm.route"); assert_eq!(summaries.len(), 2); diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 0c5be52550d..0ca371c3ce7 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -8,7 +8,9 @@ repository.workspace = true [dependencies] litellm-secret-value.workspace = true bytes.workspace = true +flate2 = { workspace = true, features = ["zlib-rs"] } futures-util.workspace = true +http.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-auth-aws.workspace = true litellm-cache-response.workspace = true @@ -30,9 +32,11 @@ tracing.workspace = true [dev-dependencies] litellm-error.workspace = true litellm-framer = { workspace = true, features = ["sse"] } +litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-host-native.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true rstest.workspace = true +thiserror.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/inference-messages/src/constants.rs b/litellm-rust/crates/inference-messages/src/constants.rs index c0b023be7c7..5fae6e45dad 100644 --- a/litellm-rust/crates/inference-messages/src/constants.rs +++ b/litellm-rust/crates/inference-messages/src/constants.rs @@ -2,3 +2,6 @@ /// seconds. Mirrors the Python Anthropic Messages default. The per-request /// timeout from the caller still overrides this on the request builder. pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600; + +pub(crate) const MESSAGES_HTTP_ERROR_ATTEMPTS: usize = 2; +pub(crate) const MESSAGES_STREAM_ERROR_BODY_READ_SECS: u64 = 5; diff --git a/litellm-rust/crates/inference-messages/src/encoding.rs b/litellm-rust/crates/inference-messages/src/encoding.rs new file mode 100644 index 00000000000..c315deb1a33 --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/encoding.rs @@ -0,0 +1,437 @@ +use bytes::Bytes; +use flate2::{Decompress, DecompressError, FlushDecompress, Status}; +use futures_util::StreamExt; +use reqwest::header::{CONTENT_ENCODING, HeaderMap}; + +type BoxError = Box; + +const MAX_WINDOW_BITS: u8 = 15; +const MAX_DECOMPRESS_SIZE: usize = 1 << 16; + +#[derive(Debug)] +pub struct DecodingError(String); + +impl std::fmt::Display for DecodingError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} + +impl std::error::Error for DecodingError {} + +impl From for DecodingError { + fn from(error: DecompressError) -> Self { + Self(match (error.needs_dictionary(), error.message()) { + (Some(_), _) => "Error 2 while decompressing data".to_owned(), + (None, Some(message)) => format!("Error -3 while decompressing data: {message}"), + (None, None) => "Error -3 while decompressing data: invalid input data".to_owned(), + }) + } +} + +enum Codec { + Gzip, + Deflate, +} + +enum FirstAttempt { + Unread, + Reading(Vec), + Over, +} + +struct Inflater { + zlib: Decompress, + first_attempt: FirstAttempt, + pending: Vec, + saturated: bool, + stalled: bool, + finished: bool, +} + +impl Inflater { + fn new(codec: Codec) -> Self { + let (zlib, first_attempt) = match codec { + Codec::Gzip => (Decompress::new_gzip(MAX_WINDOW_BITS), FirstAttempt::Over), + Codec::Deflate => (Decompress::new(true), FirstAttempt::Unread), + }; + Self { + zlib, + first_attempt, + pending: Vec::new(), + saturated: false, + stalled: false, + finished: false, + } + } + + fn push(&mut self, data: &[u8]) { + self.first_attempt = match self.first_attempt { + FirstAttempt::Unread => FirstAttempt::Reading(data.to_vec()), + FirstAttempt::Reading(_) | FirstAttempt::Over => FirstAttempt::Over, + }; + if self.finished { + return; + } + self.pending.extend_from_slice(data); + self.stalled = false; + } + + fn has_work(&self) -> bool { + !self.finished && (self.saturated || (!self.pending.is_empty() && !self.stalled)) + } + + fn step(&mut self) -> Result, DecodingError> { + match (self.inflate(), &mut self.first_attempt) { + (Err(_), FirstAttempt::Reading(first)) => { + self.pending = std::mem::take(first); + self.first_attempt = FirstAttempt::Over; + self.zlib = Decompress::new_with_window_bits(false, MAX_WINDOW_BITS); + self.saturated = false; + self.stalled = false; + self.step() + } + (inflated, attempt) => { + if self.pending.is_empty() && !self.saturated { + *attempt = FirstAttempt::Over; + } + inflated + } + } + } + + fn inflate(&mut self) -> Result, DecodingError> { + let mut output = Vec::with_capacity(MAX_DECOMPRESS_SIZE); + let (read, written) = (self.zlib.total_in(), self.zlib.total_out()); + let status = self + .zlib + .decompress_vec(&self.pending, &mut output, FlushDecompress::Sync)?; + self.pending.drain(..(self.zlib.total_in() - read) as usize); + self.finished = matches!(status, Status::StreamEnd); + if self.finished { + self.pending.clear(); + } + self.saturated = !self.finished && output.len() == MAX_DECOMPRESS_SIZE; + self.stalled = self.zlib.total_in() == read && self.zlib.total_out() == written; + Ok(output) + } +} + +struct Decoder(Vec); + +impl Decoder { + fn detect(headers: &HeaderMap) -> Option { + let inflaters: Vec = header_pairs(headers) + .into_iter() + .filter(|(name, _)| name.as_str() == CONTENT_ENCODING.as_str()) + .flat_map(|(_, value)| { + value + .split(',') + .map(|coding| coding.trim().to_lowercase()) + .collect::>() + }) + .filter_map(|coding| match coding.as_str() { + "gzip" => Some(Codec::Gzip), + "deflate" => Some(Codec::Deflate), + _ => None, + }) + .rev() + .map(Inflater::new) + .collect(); + (!inflaters.is_empty()).then_some(Self(inflaters)) + } + + fn feed(&mut self, chunk: &[u8]) { + self.0[0].push(chunk); + } + + fn next(&mut self) -> Result, DecodingError> { + self.pull(self.0.len() - 1) + .map(|decoded| decoded.map(Bytes::from)) + } + + fn pull(&mut self, layer: usize) -> Result>, DecodingError> { + loop { + if self.0[layer].has_work() { + let decoded = self.0[layer].step()?; + if !decoded.is_empty() { + return Ok(Some(decoded)); + } + continue; + } + let Some(outer) = layer.checked_sub(1) else { + return Ok(None); + }; + match self.pull(outer)? { + Some(decoded) => self.0[layer].push(&decoded), + None => return Ok(None), + } + } + } +} + +pub(crate) fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> { + let utf8 = headers + .values() + .all(|value| std::str::from_utf8(value.as_bytes()).is_ok()); + headers + .iter() + .map(|(name, value)| { + let bytes = value.as_bytes(); + let text = if utf8 { + String::from_utf8_lossy(bytes).into_owned() + } else { + bytes.iter().copied().map(char::from).collect() + }; + (name.to_string(), text) + }) + .collect() +} + +pub(crate) fn decoded(response: reqwest::Response) -> reqwest::Response { + let Some(decoder) = Decoder::detect(response.headers()) else { + return response; + }; + let status = response.status(); + let version = response.version(); + let headers = response.headers().clone(); + let body = futures_util::stream::unfold( + Some((response.bytes_stream().boxed(), decoder)), + |state| async move { + let (mut body, mut decoder) = state?; + loop { + match decoder.next() { + Ok(Some(decoded)) => return Some((Ok(decoded), Some((body, decoder)))), + Ok(None) => {} + Err(error) => return Some((Err(BoxError::from(error)), None)), + } + match body.next().await? { + Ok(chunk) => decoder.feed(&chunk), + Err(error) => return Some((Err(BoxError::from(error)), None)), + } + } + }, + ); + let mut rebuilt = http::Response::new(reqwest::Body::wrap_stream(body)); + *rebuilt.status_mut() = status; + *rebuilt.version_mut() = version; + *rebuilt.headers_mut() = headers; + reqwest::Response::from(rebuilt) +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use flate2::{Compression, write::GzEncoder, write::ZlibEncoder}; + use futures_util::TryStreamExt; + use rstest::rstest; + + use reqwest::header::{HeaderMap, HeaderValue}; + + use super::{decoded, header_pairs}; + + fn encoded(encoding: &str, body: &[u8]) -> Vec { + match encoding { + "gzip" => { + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(body).unwrap(); + encoder.finish().unwrap() + } + "deflate" => { + let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(body).unwrap(); + encoder.finish().unwrap() + } + _ => body.to_vec(), + } + } + + fn response(encoding: &str, body: Vec) -> reqwest::Response { + let length = body.len().to_string(); + let mut response = http::Response::new(reqwest::Body::from(body)); + let headers = response.headers_mut(); + headers.insert("content-encoding", encoding.parse().unwrap()); + headers.insert("content-length", length.parse().unwrap()); + reqwest::Response::from(response) + } + + #[rstest] + #[case::gzip("gzip")] + #[case::deflate("deflate")] + #[case::unknown("br")] + #[tokio::test] + async fn decoded_bodies_keep_the_head_the_provider_sent(#[case] encoding: &str) { + let body = br#"{"type":"message"}"#; + let wire = encoded(encoding, body); + let length = wire.len().to_string(); + let response = decoded(response(encoding, wire.clone())); + + assert_eq!(response.headers()["content-encoding"], encoding); + assert_eq!(response.headers()["content-length"], length.as_str()); + let expected = if encoding == "br" { + wire + } else { + body.to_vec() + }; + assert_eq!( + response.bytes().await.unwrap().as_ref(), + expected.as_slice() + ); + } + + fn raw_deflate(body: &[u8]) -> Vec { + let mut encoder = flate2::write::DeflateEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(body).unwrap(); + encoder.finish().unwrap() + } + + fn chunked(encoding: &str, wire: Vec, size: usize) -> reqwest::Response { + let chunks: Vec, std::io::Error>> = + wire.chunks(size).map(|chunk| Ok(chunk.to_vec())).collect(); + let mut response = http::Response::new(reqwest::Body::wrap_stream( + futures_util::stream::iter(chunks), + )); + response + .headers_mut() + .insert("content-encoding", encoding.parse().unwrap()); + reqwest::Response::from(response) + } + + async fn read(response: reqwest::Response) -> Result, String> { + response + .bytes() + .await + .map(|bytes| bytes.to_vec()) + .map_err(|error| { + litellm_error::find::(&error) + .map_or_else(|| error.to_string(), ToString::to_string) + }) + } + + const BODY: &[u8] = br#"{"type":"message","content":[{"type":"text","text":"Paris"}]}"#; + + #[rstest] + #[case::padded_upper_case(" GZIP ", encoded("gzip", BODY))] + #[case::listed_with_an_unknown_coding("gzip, br", encoded("gzip", BODY))] + #[case::applied_twice("gzip, gzip", encoded("gzip", &encoded("gzip", BODY)))] + #[case::deflate_without_its_zlib_header("deflate", raw_deflate(BODY))] + #[case::gzip_with_data_past_its_end("gzip", [encoded("gzip", BODY), b"\x1f\x8b trailing".to_vec()].concat())] + #[tokio::test] + async fn bodies_decode_as_httpx_decodes_them( + #[case] encoding: &str, + #[case] wire: Vec, + #[values(7, 4096)] chunk: usize, + ) { + assert_eq!( + read(decoded(chunked(encoding, wire, chunk))).await.unwrap(), + BODY + ); + } + + #[rstest] + #[case::x_gzip_alias(" X-Gzip ")] + #[case::unknown_only("br, identity")] + #[tokio::test] + async fn codings_httpx_has_no_decoder_for_pass_through(#[case] encoding: &str) { + let wire = encoded("gzip", BODY); + assert_eq!( + read(decoded(chunked(encoding, wire.clone(), 4096))) + .await + .unwrap(), + wire + ); + } + + #[rstest] + #[case::not_deflate("deflate", b"not deflate at all".to_vec(), "Error -3 while decompressing data: invalid block type")] + #[case::bad_gzip_trailer( + "gzip", + { let wire = encoded("gzip", BODY); [&wire[..wire.len() - 8], &[0; 8]].concat() }, + "Error -3 while decompressing data: incorrect data check" + )] + #[case::not_gzip("gzip", b"\x1f\x8b\x08\x00not really gzip data at all".to_vec(), "Error -3 while decompressing data: invalid stored block lengths")] + #[case::no_gzip_magic("gzip", b"plain text".to_vec(), "Error -3 while decompressing data: incorrect header check")] + #[tokio::test] + async fn corrupt_bodies_fail_with_zlib_s_error( + #[case] encoding: &str, + #[case] wire: Vec, + #[case] expected: &str, + ) { + assert_eq!( + read(decoded(chunked(encoding, wire, 4096))) + .await + .unwrap_err(), + expected + ); + } + + #[rstest] + #[case::gzip_in_gzip("gzip, gzip")] + #[case::deflate_in_gzip("deflate, gzip")] + #[tokio::test] + async fn a_nested_bomb_is_decoded_in_bounded_reads(#[case] encoding: &str) { + let zeros = vec![0; 8 << 20]; + let inner = if encoding.starts_with("deflate") { + encoded("deflate", &zeros) + } else { + encoded("gzip", &zeros) + }; + let wire = encoded("gzip", &inner); + let sizes: Vec = decoded(chunked(encoding, wire, 4096)) + .bytes_stream() + .map_ok(|chunk| chunk.len()) + .try_collect() + .await + .unwrap(); + assert!( + sizes.iter().all(|size| *size <= super::MAX_DECOMPRESS_SIZE), + "{sizes:?}" + ); + assert_eq!(sizes.iter().sum::(), zeros.len()); + } + + #[rstest] + #[tokio::test] + async fn a_truncated_body_yields_what_it_holds_without_an_error() { + let wire = encoded("gzip", BODY); + let decoded_body = read(decoded(chunked( + "gzip", + wire[..wire.len() - 8].to_vec(), + 4096, + ))) + .await + .unwrap(); + assert_eq!(decoded_body, BODY); + } + + #[rstest] + #[case::ascii(&[b"plain".as_slice()], &["plain"])] + #[case::utf_8(&[b"plain".as_slice(), b"caf\xc3\xa9"], &["plain", "caf\u{e9}"])] + #[case::latin_1_once_any_value_is_not_utf_8( + &[b"caf\xc3\xa9".as_slice(), b"caf\xe9"], + &["caf\u{c3}\u{a9}", "caf\u{e9}"] + )] + fn header_values_decode_as_httpx_decodes_them( + #[case] values: &[&[u8]], + #[case] expected: &[&str], + ) { + let headers: HeaderMap = values + .iter() + .enumerate() + .map(|(index, value)| { + ( + format!("x-fz-{index}") + .parse::() + .unwrap(), + HeaderValue::from_bytes(value).unwrap(), + ) + }) + .collect(); + let decoded: Vec = header_pairs(&headers) + .into_iter() + .map(|(_, value)| value) + .collect(); + assert_eq!(decoded, expected); + } +} diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index c7a7d93bf98..4796a640e25 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -18,7 +18,7 @@ use litellm_llms::base_llm::{ messages::{ context::MessagesTransformContext, streaming::{self, ByteStream, StreamDecoder}, - transformation::BaseMessagesConfig, + transformation::{BaseMessagesConfig, FederatedEnvironment}, }, }; use litellm_llms_types::formats::messages::{ @@ -29,13 +29,21 @@ use litellm_tracing::ByteChunk; use serde_json::{Map, Value}; use super::{ - Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, + Delivered, Error, HostAnswer, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute, + Transport, TransportBody, TransportRequest, TransportResponse, UnencodableHeader, + WorkloadIdentity, common_utils::{MessagesProvider, merge_messages_header_sources, string_headers}, - constants::MESSAGES_TIMEOUT_SECS, + constants::{ + MESSAGES_HTTP_ERROR_ATTEMPTS, MESSAGES_STREAM_ERROR_BODY_READ_SECS, MESSAGES_TIMEOUT_SECS, + }, route::Messages, types::invalid_request, }; -use litellm_inference::{caching::CallCache, context::CallContext, outbound::outbound_request}; +use litellm_inference::{ + caching::{CacheRequest, CallCache}, + context::CallContext, + outbound::raw_leaf_outbound_request, +}; struct ProviderCall { identity: ProviderIdentity, @@ -44,6 +52,12 @@ struct ProviderCall { signer: Option, timeout: Option, stream: bool, + whole_sse_frames: bool, + joined_reads: bool, + transform_context: MessagesTransformContext, + ascii_response_headers: bool, + unencodable_request_header: Option>, + ascii_request_headers: bool, } impl ProviderCall { @@ -52,16 +66,122 @@ impl ProviderCall { } } +pub(super) async fn cached( + route: &MessagesRoute, + provider: MessagesProvider, + model: &str, + context: &CallContext<'_, impl Interceptors>, +) -> Result, Error> { + let Some(cache) = caller_keyed(route) else { + return Ok(None); + }; + let identity = ProviderIdentity { + model: model.to_owned(), + provider: provider.as_str().to_string(), + }; + let Some((output, source)) = CallCache::::new( + Some(cache), + context.cache, + CacheRequest::from_wire(identity.clone(), None), + ) + .lookup() + .await + else { + return Ok(None); + }; + context + .result_ready(execution_facts(provider, identity, &output, &source)) + .await?; + Ok(Some(output)) +} + +pub(super) async fn mocked( + route: &MessagesRoute, + provider: MessagesProvider, + model: String, + context: &CallContext<'_, impl Interceptors>, + mock: Map, + stream: bool, +) -> MessagesCallResponse { + let output = MessagesCallResponse::Complete(MessagesBody::Unchecked(mock)); + let Some(cache) = caller_keyed(route).filter(|_| !stream) else { + return output; + }; + let identity = ProviderIdentity { + model, + provider: provider.as_str().to_string(), + }; + CallCache::::new( + Some(cache), + context.cache, + CacheRequest::from_wire(identity, None), + ) + .finish(output, &ResultSource::Provider) + .await +} + +pub(super) async fn answered( + route: &MessagesRoute, + host_answer: &dyn HostAnswer, + provider: MessagesProvider, + model: String, + stream: bool, + context: &CallContext<'_, impl Interceptors>, +) -> Result { + host_answer.answer(stream).await?; + deliver_answer(route, provider, model, stream, context).await +} + +pub(super) async fn deliver_answer( + route: &MessagesRoute, + provider: MessagesProvider, + model: String, + stream: bool, + context: &CallContext<'_, impl Interceptors>, +) -> Result { + let identity = ProviderIdentity { + model, + provider: provider.as_str().to_string(), + }; + let cache = CallCache::::new( + caller_keyed(route), + context.cache, + CacheRequest::from_wire(identity.clone(), None), + ); + context + .result_ready(ExecutionFacts { + provider: identity, + source: ResultSource::Host, + cost: Cost::Deferred, + }) + .await?; + let output = MessagesCallResponse::Complete(MessagesBody::Unchecked(Map::new())); + if stream { + return Ok(output); + } + Ok(cache.finish(output, &ResultSource::Host).await) +} + +fn caller_keyed(route: &MessagesRoute) -> Option<&litellm_cache_response::ScopedCache> { + route + .cache + .as_ref() + .filter(|cache| cache.service.keyed_by_caller()) +} + pub(super) async fn execute( route: &MessagesRoute, provider: MessagesProvider, model: String, call: MessagesCall, secrets: Secrets, + requested_stream: bool, context: &CallContext<'_, impl Interceptors>, ) -> Result { + let caller_keyed = caller_keyed(route).is_some(); let request = provider_call( &route.auth, + route.workload_identity.as_deref(), provider, model, call, @@ -71,38 +191,64 @@ pub(super) async fn execute( .boxed() .await?; let cache = CallCache::::from_wire( - route.cache.as_ref().filter(|_| request.cacheable()), + route + .cache + .as_ref() + .filter(|_| caller_keyed || request.cacheable()), context.cache, &request.identity, &request.wire, ); let identity = request.identity.clone(); - let (output, source) = match cache.lookup().await { + let hit = match caller_keyed { + true => None, + false => cache.lookup().await, + }; + let (output, source) = match hit { Some(hit) => hit, None => ( - call_provider(&route.http, request, context).await?, + call_provider(&route.http, route.transport.as_deref(), request, context).await?, ResultSource::Provider, ), }; - let reported = match &output { + context + .result_ready(execution_facts(provider, identity, &output, &source)) + .await?; + let output = match source { + ResultSource::Provider => { + super::safeguards::stamped(route.safeguards.as_ref(), output).await? + } + ResultSource::Cache { .. } | ResultSource::Host => output, + }; + if !requested_stream && matches!(output, MessagesCallResponse::Stream { .. }) { + return Ok(output); + } + Ok(cache.finish(output, &source).await) +} + +fn execution_facts( + provider: MessagesProvider, + identity: ProviderIdentity, + output: &MessagesCallResponse, + source: &ResultSource, +) -> ExecutionFacts { + let reported = match output { MessagesCallResponse::Complete(response) => match response { MessagesBody::Message(message) => provider.config().reported_cost(message), MessagesBody::Unchecked(_) => None, }, MessagesCallResponse::Stream { .. } => None, }; - context - .result_ready(ExecutionFacts { - provider: identity, - cost: Cost::settle(&source, reported), - source: source.clone(), - }) - .await?; - Ok(cache.finish(output, &source).await) + ExecutionFacts { + provider: identity, + cost: Cost::settle(source, reported), + source: source.clone(), + } } async fn provider_call( auth: &AuthServices, + workload_identity: Option<&dyn WorkloadIdentity>, provider: MessagesProvider, model: String, call: MessagesCall, @@ -129,14 +275,18 @@ async fn provider_call( let sanitized = config.shape_request( MessagesRequest { model, ..body }, - shaping.settings.reasoning_auto_summary, + shaping.settings.shape_options(), )?; - let transform_context = - MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); - let without_listed = - without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; - let trimmed = requested_optional_params(without_listed, provider, &transform_context)?; - let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?; + let transform_context = MessagesTransformContext { + drops_unsupported_params: shaping.settings.drops_unsupported_params, + ..MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params) + }; + let dropped_keys = top_level_paths(&shaping.settings.additional_drop_params); + let without_keys = without_additional_drop_params(sanitized, &dropped_keys)?; + let requested = requested_optional_params(without_keys, provider, &transform_context)?; + let requested_params = serde_json::to_value(&requested.params).map_err(serialize_failure)?; + let trimmed = + without_additional_drop_params(requested, &shaping.settings.additional_drop_params)?; let scoped = get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str()); @@ -144,35 +294,58 @@ async fn provider_call( extra_headers.unwrap_or_default(), scoped, ])))?; - let validated = config.validate_environment( + let validated = match config.validate_federated_environment( forwarded, api_key.as_deref(), &trimmed.model, &trimmed.params, &connection_params, &env_lookup, - )?; + )? { + FederatedEnvironment::Validated(validated) => validated, + FederatedEnvironment::Unauthenticated { + headers, + missing, + federate, + } => match workload_identity { + Some(identity) => match identity + .token(&connection_params, api_base.as_deref()) + .await? + { + Some(token) => federate(headers, token), + None => return Err(missing.into()), + }, + None => return Err(missing.into()), + }, + }; + let tool_custom_field = config.tool_custom_field(&trimmed.model); + let requested_output_format = trimmed.params.output_format.clone(); let body = config.request_body( - config.transform_anthropic_messages_request(trimmed, &transform_context)?, + tool_custom_field.apply( + config.transform_anthropic_messages_request(trimmed, &transform_context)?, + shaping.settings.drops_unsupported_params, + ), &validated.headers, shaping.settings.drop_params, )?; - let environment = ValidatedEnvironment { - headers: config.request_headers( - with_python_default_headers(validated.headers, config.default_headers()), - &body, - ), - auth: validated.auth, - }; - let url = config.get_complete_url( api_base.as_deref(), &body.model, &connection_params, - body.params.stream == Some(true), + body.params.streams(), &env_lookup, )?; + let environment = ValidatedEnvironment { + headers: config.headers_for_url( + config.request_headers( + with_python_default_headers(validated.headers, config.default_headers()), + &with_requested_output_format(&body, requested_output_format), + ), + &url, + ), + auth: validated.auth, + }; let request_context = RequestContext { model: body.model.clone(), @@ -197,19 +370,10 @@ async fn provider_call( request_context, ) .await?; - let stream = match wire.body.get("stream") { - None | Some(Value::Null) => false, - Some(Value::Bool(stream)) => *stream, - Some(value) => { - return Err(Error::invalid_request( - litellm_llms::ErrorDetail::InvalidValue { - field: "stream", - expected: "a boolean", - actual: value.clone(), - }, - )); - } - }; + let stream = wire + .body + .get("stream") + .is_some_and(litellm_llms_types::recognized::python_truthy); Ok(ProviderCall { identity, wire, @@ -217,9 +381,33 @@ async fn provider_call( signer: authenticated.signer, timeout: timeout.or_else(|| shaping.settings.request_timeout()), stream, + whole_sse_frames: shaping.settings.whole_sse_frames, + joined_reads: shaping.settings.joined_reads, + transform_context, + ascii_response_headers: shaping.settings.ascii_response_headers, + unencodable_request_header: shaping.settings.unencodable_request_header, + ascii_request_headers: shaping.settings.ascii_request_headers, }) } +fn with_requested_output_format( + body: &MessagesRequest, + requested: Option, +) -> std::borrow::Cow<'_, MessagesRequest> { + match requested { + Some(output_format) if body.params.output_format.is_none() => { + std::borrow::Cow::Owned(MessagesRequest { + params: MessagesOptionalParams { + output_format: Some(output_format), + ..body.params.clone() + }, + ..body.clone() + }) + } + _ => std::borrow::Cow::Borrowed(body), + } +} + /// Python sends its header dict through `httpx.Headers`, which folds names that differ only in /// case into one lowercase header whose values are joined with `", "`. fn joined_header_values(headers: Vec<(String, String)>) -> Vec<(String, String)> { @@ -278,6 +466,14 @@ fn requested_optional_params( }) } +fn top_level_paths(paths: &[String]) -> Vec { + paths + .iter() + .filter(|path| !path.contains(['.', '['])) + .cloned() + .collect() +} + fn without_additional_drop_params( request: MessagesRequest, paths: &[String], @@ -295,52 +491,57 @@ fn without_additional_drop_params( }) } +enum Upstream { + Native(reqwest::Response), + Hosted(TransportResponse), +} + async fn call_provider( http: &litellm_http::Client, + transport: Option<&dyn Transport>, request: ProviderCall, context: &CallContext<'_, impl Interceptors>, ) -> Result { + let timeout = request + .timeout + .unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); + let response = match send_recovering(http, transport, &request, timeout, context).await? { + Upstream::Native(response) => response, + Upstream::Hosted(response) => return hosted_response(response, request, context).await, + }; let ProviderCall { identity, - wire, provider, - signer, - timeout, stream, + whole_sse_frames, + joined_reads, + .. } = request; let provider_name = provider.as_str(); - log_request_body(provider_name, stream, &wire.body); - let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); - let response = send( - http, - Authenticated { - headers: wire.headers, - signer, - }, - &wire.url, - &wire.body, - timeout, - stream, - ) - .await?; - if !response.status().is_success() { - return Err(provider_error(response, timeout).await); - } let config = provider.config(); + let head = ResponseHead { + status: response.status().as_u16(), + headers: super::encoding::header_pairs(response.headers()), + }; if stream { - return Ok(streaming_response( + context + .raw_response_received(RawResponse::new(String::new()).with_head(head)) + .await?; + let streamed = streaming_response( response, config.stream_decoder(), provider_name, timeout, - )); + joined_reads, + ); + return Ok(if whole_sse_frames { + framed(streamed) + } else { + streamed + }); } - let head = ResponseHead { - status: response.status().as_u16(), - headers: litellm_http::request::header_pairs(response.headers()), - }; let status = head.status; - let text = response.text().await.map_err(litellm_http::Error::from)?; + let text = read_body(response, timeout).await?; log_response_body(&text); context .raw_response_received(RawResponse::new(text.clone()).with_head(head)) @@ -353,6 +554,197 @@ async fn call_provider( decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete) } +async fn hosted_response( + response: TransportResponse, + request: ProviderCall, + context: &CallContext<'_, impl Interceptors>, +) -> Result { + let TransportResponse { + status, + headers, + delivered, + } = response; + let ProviderCall { + identity, + provider, + whole_sse_frames, + .. + } = request; + let provider_name = provider.as_str(); + let config = provider.config(); + let text = match delivered { + Delivered::Complete(text) => text, + Delivered::Stream(chunks) => { + context + .raw_response_received(RawResponse::new(String::new()).with_head(ResponseHead { + status, + headers: headers.clone(), + })) + .await?; + let streamed = + hosted_streaming_response(headers, chunks, config.stream_decoder(), provider_name); + return Ok(if whole_sse_frames { + framed(streamed) + } else { + streamed + }); + } + Delivered::Rejected { error, .. } => return Err(error), + }; + log_response_body(&text); + context + .raw_response_received( + RawResponse::new(text.clone()).with_head(ResponseHead { status, headers }), + ) + .await?; + if serde_json::from_str::(&text).is_err() { + return Err(litellm_http::Error::upstream(status, text).into()); + } + decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete) +} + +fn hosted_request( + authenticated: Authenticated, + url: &str, + body: &Value, + stream: bool, +) -> Result { + if authenticated.signer.is_none() { + return Ok(TransportRequest { + url: url.to_owned(), + headers: authenticated.headers, + body: TransportBody::Json(body.clone()), + stream, + }); + } + let names: Vec = authenticated + .headers + .iter() + .map(|(name, _)| name.clone()) + .collect(); + let signed = raw_leaf_outbound_request(authenticated, url, body, None)?; + let headers = signed + .headers() + .iter() + .filter_map(|(name, value)| { + let spelled = names + .iter() + .find(|known| known.eq_ignore_ascii_case(name.as_str())) + .cloned() + .unwrap_or_else(|| name.to_string()); + Some((spelled, value.to_str().ok()?.to_owned())) + }) + .collect(); + let bytes = signed + .body() + .and_then(reqwest::Body::as_bytes) + .map(Bytes::copy_from_slice) + .unwrap_or_default(); + Ok(TransportRequest { + url: signed.url().to_string(), + headers, + body: TransportBody::Signed(bytes), + stream, + }) +} + +async fn send_recovering( + http: &litellm_http::Client, + transport: Option<&dyn Transport>, + request: &ProviderCall, + timeout: Duration, + context: &CallContext<'_, impl Interceptors>, +) -> Result { + let config = request.provider.config(); + let mut body = std::borrow::Cow::Borrowed(&request.wire.body); + let mut attempt = 1; + loop { + log_request_body(request.provider.as_str(), request.stream, &body); + if let Some(value) = &request.unencodable_request_header { + return Err(Error::transport(UnencodableHeader::Sent(value.clone()))); + } + // The credential header built from the caller's api_key never passes the caller header + // check, so it is checked on the wire headers as Python's client encodes them. + if request.ascii_request_headers + && let Some((_, value)) = request + .wire + .headers + .iter() + .find(|(_, value)| !value.is_ascii()) + { + return Err(Error::transport(UnencodableHeader::Sent( + value.as_bytes().to_vec(), + ))); + } + let authenticated = Authenticated { + headers: request.wire.headers.clone(), + signer: request.signer.clone(), + }; + let response = match transport { + None => Upstream::Native( + send(http, authenticated, &request.wire.url, &body, timeout).await?, + ), + Some(transport) => Upstream::Hosted( + transport + .send(hosted_request( + authenticated, + &request.wire.url, + &body, + request.stream, + )?) + .await?, + ), + }; + if let Upstream::Native(response) = &response + && request.ascii_response_headers + && let Some(unencodable) = UnencodableHeader::received(response.headers()) + { + return Err(Error::transport(unencodable)); + } + let failure = match response { + Upstream::Native(response) if response.status().is_success() => { + return Ok(Upstream::Native(response)); + } + Upstream::Native(response) => { + provider_failure(response, timeout, request.stream).await? + } + Upstream::Hosted(TransportResponse { + status, + headers, + delivered: Delivered::Rejected { text, error }, + }) => { + log_error_body(status, &text); + ProviderFailure { + status, + text, + headers, + raised: Some(error), + } + } + hosted @ Upstream::Hosted(_) => return Ok(hosted), + }; + let retried = if attempt == MESSAGES_HTTP_ERROR_ATTEMPTS { + None + } else { + config.retry_after_http_error( + failure.status, + &failure.text, + &body, + &request.transform_context, + ) + }; + let Some(retried) = retried else { + return Err(failure.into()); + }; + attempt += 1; + tracing::debug!( + "Anthropic /v1/messages: retrying after a recoverable HTTP error (attempt {attempt}/{MESSAGES_HTTP_ERROR_ATTEMPTS})." + ); + body = std::borrow::Cow::Owned(retried.body.clone()); + context.provider_request_retried(retried).await?; + } +} + fn serialize_failure(err: serde_json::Error) -> Error { Error::invalid_request(litellm_llms::ErrorDetail::failed( "Anthropic messages request serialization", @@ -360,27 +752,37 @@ fn serialize_failure(err: serde_json::Error) -> Error { )) } -/// Python's httpx timeout bounds each phase and each read, never the whole exchange. A complete -/// response keeps the client's overall deadline; a stream is bounded until its head arrives and -/// then per read (see [`timed_chunks`]), so a long stream that keeps sending is not cut off. async fn send( http: &litellm_http::Client, authenticated: Authenticated, url: &str, body: &Value, timeout: Duration, - stream: bool, ) -> Result { - let request = outbound_request(authenticated, url, body, (!stream).then_some(timeout))?; - if !stream { - return Ok(litellm_inference::outbound::send(request, http).await?); - } + let request = raw_leaf_outbound_request(authenticated, url, body, None)?; tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http)) .await .map_err(litellm_http::Error::read_timeout)? + .map(super::encoding::decoded) .map_err(Error::from) } +async fn read_body( + mut response: reqwest::Response, + timeout: Duration, +) -> Result { + let expected = response.content_length(); + let mut body = Vec::new(); + while let Some(chunk) = tokio::time::timeout(timeout, response.chunk()) + .await + .map_err(litellm_http::Error::read_timeout)? + .map_err(|error| litellm_http::Error::body_read(error, body.len() as u64, expected))? + { + body.extend_from_slice(&chunk); + } + Ok(String::from_utf8_lossy(&body).into_owned()) +} + /// The stream's body, each read bounded by `timeout` as httpx bounds each read. fn timed_chunks( response: reqwest::Response, @@ -402,30 +804,65 @@ fn timed_chunks( ) } -/// The provider's error, its body read under `timeout`: a stream has no overall deadline once -/// its head arrives, so an error body that stalls would otherwise hold the call forever. -async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error { - let status = response.status().as_u16(); - let headers = litellm_http::request::header_pairs(response.headers()); - match tokio::time::timeout(timeout, response.text()).await { - Ok(Ok(text)) => { - log_error_body(status, &text); - litellm_http::Error::upstream_with_headers(status, text, headers).into() +struct ProviderFailure { + status: u16, + text: String, + headers: Vec<(String, String)>, + raised: Option, +} + +impl From for Error { + fn from(failure: ProviderFailure) -> Self { + match failure.raised { + Some(raised) => raised, + None => litellm_http::Error::upstream_with_headers( + failure.status, + failure.text, + failure.headers, + ) + .into(), } - Ok(Err(error)) => litellm_http::Error::from(error).into(), - Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(), } } +async fn provider_failure( + response: reqwest::Response, + timeout: Duration, + stream: bool, +) -> Result { + let status = response.status().as_u16(); + let headers = super::encoding::header_pairs(response.headers()); + let text = if stream { + tokio::time::timeout( + Duration::from_secs(MESSAGES_STREAM_ERROR_BODY_READ_SECS), + read_body(response, timeout), + ) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default() + } else { + read_body(response, timeout).await? + }; + log_error_body(status, &text); + Ok(ProviderFailure { + status, + text, + headers, + raised: None, + }) +} + fn decode_response( config: &dyn BaseMessagesConfig, model: &str, text: &str, ) -> Result { - let response = match serde_json::from_str::(text) { + let marked = litellm_host::raw_json::mark_wide_numbers(text); + let response = match serde_json::from_str::(&marked) { Ok(response) => response, Err(err) => { - let Ok(mut body) = serde_json::from_str::>(text) else { + let Ok(mut body) = serde_json::from_str::>(&marked) else { return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid( "messages response JSON", err, @@ -435,11 +872,34 @@ fn decode_response( return Ok(MessagesBody::Unchecked(body)); } }; - config + let message = config .transform_anthropic_messages_response(model, response) - .map(|message| MessagesBody::Message(Box::new(message))) .map_err(Error::from) - .map_err(Error::after_send) + .map_err(Error::after_send)?; + if text_has_stop_reason(text) { + return Ok(MessagesBody::Message(Box::new(message))); + } + match serde_json::to_value(message).map_err(serialize_failure)? { + Value::Object(mut body) => { + body.remove("stop_reason"); + Ok(MessagesBody::Unchecked(body)) + } + _ => unreachable!("a message serializes as a JSON object"), + } +} + +#[derive(serde::Deserialize)] +struct StopReasonKey { + #[serde(default, deserialize_with = "key_present")] + stop_reason: bool, +} + +fn key_present<'de, D: serde::Deserializer<'de>>(deserializer: D) -> Result { + >::deserialize(deserializer).map(|_| true) +} + +fn text_has_stop_reason(text: &str) -> bool { + serde_json::from_str::(text).is_ok_and(|key| key.stop_reason) } fn streaming_response( @@ -447,14 +907,15 @@ fn streaming_response( decoder: Option, provider: &'static str, timeout: Duration, + joined: bool, ) -> MessagesCallResponse { - let headers = response - .headers() - .iter() - .filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string()))) - .collect(); - let chunks = timed_chunks(response, timeout) - .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); + let headers = super::encoding::header_pairs(response.headers()); + let chunks = timed_chunks(response, timeout); + let chunks = match (joined, &decoder) { + (true, None) => super::reads::joined_reads(chunks), + _ => chunks.boxed(), + } + .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); let chunks = match decoder { None => chunks.map_err(Error::from).boxed(), Some(decode) => decoded_chunks( @@ -469,6 +930,73 @@ fn streaming_response( } } +fn hosted_streaming_response( + headers: Vec<(String, String)>, + chunks: BoxStream<'static, Result>, + decoder: Option, + provider: &'static str, +) -> MessagesCallResponse { + let chunks = chunks.inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); + let chunks = match decoder { + None => chunks.boxed(), + Some(decode) => decoded_chunks( + chunks.map_err(std::io::Error::other).boxed(), + decode, + provider, + ), + }; + MessagesCallResponse::Stream { + head: super::route::MessagesStreamHead { headers }, + chunks, + } +} + +fn framed(response: MessagesCallResponse) -> MessagesCallResponse { + let MessagesCallResponse::Stream { head, chunks } = response else { + return response; + }; + MessagesCallResponse::Stream { + head, + chunks: whole_sse_frames(chunks), + } +} + +fn whole_sse_frames( + chunks: BoxStream<'static, Result>, +) -> BoxStream<'static, Result> { + futures_util::stream::try_unfold( + (Some(chunks), Vec::new()), + |(chunks, mut pending)| async move { + let Some(mut chunks) = chunks else { + return Ok(None); + }; + while let Some(chunk) = chunks.try_next().await? { + pending.extend_from_slice(&chunk); + let end = complete_frames_end(&pending); + if end > 0 { + let rest = pending.split_off(end); + return Ok(Some((Bytes::from(pending), (Some(chunks), rest)))); + } + } + Ok((!pending.is_empty()).then(|| (Bytes::from(pending), (None, Vec::new())))) + }, + ) + .boxed() +} + +fn complete_frames_end(pending: &[u8]) -> usize { + [b"\r\n\r\n".as_slice(), b"\n\n", b"\r\r"] + .iter() + .filter_map(|delimiter| { + pending + .windows(delimiter.len()) + .rposition(|window| window == *delimiter) + .map(|at| at + delimiter.len()) + }) + .max() + .unwrap_or(0) +} + fn decoded_chunks( bytes: ByteStream, decode: StreamDecoder, @@ -529,6 +1057,14 @@ mod tests { async fn provider_call_with_secrets( call: MessagesCall, secrets: &(dyn Lookup + Send + Sync), + ) -> Result { + provider_call_federating(call, secrets, None).await + } + + async fn provider_call_federating( + call: MessagesCall, + secrets: &(dyn Lookup + Send + Sync), + workload_identity: Option<&dyn WorkloadIdentity>, ) -> Result { let provider = resolve_llm_provider( &call.body.model, @@ -542,6 +1078,7 @@ mod tests { let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default()); provider_call( &AuthServices::default(), + workload_identity, config, model, call, @@ -735,6 +1272,178 @@ mod tests { assert!(error.is_missing_credential(), "{error:?}"); } + struct ScriptedIdentity { + minted: Result, ()>, + asked: std::sync::Mutex, Option)>>, + } + + impl ScriptedIdentity { + fn new(minted: Result, ()>) -> Self { + Self { + minted, + asked: std::sync::Mutex::new(Vec::new()), + } + } + + fn asked(&self) -> Vec<(Option, Option)> { + self.asked.lock().unwrap().clone() + } + } + + impl WorkloadIdentity for ScriptedIdentity { + fn token<'a>( + &'a self, + litellm_params: &'a litellm_router_types::LitellmParams, + api_base: Option<&'a str>, + ) -> futures_util::future::BoxFuture<'a, Result, litellm_auth::Error>> + { + self.asked.lock().unwrap().push(( + litellm_params + .extra + .get("anthropic_federation_rule_id") + .cloned(), + api_base.map(str::to_owned), + )); + let minted = self.minted; + Box::pin(async move { + minted + .map(|token| token.map(SecretValue::new)) + .map_err(|()| litellm_auth::Error::InvalidHeader) + }) + } + } + + fn federated_call( + shaping: MessagesShaping, + api_key: Option<&str>, + extra_headers: Option, + ) -> MessagesCall { + MessagesCall { + body: body( + json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + ), + api_key: api_key.map(str::to_owned), + api_base: Some("https://anthropic.test".into()), + custom_llm_provider: Some("anthropic".into()), + litellm_params: crate::litellm_params( + serde_json::from_value( + json!({"model": "claude-test", "anthropic_federation_rule_id": "fdrl_1"}), + ) + .unwrap(), + ) + .unwrap(), + extra_headers: extra_headers.map(|headers| serde_json::from_value(headers).unwrap()), + provider_specific_header: None, + timeout: None, + shaping, + } + } + + fn sent_headers<'a>(call: &'a ProviderCall, names: &[&str]) -> Vec<(&'a str, &'a str)> { + call.wire + .headers + .iter() + .filter(|(name, _)| names.contains(&name.as_str())) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect() + } + + #[rstest] + #[case::alone(None, &[("anthropic-beta", "oauth-2025-04-20"), ("authorization", "Bearer sk-ant-oat01-minted")])] + #[case::beside_a_caller_beta( + Some(json!({"anthropic-beta": "files-api-2025-04-14"})), + &[("anthropic-beta", "files-api-2025-04-14,oauth-2025-04-20"), ("authorization", "Bearer sk-ant-oat01-minted")] + )] + #[tokio::test] + async fn a_minted_federation_token_is_sent_as_an_oauth_bearer( + shaping: MessagesShaping, + #[case] extra_headers: Option, + #[case] expected: &[(&str, &str)], + ) { + let identity = ScriptedIdentity::new(Ok(Some("sk-ant-oat01-minted"))); + let call = provider_call_federating( + federated_call(shaping, None, extra_headers), + &|_: &str| None, + Some(&identity), + ) + .await + .unwrap(); + let mut sent = sent_headers( + &call, + &[ + "anthropic-beta", + "authorization", + "x-api-key", + "anthropic-dangerous-direct-browser-access", + ], + ); + sent.sort_unstable(); + assert_eq!( + (sent.as_slice(), identity.asked()), + ( + expected, + vec![(Some(json!("fdrl_1")), Some("https://anthropic.test".into()))] + ) + ); + } + + #[rstest] + #[case::static_key(Some("sk-static"), None, &[("x-api-key", "sk-static")])] + #[case::caller_credential( + None, + Some(json!({"x-api-key": "sk-caller"})), + &[("x-api-key", "sk-caller")] + )] + #[tokio::test] + async fn a_configured_credential_is_never_federated( + shaping: MessagesShaping, + #[case] api_key: Option<&str>, + #[case] extra_headers: Option, + #[case] expected: &[(&str, &str)], + ) { + let identity = ScriptedIdentity::new(Ok(Some("sk-ant-oat01-minted"))); + let call = provider_call_federating( + federated_call(shaping, api_key, extra_headers), + &|_: &str| None, + Some(&identity), + ) + .await + .unwrap(); + assert_eq!( + ( + sent_headers(&call, &["authorization", "x-api-key"]), + identity.asked() + ), + (expected.to_vec(), Vec::new()) + ); + } + + #[rstest] + #[case::nothing_minted(Ok(None), true)] + #[case::minting_failed(Err(()), false)] + #[tokio::test] + async fn an_unminted_federation_token_fails_the_call( + shaping: MessagesShaping, + #[case] minted: Result, ()>, + #[case] missing_credential: bool, + ) { + let identity = ScriptedIdentity::new(minted); + let Err(error) = provider_call_federating( + federated_call(shaping, None, None), + &|_: &str| None, + Some(&identity), + ) + .await + else { + panic!("a call without a minted token fails"); + }; + assert_eq!( + error.is_missing_credential(), + missing_credential, + "{error:?}" + ); + } + #[rstest] #[case::nothing_configured(None, None, None)] #[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))] @@ -1037,7 +1746,7 @@ mod tests { .await .unwrap(); assert_eq!(call.wire.url, format!("{base}{path}")); - let signed = outbound_request( + let signed = raw_leaf_outbound_request( Authenticated { headers: call.wire.headers, signer: call.signer, @@ -1081,6 +1790,7 @@ mod tests { Some(streaming::event_stream), "test", Duration::from_secs(5), + false, ) else { panic!("a streaming response returns chunks"); }; @@ -1092,4 +1802,48 @@ mod tests { } assert!(chunks.next().await.is_none()); } + + #[rstest] + #[case::no_delimiter(b"data: a".as_slice(), 0)] + #[case::lf(b"data: a\n\ndata: b", 9)] + #[case::crlf(b"data: a\r\n\r\ndata: b", 11)] + #[case::cr(b"data: a\r\rdata: b", 9)] + #[case::last_of_several(b"data: a\n\ndata: b\n\ndata: c", 18)] + fn complete_frames_end_is_after_the_last_delimiter(#[case] pending: &[u8], #[case] end: usize) { + assert_eq!(complete_frames_end(pending), end); + } + + #[rstest] + #[case::split_frame( + vec![&b"data: a\n"[..], b"\ndata: ", b"b\n\n"], + vec![&b"data: a\n\n"[..], b"data: b\n\n"] + )] + #[case::several_frames_in_one_chunk( + vec![&b"data: a\n\ndata: b\n\nda"[..], b"ta: c\n\n"], + vec![&b"data: a\n\ndata: b\n\n"[..], b"data: c\n\n"] + )] + #[case::unterminated_tail_is_flushed( + vec![&b"data: a\n\nda"[..], b"ta: b"], + vec![&b"data: a\n\n"[..], b"data: b"] + )] + #[tokio::test] + async fn whole_sse_frames_regroups_chunks_at_frame_boundaries( + #[case] chunks: Vec<&'static [u8]>, + #[case] expected: Vec<&'static [u8]>, + ) { + let upstream = futures_util::stream::iter( + chunks + .into_iter() + .map(|chunk| Ok::<_, Error>(Bytes::from_static(chunk))), + ) + .boxed(); + let framed: Vec = whole_sse_frames(upstream).try_collect().await.unwrap(); + assert_eq!( + framed, + expected + .into_iter() + .map(Bytes::from_static) + .collect::>() + ); + } } diff --git a/litellm-rust/crates/inference-messages/src/host_answer.rs b/litellm-rust/crates/inference-messages/src/host_answer.rs new file mode 100644 index 00000000000..07d183406ea --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/host_answer.rs @@ -0,0 +1,11 @@ +use futures_util::future::BoxFuture; + +use crate::Error; + +pub trait HostAnswer: Send + Sync { + fn answer(&self, stream: bool) -> BoxFuture<'_, Result<(), Error>>; + + fn intercept(&self, stream: bool) -> BoxFuture<'_, Result> { + Box::pin(async move { self.answer(stream).await.map(|()| true) }) + } +} diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 25ebeaacb49..cd312ee61b1 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -1,25 +1,41 @@ mod common_utils; mod constants; +mod encoding; mod handler; +mod host_answer; +mod pre_request; +mod reads; pub mod route; +mod safeguards; +mod transport; mod types; +mod workload_identity; use futures_util::FutureExt; use litellm_auth::AuthServices; use litellm_host::interceptors::Interceptors; use litellm_inference::context::CallContext; use litellm_inference::provider::resolve_llm_provider; +use litellm_llms_types::recognized::Recognized; use litellm_secrets::source::SecretSource; use std::sync::Arc; use crate::common_utils::messages_provider; pub use common_utils::merge_messages_header_sources; +pub use encoding::DecodingError; +pub use host_answer::HostAnswer; pub use litellm_inference::Error; +pub use pre_request::PreRequestHooks; +pub use safeguards::Safeguards; +pub use transport::{ + Delivered, Transport, TransportBody, TransportRequest, TransportResponse, UnencodableHeader, +}; pub use types::{ MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, messages_body, }; +pub use workload_identity::WorkloadIdentity; #[derive(Clone)] pub struct MessagesRoute { @@ -27,6 +43,11 @@ pub struct MessagesRoute { auth: Arc, secrets: Arc, cache: Option, + workload_identity: Option>, + transport: Option>, + pre_request_hooks: Option>, + safeguards: Option>, + host_answer: Option>, } impl MessagesRoute { @@ -40,6 +61,51 @@ impl MessagesRoute { auth, secrets, cache: None, + workload_identity: None, + transport: None, + pre_request_hooks: None, + safeguards: None, + host_answer: None, + } + } + + #[must_use] + pub fn with_workload_identity(self, workload_identity: Arc) -> Self { + Self { + workload_identity: Some(workload_identity), + ..self + } + } + + #[must_use] + pub fn with_transport(self, transport: Arc) -> Self { + Self { + transport: Some(transport), + ..self + } + } + + #[must_use] + pub fn with_pre_request_hooks(self, hooks: Arc) -> Self { + Self { + pre_request_hooks: Some(hooks), + ..self + } + } + + #[must_use] + pub fn with_safeguards(self, safeguards: Arc) -> Self { + Self { + safeguards: Some(safeguards), + ..self + } + } + + #[must_use] + pub fn with_host_answer(self, host_answer: Arc) -> Self { + Self { + host_answer: Some(host_answer), + ..self } } @@ -67,7 +133,7 @@ impl MessagesRoute { model = %call.body.model, provider, resolved_model, - stream = call.body.params.stream == Some(true), + stream = call.body.params.streams(), outcome ))] async fn run( @@ -76,24 +142,69 @@ impl MessagesRoute { context: CallContext<'_, impl Interceptors>, ) -> Result { litellm_inference::diagnostic::call(async { - let provider = resolve_llm_provider( - &call.body.model, - call.custom_llm_provider.as_deref(), - "messages", - )?; - let config = messages_provider(provider.provider, provider.model).ok_or_else(|| { - Error::InvalidProvider(<&str>::from(provider.provider).to_string()) - })?; - litellm_inference::diagnostic::provider( - provider.model, - <&'static str>::from(provider.provider), - ); - let model = provider.model.to_owned(); + let requested_stream = call.body.params.streams(); + let client_streams = call.body.params.stream == Some(Recognized::Known(true)); + let call = match &self.pre_request_hooks { + None => call, + Some(hooks) => { + let (config, model, _) = Self::provider(&call)?; + if let Some(output) = handler::cached(self, config, &model, &context).await? { + return Ok(output); + } + hooks.run(call).await? + } + }; + let (config, model, provider) = Self::provider(&call)?; + litellm_inference::diagnostic::provider(&model, provider); + if self.pre_request_hooks.is_none() + && let Some(output) = handler::cached(self, config, &model, &context).await? + { + return Ok(output); + } + if let Some(host_answer) = self + .host_answer + .as_deref() + .filter(|_| call.shaping.settings.intercepted) + && host_answer.intercept(requested_stream).await? + { + return handler::deliver_answer(self, config, model, requested_stream, &context) + .await; + } + if let Some(mock) = call.shaping.settings.mock_response { + return Ok( + handler::mocked(self, config, model, &context, mock, requested_stream).await, + ); + } + if let Some(host_answer) = self + .host_answer + .as_deref() + .filter(|_| call.shaping.settings.mcp_gateway) + { + let stream = call.body.params.streams(); + return handler::answered(self, host_answer, config, model, stream, &context).await; + } let secrets = self.secrets.resolve(config.config().secret_names()).await?; - handler::execute(self, config, model, call, secrets, &context) + handler::execute(self, config, model, call, secrets, client_streams, &context) .boxed() .await }) .await } + + fn provider( + call: &MessagesCall, + ) -> Result<(common_utils::MessagesProvider, String, &'static str), Error> { + let provider = resolve_llm_provider( + &call.body.model, + call.custom_llm_provider.as_deref(), + "messages", + )?; + let config = messages_provider(provider.provider, provider.model) + .ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?; + Ok(( + config, + provider.model.to_owned(), + <&'static str>::from(provider.provider), + )) + } } diff --git a/litellm-rust/crates/inference-messages/src/pre_request.rs b/litellm-rust/crates/inference-messages/src/pre_request.rs new file mode 100644 index 00000000000..949dcd6585b --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/pre_request.rs @@ -0,0 +1,7 @@ +use futures_util::future::BoxFuture; + +use super::{Error, MessagesCall}; + +pub trait PreRequestHooks: Send + Sync { + fn run(&self, call: MessagesCall) -> BoxFuture<'_, Result>; +} diff --git a/litellm-rust/crates/inference-messages/src/reads.rs b/litellm-rust/crates/inference-messages/src/reads.rs new file mode 100644 index 00000000000..0e127a42a13 --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/reads.rs @@ -0,0 +1,394 @@ +use std::{collections::VecDeque, sync::Arc}; + +use bytes::Bytes; +use futures_util::{FutureExt, Stream, StreamExt, stream::BoxStream}; +use std::sync::Mutex; +use tokio::{sync::Notify, task::JoinHandle}; + +const READ_SIZE: usize = 16 * 1024; +const LOW_WATER: usize = 64 * 1024; +const HIGH_WATER: usize = 2 * LOW_WATER; +const SETTLE_YIELDS: usize = 64; + +struct Buffer { + chunks: VecDeque, + size: usize, + quiet: bool, + failure: Option, + ended: bool, +} + +impl Buffer { + fn take(&mut self) -> Bytes { + let mut joined = Vec::new(); + while let Some(mut chunk) = self.chunks.pop_front() { + let room = READ_SIZE - joined.len(); + if chunk.len() > room { + let rest = chunk.split_off(room); + self.chunks.push_front(rest); + } + if joined.is_empty() && (chunk.len() == READ_SIZE || self.chunks.is_empty()) { + self.size -= chunk.len(); + return chunk; + } + joined.extend_from_slice(&chunk); + self.size -= chunk.len(); + if joined.len() == READ_SIZE { + break; + } + } + Bytes::from(joined) + } +} + +struct Reads { + buffer: Mutex>, + received: Notify, + drained: Notify, +} + +enum Read { + Chunk(Bytes), + Failure(E), + End, +} + +impl Reads { + fn readable(&self) -> bool { + let buffer = self.buffer(); + buffer.failure.is_some() + || buffer.ended + || (!buffer.chunks.is_empty() && (buffer.quiet || buffer.size >= READ_SIZE)) + } + + fn quiet(&self) { + self.buffer().quiet = true; + self.received.notify_one(); + } + + fn read(&self) -> Option> { + let mut buffer = self.buffer(); + if !buffer.chunks.is_empty() { + let chunk = buffer.take(); + drop(buffer); + self.drained.notify_one(); + return Some(Read::Chunk(chunk)); + } + if let Some(failure) = buffer.failure.take() { + return Some(Read::Failure(failure)); + } + buffer.ended.then_some(Read::End) + } + + fn buffer(&self) -> std::sync::MutexGuard<'_, Buffer> { + self.buffer + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +struct Reader(JoinHandle<()>); + +impl Drop for Reader { + fn drop(&mut self) { + self.0.abort(); + } +} + +pub(crate) fn joined_reads( + chunks: impl Stream> + Send + 'static, +) -> BoxStream<'static, Result> { + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return chunks.boxed(); + }; + let reads = Arc::new(Reads { + buffer: Mutex::new(Buffer { + chunks: VecDeque::new(), + size: 0, + quiet: false, + failure: None, + ended: false, + }), + received: Notify::new(), + drained: Notify::new(), + }); + let reader = Reader(runtime.spawn(fill(chunks.boxed(), Arc::clone(&reads)))); + futures_util::stream::unfold(Some((reads, reader)), |state| async move { + let (reads, reader) = state?; + loop { + let received = reads.received.notified(); + if reads.readable() { + break; + } + received.await; + } + let read = reads.read()?; + match read { + Read::Chunk(chunk) => Some((Ok(chunk), Some((reads, reader)))), + Read::Failure(failure) => Some((Err(failure), None)), + Read::End => None, + } + }) + .boxed() +} + +async fn fill(mut chunks: BoxStream<'static, Result>, reads: Arc>) { + loop { + let next = match settled_next(&mut chunks).await { + Some(next) => next, + None => { + reads.quiet(); + chunks.next().await + } + }; + let paused = { + let mut buffer = reads.buffer(); + buffer.quiet = false; + match next { + Some(Ok(chunk)) => { + buffer.size += chunk.len(); + if !chunk.is_empty() { + buffer.chunks.push_back(chunk); + } + } + Some(Err(failure)) => { + buffer.failure = Some(failure); + buffer.ended = true; + } + None => buffer.ended = true, + } + if buffer.ended { + drop(buffer); + reads.received.notify_one(); + return; + } + buffer.size > HIGH_WATER + }; + reads.received.notify_one(); + if paused { + reads.quiet(); + loop { + let drained = reads.drained.notified(); + if reads.buffer().size < LOW_WATER { + break; + } + drained.await; + } + } + } +} + +async fn settled_next( + chunks: &mut BoxStream<'static, Result>, +) -> Option>> { + for _ in 0..SETTLE_YIELDS { + if let Some(next) = chunks.next().now_or_never() { + return Some(next); + } + tokio::task::yield_now().await; + } + None +} + +#[cfg(test)] +mod tests { + use std::{ + pin::Pin, + sync::atomic::{AtomicUsize, Ordering}, + task::{Context, Poll}, + time::Duration, + }; + + use futures_util::{FutureExt, StreamExt, TryStreamExt}; + use tokio::sync::mpsc; + + use super::*; + + struct Upstream { + receiver: mpsc::UnboundedReceiver>, + pulled: Arc, + } + + impl Stream for Upstream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let next = self.receiver.poll_recv(cx); + if let Poll::Ready(Some(_)) = next { + self.pulled.fetch_add(1, Ordering::SeqCst); + } + next + } + } + + type Item = Result; + + fn upstream() -> ( + mpsc::UnboundedSender, + BoxStream<'static, Item>, + Arc, + ) { + let (sender, receiver) = mpsc::unbounded_channel(); + let pulled = Arc::new(AtomicUsize::new(0)); + let reads = joined_reads(Upstream { + receiver, + pulled: Arc::clone(&pulled), + }); + (sender, reads, pulled) + } + + async fn settled() { + tokio::time::sleep(Duration::from_millis(20)).await; + } + + #[rstest::rstest] + #[tokio::test] + async fn chunks_that_arrived_together_are_read_as_one() { + let (sender, mut reads, _) = upstream(); + sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap(); + sender.send(Ok(Bytes::from("event: b\n\n"))).unwrap(); + sender.send(Ok(Bytes::from("event: c\n\n"))).unwrap(); + settled().await; + assert_eq!( + reads.try_next().await.unwrap().unwrap(), + "event: a\n\nevent: b\n\nevent: c\n\n" + ); + sender.send(Ok(Bytes::from("event: d\n\n"))).unwrap(); + settled().await; + assert_eq!(reads.try_next().await.unwrap().unwrap(), "event: d\n\n"); + drop(sender); + assert!(reads.try_next().await.unwrap().is_none()); + } + + #[rstest::rstest] + #[tokio::test] + async fn chunks_still_being_handed_over_join_the_read() { + let handed_over = futures_util::stream::unfold(0, |sent| async move { + tokio::task::yield_now().await; + (sent < 50).then(|| (Ok::<_, &str>(Bytes::from("event: ping\n\n")), sent + 1)) + }); + let read: Vec = joined_reads(handed_over) + .map_ok(|chunk| chunk.len()) + .try_collect() + .await + .unwrap(); + assert_eq!(read, [50 * "event: ping\n\n".len()]); + } + + #[rstest::rstest] + #[tokio::test] + async fn a_read_returns_the_first_chunk_without_waiting_for_more() { + let (sender, mut reads, _) = upstream(); + let first = tokio::spawn(async move { reads.try_next().await }); + settled().await; + sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap(); + assert_eq!(first.await.unwrap().unwrap().unwrap(), "event: a\n\n"); + } + + #[rstest::rstest] + #[tokio::test] + async fn reads_are_cut_at_sixteen_kibibytes() { + let (sender, reads, _) = upstream(); + sender.send(Ok(Bytes::from(vec![b'a'; 10_000]))).unwrap(); + sender.send(Ok(Bytes::from(vec![b'b'; 30_000]))).unwrap(); + drop(sender); + settled().await; + let sizes: Vec = reads + .map_ok(|chunk| chunk.len()) + .try_collect() + .await + .unwrap(); + assert_eq!(sizes, [READ_SIZE, READ_SIZE, 40_000 - 2 * READ_SIZE]); + } + + #[rstest::rstest] + #[tokio::test] + async fn the_bytes_still_buffered_are_read_before_a_failure() { + let (sender, mut reads, _) = upstream(); + sender.send(Ok(Bytes::from("event: a\n\n"))).unwrap(); + sender.send(Ok(Bytes::from("event: b\n\n"))).unwrap(); + sender.send(Err("reset")).unwrap(); + settled().await; + assert_eq!( + reads.try_next().await.unwrap().unwrap(), + "event: a\n\nevent: b\n\n" + ); + assert_eq!(reads.try_next().await.unwrap_err(), "reset"); + assert!(reads.next().await.is_none()); + } + + #[rstest::rstest] + #[tokio::test] + async fn reading_pauses_above_the_high_water_mark_until_below_the_low_one() { + let (sender, mut reads, pulled) = upstream(); + std::iter::repeat_with(|| Bytes::from(vec![b'a'; READ_SIZE])) + .take(12) + .try_for_each(|chunk| sender.send(Ok(chunk))) + .unwrap(); + settled().await; + assert_eq!(pulled.load(Ordering::SeqCst), 9); + let read: Vec = (&mut reads).take(5).try_collect().await.unwrap(); + assert_eq!(read.len(), 5); + settled().await; + assert_eq!(pulled.load(Ordering::SeqCst), 9); + reads.try_next().await.unwrap(); + settled().await; + assert_eq!(pulled.load(Ordering::SeqCst), 12); + } + + async fn chunked_upstream(frames: usize) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 4096]; + let _ = socket.read(&mut request).await.unwrap(); + let frame = "event: ping\ndata: {\"type\": \"ping\"}\n\n"; + let body: String = + std::iter::repeat_n(format!("{:x}\r\n{frame}\r\n", frame.len()), frames).collect(); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\nconnection: close\r\n\r\n{body}0\r\n\r\n" + ); + socket.write_all(response.as_bytes()).await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + }); + format!("http://{address}") + } + + #[rstest::rstest] + #[case::fifty(50)] + #[case::a_thousand(1000)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn http_chunks_written_together_are_read_together(#[case] frames: usize) { + let url = chunked_upstream(frames).await; + let response = litellm_http::Client::for_test() + .get(url) + .send() + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + let sizes: Vec = joined_reads(response.bytes_stream()) + .map_ok(|chunk| chunk.len()) + .try_collect() + .await + .unwrap(); + let total = frames * "event: ping\ndata: {\"type\": \"ping\"}\n\n".len(); + let expected: Vec = std::iter::repeat_n(READ_SIZE, total / READ_SIZE) + .chain((!total.is_multiple_of(READ_SIZE)).then_some(total % READ_SIZE)) + .collect(); + assert_eq!(sizes, expected); + } + + #[rstest::rstest] + fn without_a_runtime_the_chunks_pass_through() { + let chunks = + futures_util::stream::iter([Ok::<_, &str>(Bytes::from("a")), Ok(Bytes::from("b"))]); + let read: Vec = joined_reads(chunks) + .try_collect() + .now_or_never() + .unwrap() + .unwrap(); + assert_eq!(read, [Bytes::from("a"), Bytes::from("b")]); + } +} diff --git a/litellm-rust/crates/inference-messages/src/route.rs b/litellm-rust/crates/inference-messages/src/route.rs index 5f06d16d49f..fa850b23bf2 100644 --- a/litellm-rust/crates/inference-messages/src/route.rs +++ b/litellm-rust/crates/inference-messages/src/route.rs @@ -74,4 +74,103 @@ impl litellm_inference::caching::StreamCachable for Messages { fn bytes(chunk: &Self::Chunk) -> &[u8] { chunk.as_ref() } + + fn storable_stream(text: &str) -> bool { + let lines: Vec<&str> = sse_lines(text).collect(); + lines.contains(&"event: message_stop") + && !carries_provider_error(&lines) + && lines.iter().any(|line| line_carries_output(line)) + } +} + +fn sse_lines(text: &str) -> impl Iterator { + let mut rest = Some(text); + std::iter::from_fn(move || { + let current = rest?; + match current.find(['\r', '\n']) { + Some(at) => { + let skip = if current[at..].starts_with("\r\n") { + 2 + } else { + 1 + }; + rest = Some(¤t[at + skip..]); + Some(¤t[..at]) + } + None => { + rest = None; + (!current.is_empty()).then_some(current) + } + } + }) +} + +fn sse_data(line: &str, stripped: bool) -> Option> { + let data = line.strip_prefix("data:")?; + let data = if stripped { + data.trim_matches(|c: char| matches!(c, ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c')) + } else { + data + }; + match serde_json::from_str(data) { + Ok(serde_json::Value::Object(map)) => Some(map), + _ => None, + } +} + +fn carries_provider_error(lines: &[&str]) -> bool { + lines + .iter() + .filter_map(|line| sse_data(line, true)) + .find(|event| event.get("type").and_then(serde_json::Value::as_str) == Some("error")) + .is_some_and(|event| event.get("error").is_some_and(serde_json::Value::is_object)) +} + +fn line_carries_output(line: &str) -> bool { + let Some(event) = sse_data(line, false) else { + return false; + }; + let block = match event.get("type").and_then(serde_json::Value::as_str) { + Some("content_block_start") => event.get("content_block"), + Some("content_block_delta") => event.get("delta"), + _ => None, + }; + block + .and_then(serde_json::Value::as_object) + .is_some_and(|fields| { + fields.iter().any(|(name, value)| { + name != "type" + && !matches!(value, serde_json::Value::Null) + && value.as_str() != Some("") + && value.as_array().is_none_or(|items| !items.is_empty()) + }) + }) +} + +#[cfg(test)] +mod tests { + use litellm_inference::caching::StreamCachable; + use rstest::rstest; + + use super::Messages; + + const TEXT: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Paris\"}}\n\n"; + const EMPTY: &str = "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"\"}}\n\n"; + const STOP: &str = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + + #[rstest] + #[case::complete(format!("{TEXT}{STOP}"), true)] + #[case::crlf_lines(format!("{TEXT}{STOP}").replace('\n', "\r\n"), true)] + #[case::done_sentinel(format!("{TEXT}{STOP}data: [DONE]\n\n"), true)] + #[case::data_only_frames(format!("{TEXT}{STOP}").replace("event: message_stop\n", ""), false)] + #[case::event_without_space(format!("{TEXT}{STOP}").replace("event: message_stop", "event:message_stop"), false)] + #[case::provider_error(format!("{TEXT}{STOP}event: error\ndata: {{\"type\":\"error\",\"error\":{{\"type\":\"overloaded_error\"}}}}\n\n"), false)] + #[case::error_without_body(format!("{TEXT}{STOP}data: {{\"type\":\"error\",\"error\":\"overloaded\"}}\n\n"), true)] + #[case::no_output(format!("{EMPTY}{STOP}"), false)] + fn a_stream_is_stored_by_python_s_message_stop_error_and_output_tests( + #[case] text: String, + #[case] stored: bool, + ) { + assert_eq!(Messages::storable_stream(&text), stored); + } } diff --git a/litellm-rust/crates/inference-messages/src/safeguards.rs b/litellm-rust/crates/inference-messages/src/safeguards.rs new file mode 100644 index 00000000000..78e4a752536 --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/safeguards.rs @@ -0,0 +1,248 @@ +use std::{collections::VecDeque, sync::Arc}; + +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, future::BoxFuture, stream::BoxStream}; +use serde_json::Value; + +use super::{Error, MessagesBody, MessagesCallResponse}; + +const SAFEGUARD_RESULTS: &str = "safeguard_results"; + +pub trait Safeguards: Send + Sync { + fn results<'a>( + &'a self, + response: &'a MessagesBody, + ) -> BoxFuture<'a, Result, Error>>; + + fn open_stream(&self) -> BoxFuture<'_, Result>; + + fn frames(&self, chunk: Bytes) -> BoxFuture<'_, Result, Error>>; + + fn flush(&self) -> BoxFuture<'_, Result, Error>>; +} + +pub(super) async fn stamped( + safeguards: Option<&Arc>, + output: MessagesCallResponse, +) -> Result { + let Some(safeguards) = safeguards else { + return Ok(output); + }; + match output { + MessagesCallResponse::Complete(response) => { + let results = safeguards.results(&response).await?; + Ok(MessagesCallResponse::Complete(match results { + Some(results) => with_results(response, results), + None => response, + })) + } + MessagesCallResponse::Stream { head, chunks } => { + if !safeguards.open_stream().await? { + return Ok(MessagesCallResponse::Stream { head, chunks }); + } + Ok(MessagesCallResponse::Stream { + head, + chunks: stamped_frames(Arc::clone(safeguards), chunks), + }) + } + } +} + +fn with_results(response: MessagesBody, results: Value) -> MessagesBody { + match response { + MessagesBody::Message(mut message) => { + message.extra.insert(SAFEGUARD_RESULTS.to_owned(), results); + MessagesBody::Message(message) + } + MessagesBody::Unchecked(mut body) => { + body.insert(SAFEGUARD_RESULTS.to_owned(), results); + MessagesBody::Unchecked(body) + } + } +} + +struct Framing { + chunks: Option>>, + safeguards: Arc, + ready: VecDeque, +} + +fn stamped_frames( + safeguards: Arc, + chunks: BoxStream<'static, Result>, +) -> BoxStream<'static, Result> { + let framing = Framing { + chunks: Some(chunks), + safeguards, + ready: VecDeque::new(), + }; + futures_util::stream::try_unfold(framing, |mut framing| async move { + loop { + if let Some(frame) = framing.ready.pop_front() { + return Ok(Some((frame, framing))); + } + let Some(mut chunks) = framing.chunks.take() else { + return Ok(None); + }; + let frames = match chunks.try_next().await? { + Some(chunk) => { + framing.chunks = Some(chunks); + framing.safeguards.frames(chunk).await? + } + None => framing.safeguards.flush().await?, + }; + framing.ready.extend(frames); + } + }) + .boxed() +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use litellm_llms_types::formats::messages::MessagesResponse; + use rstest::rstest; + use serde_json::{Map, json}; + + use super::*; + use crate::route::MessagesStreamHead; + + struct Scripted { + results: Option, + streams: bool, + fed: Mutex>, + } + + impl Scripted { + fn new(results: Option, streams: bool) -> Arc { + Arc::new(Self { + results, + streams, + fed: Mutex::new(Vec::new()), + }) + } + } + + impl Safeguards for Scripted { + fn results<'a>( + &'a self, + _: &'a MessagesBody, + ) -> BoxFuture<'a, Result, Error>> { + Box::pin(async move { Ok(self.results.clone()) }) + } + + fn open_stream(&self) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(self.streams) }) + } + + fn frames(&self, chunk: Bytes) -> BoxFuture<'_, Result, Error>> { + Box::pin(async move { + self.fed.lock().unwrap().push(chunk.clone()); + Ok(chunk + .split_inclusive(|byte| *byte == b';') + .map(Bytes::copy_from_slice) + .collect()) + }) + } + + fn flush(&self) -> BoxFuture<'_, Result, Error>> { + Box::pin(async move { Ok(vec![Bytes::from_static(b"tail")]) }) + } + } + + fn message() -> MessagesBody { + let response: MessagesResponse = serde_json::from_value(json!({ + "id": "msg_1", "type": "message", "role": "assistant", "model": "m", + "content": [], "stop_reason": "end_turn", "stop_sequence": null, + })) + .unwrap(); + MessagesBody::Message(Box::new(response)) + } + + fn stream(chunks: &[&'static [u8]]) -> MessagesCallResponse { + let chunks: Vec> = chunks + .iter() + .map(|chunk| Ok(Bytes::from_static(chunk))) + .collect(); + MessagesCallResponse::Stream { + head: MessagesStreamHead { + headers: Vec::new(), + }, + chunks: futures_util::stream::iter(chunks).boxed(), + } + } + + async fn collected(output: MessagesCallResponse) -> Vec { + let MessagesCallResponse::Stream { chunks, .. } = output else { + panic!("expected a stream"); + }; + chunks.try_collect().await.unwrap() + } + + #[rstest] + #[case::message_gains_the_results(message(), Some(json!([{"type": "dangerous_tool_use"}])))] + #[case::unchecked_gains_the_results( + MessagesBody::Unchecked(Map::new()), + Some(json!([{"type": "dangerous_tool_use"}])) + )] + #[case::declined_leaves_it(message(), None)] + #[tokio::test] + async fn a_complete_response_carries_what_the_classifier_answered( + #[case] response: MessagesBody, + #[case] results: Option, + ) { + let safeguards: Arc = Scripted::new(results.clone(), false); + let MessagesCallResponse::Complete(stamped) = stamped( + Some(&safeguards), + MessagesCallResponse::Complete(response.clone()), + ) + .await + .unwrap() else { + panic!("expected a complete response"); + }; + let body = serde_json::to_value(&stamped).unwrap(); + assert_eq!(body.get(SAFEGUARD_RESULTS), results.as_ref()); + let mut unstamped = body; + unstamped.as_object_mut().unwrap().remove(SAFEGUARD_RESULTS); + assert_eq!(unstamped, serde_json::to_value(&response).unwrap()); + } + + #[tokio::test] + async fn a_stream_is_reframed_through_the_classifier_and_flushed_at_its_end() { + let scripted = Scripted::new(None, true); + let safeguards: Arc = scripted.clone(); + let frames = collected( + stamped(Some(&safeguards), stream(&[b"a;b", b";c;"])) + .await + .unwrap(), + ) + .await; + assert_eq!( + frames, + ["a;", "b", ";", "c;", "tail"] + .map(|frame| Bytes::from_static(frame.as_bytes())) + .to_vec() + ); + assert_eq!(scripted.fed.lock().unwrap().len(), 2); + } + + #[rstest] + #[case::declined(Some(Scripted::new(None, false) as Arc))] + #[case::not_classified(None)] + #[tokio::test] + async fn a_stream_the_classifier_does_not_answer_passes_through( + #[case] safeguards: Option>, + ) { + let frames = collected( + stamped(safeguards.as_ref(), stream(&[b"a;b", b";c;"])) + .await + .unwrap(), + ) + .await; + assert_eq!( + frames, + [b"a;b".as_slice(), b";c;"].map(Bytes::from_static).to_vec() + ); + } +} diff --git a/litellm-rust/crates/inference-messages/src/transport.rs b/litellm-rust/crates/inference-messages/src/transport.rs new file mode 100644 index 00000000000..1353592447b --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/transport.rs @@ -0,0 +1,59 @@ +use bytes::Bytes; +use futures_util::{future::BoxFuture, stream::BoxStream}; +use serde_json::Value; + +use crate::Error; + +pub enum TransportBody { + Json(Value), + Signed(Bytes), +} + +pub struct TransportRequest { + pub url: String, + pub headers: Vec<(String, String)>, + pub body: TransportBody, + pub stream: bool, +} + +pub enum Delivered { + Complete(String), + Stream(BoxStream<'static, Result>), + Rejected { text: String, error: Error }, +} + +pub struct TransportResponse { + pub status: u16, + pub headers: Vec<(String, String)>, + pub delivered: Delivered, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum UnencodableHeader { + Sent(Vec), + Received(Vec), +} + +impl std::fmt::Display for UnencodableHeader { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("a header value is not ASCII, which Python's HTTP client cannot encode") + } +} + +impl std::error::Error for UnencodableHeader {} + +impl UnencodableHeader { + pub(crate) fn received(headers: &reqwest::header::HeaderMap) -> Option { + headers + .values() + .find(|value| !value.as_bytes().is_ascii()) + .map(|value| Self::Received(value.as_bytes().to_vec())) + } +} + +pub trait Transport: Send + Sync { + fn send( + &self, + request: TransportRequest, + ) -> BoxFuture<'static, Result>; +} diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 772ea6f7292..9dbc171ecb0 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -2,7 +2,10 @@ use std::time::Duration; use bytes::Bytes; use litellm_host::call::CallOutput; -use litellm_llms::{ErrorDetail, base_llm::messages::context::MessagesModelCapabilities}; +use litellm_llms::{ + ErrorDetail, + base_llm::messages::context::{MessagesModelCapabilities, MessagesShapeOptions}, +}; use litellm_llms_types::{ formats::messages::{MessagesRequest, MessagesResponse}, headers::ProviderSpecificHeaders, @@ -71,9 +74,38 @@ pub struct MessagesSettings { pub additional_drop_params: Vec, #[serde(default)] pub request_timeout: Option, + #[serde(default)] + pub mock_response: Option>, + #[serde(default)] + pub whole_sse_frames: bool, + #[serde(default)] + pub joined_reads: bool, + #[serde(default)] + pub presanitized: bool, + #[serde(default)] + pub drops_unsupported_params: bool, + #[serde(default)] + pub intercepted: bool, + #[serde(default)] + pub mcp_gateway: bool, + #[serde(default)] + pub ascii_response_headers: bool, + #[serde(default)] + pub unencodable_request_header: Option>, + /// Python's HTTP client encodes every header value it sends as ASCII, so a value it cannot + /// encode fails the call before the send. + #[serde(default)] + pub ascii_request_headers: bool, } impl MessagesSettings { + pub fn shape_options(&self) -> MessagesShapeOptions { + MessagesShapeOptions { + reasoning_auto_summary: self.reasoning_auto_summary, + presanitized: self.presanitized, + } + } + pub fn request_timeout(&self) -> Option { self.request_timeout .and_then(|secs| Duration::try_from_secs_f64(secs).ok()) @@ -140,7 +172,8 @@ mod tests { "drop_params": true, "reasoning_auto_summary": true, "additional_drop_params": ["metadata.user_id", "thinking"], - "request_timeout": 1800.0 + "request_timeout": 1800.0, + "drops_unsupported_params": true }), MessagesShaping { settings: MessagesSettings { @@ -148,6 +181,8 @@ mod tests { reasoning_auto_summary: true, additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], request_timeout: Some(1800.0), + drops_unsupported_params: true, + ..MessagesSettings::default() }, capabilities: MessagesModelCapabilities { supports_reasoning: true, @@ -185,6 +220,10 @@ mod tests { serialized["reasoning_auto_summary"], json!(expected.settings.reasoning_auto_summary) ); + assert_eq!( + serialized["drops_unsupported_params"], + json!(expected.settings.drops_unsupported_params) + ); assert_eq!( serialized["additional_drop_params"], json!(expected.settings.additional_drop_params) diff --git a/litellm-rust/crates/inference-messages/src/workload_identity.rs b/litellm-rust/crates/inference-messages/src/workload_identity.rs new file mode 100644 index 00000000000..56445de6eea --- /dev/null +++ b/litellm-rust/crates/inference-messages/src/workload_identity.rs @@ -0,0 +1,11 @@ +use futures_util::future::BoxFuture; +use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; + +pub trait WorkloadIdentity: Send + Sync { + fn token<'a>( + &'a self, + litellm_params: &'a LitellmParams, + api_base: Option<&'a str>, + ) -> BoxFuture<'a, Result, litellm_auth::Error>>; +} diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 787501a8fed..ff83f4e41e1 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -301,3 +301,135 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ assert!(matches!(cached.source, ResultSource::Cache { .. })); assert_eq!(cached.cost, Cost::CacheHit); } + +#[derive(Default)] +struct CallerKeyed { + config: ResponseCacheConfig, + entry: std::sync::Mutex>, +} + +impl ResponseCacheService for CallerKeyed { + fn config(&self) -> &ResponseCacheConfig { + &self.config + } + + fn keyed_by_caller(&self) -> bool { + true + } + + fn lookup<'a>( + &'a self, + _: &'a litellm_cache_response::ResponseCacheRequest, + _: Duration, + ) -> std::pin::Pin< + Box< + dyn std::future::Future, litellm_cache::Error>> + + Send + + 'a, + >, + > { + Box::pin(async move { Ok(self.entry.lock().unwrap().clone()) }) + } + + fn store<'a>( + &'a self, + _: &'a litellm_cache_response::ResponseCacheRequest, + value: Value, + _: Duration, + ) -> std::pin::Pin< + Box> + Send + 'a>, + > { + Box::pin(async move { + *self.entry.lock().unwrap() = Some(value); + Ok(()) + }) + } +} + +#[tokio::test] +async fn a_caller_keyed_hit_is_served_before_the_request_is_built_or_intercepted() { + use litellm_inference_messages::{MessagesCall, MessagesCallResponse}; + + let response = json!({"id": "msg_caller", "type": "message", "role": "assistant", "model": "native-test", + "content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 3, "output_tokens": 2}}); + let upstream = support::upstream([support::json_response(response)]).await; + let hooks = ChangingHooks::default(); + let route = support::messages_route(litellm_inference_testing::no_secrets()).with_cache( + ScopedCache::new(Arc::new(CallerKeyed::default()), CacheScope::Shared), + ); + let request = |api_key: Option<&str>| { + MessagesCall { + body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(), + api_key: api_key.map(Into::into), api_base: Some(upstream.uri()), custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, + timeout: None, shaping: Default::default(), + } + }; + let first = route + .execute(request(Some("key")), &hooks, None) + .await + .unwrap(); + let second = route.execute(request(None), &hooks, None).await.unwrap(); + let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) = + (first, second) + else { + panic!("expected Messages responses"); + }; + assert_eq!(first, second); + assert_eq!(support::received(&upstream).await.len(), 1); + assert_eq!(hooks.calls.load(Ordering::SeqCst), 1); + let facts = hooks.facts.lock().unwrap(); + let [provider, cached] = facts.as_slice() else { + panic!("expected provider and cache results"); + }; + assert_eq!(provider.source, ResultSource::Provider); + assert!(matches!(cached.source, ResultSource::Cache { .. })); + assert_eq!(cached.cost, Cost::CacheHit); +} + +#[rstest] +#[case::unary(false, true)] +#[case::streamed(true, false)] +#[tokio::test] +async fn a_mock_is_stored_under_the_caller_key_only_when_unstreamed( + #[case] stream: bool, + #[case] stored: bool, +) { + use litellm_inference_messages::{ + MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, + }; + + let mock = json!({"id": "msg_mock", "type": "message", "role": "assistant", "model": "native-test", + "content": [{"type": "text", "text": "mocked"}], "stop_reason": "end_turn", "stop_sequence": null, + "usage": {"input_tokens": 3, "output_tokens": 2}}); + let upstream = support::upstream([]).await; + let hooks = ChangingHooks::default(); + let service = Arc::new(CallerKeyed::default()); + let route = support::messages_route(litellm_inference_testing::no_secrets()) + .with_cache(ScopedCache::new(service.clone(), CacheScope::Shared)); + let call = MessagesCall { + body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": stream})).unwrap(), + api_key: Some("key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, + timeout: None, + shaping: MessagesShaping { + settings: MessagesSettings { + mock_response: serde_json::from_value(mock.clone()).unwrap(), + ..Default::default() + }, + ..Default::default() + }, + }; + + let output = route.execute(call, &hooks, None).await.unwrap(); + + let MessagesCallResponse::Complete(MessagesBody::Unchecked(body)) = output else { + panic!("expected the mock as a whole response"); + }; + assert_eq!(Value::Object(body), mock); + assert!(support::received(&upstream).await.is_empty()); + assert_eq!(hooks.calls.load(Ordering::SeqCst), 0); + assert!(hooks.facts.lock().unwrap().is_empty()); + assert_eq!(service.entry.lock().unwrap().is_some(), stored); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/gateway.rs b/litellm-rust/crates/inference-messages/tests/messages/gateway.rs new file mode 100644 index 00000000000..d6e591c723c --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/messages/gateway.rs @@ -0,0 +1,99 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use futures_util::future::BoxFuture; +use litellm_inference_messages::{HostAnswer, MessagesRoute}; +use rstest::rstest; + +use super::*; + +#[derive(Default)] +struct CountingGateway(AtomicUsize); + +impl HostAnswer for CountingGateway { + fn answer(&self, _stream: bool) -> BoxFuture<'_, Result<(), Error>> { + self.0.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(()) }) + } +} + +fn gateway_call(call: MessagesCall, api_base: String, mcp_gateway: bool) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant-test".into()), + api_base: Some(api_base), + shaping: MessagesShaping { + settings: MessagesSettings { + mcp_gateway, + ..MessagesSettings::default() + }, + ..MessagesShaping::default() + }, + ..call + } +} + +async fn run_through(route: MessagesRoute, call: MessagesCall) -> Result { + let host = LocalMessagesHost::new(call); + litellm_host_native::in_process::run_hosted( + route.machine(host.request()?, None), + host.runtime(), + ) + .await +} + +#[rstest] +#[tokio::test] +async fn a_call_the_host_hands_to_the_mcp_gateway_never_reaches_the_provider(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let gateway = Arc::new(CountingGateway::default()); + let route = + messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone()); + + let output = run_through(route, gateway_call(call, upstream.uri(), true)) + .await + .unwrap(); + + let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else { + panic!("the host answers in place of the provider"); + }; + assert!(body.is_empty()); + assert_eq!(gateway.0.load(Ordering::SeqCst), 1); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_call_the_host_keeps_reaches_the_provider_past_the_gateway(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let gateway = Arc::new(CountingGateway::default()); + let route = + messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone()); + + let output = run_through(route, gateway_call(call, upstream.uri(), false)) + .await + .unwrap(); + + assert!(matches!( + output, + MessagesOutput::Complete(MessagesBody::Message(_)) + )); + assert_eq!(gateway.0.load(Ordering::SeqCst), 0); + only_request(&upstream).await; +} + +#[rstest] +#[tokio::test] +async fn a_mock_answers_before_the_gateway(call: MessagesCall) { + let gateway = Arc::new(CountingGateway::default()); + let route = + messages_route(Arc::new(RecordingSecrets::empty())).with_host_answer(gateway.clone()); + let mut mocked = gateway_call(call, UNREACHABLE_BASE.into(), true); + mocked.shaping.settings.mock_response = Some(object(json!({"mocked": true}))); + + let output = run_through(route, mocked).await.unwrap(); + + let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else { + panic!("the mock is the call's response"); + }; + assert_eq!(body, object(json!({"mocked": true}))); + assert_eq!(gateway.0.load(Ordering::SeqCst), 0); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index 063e9d18738..7254830a2b0 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -3,7 +3,7 @@ use litellm_host::lifecycle::ExecutionEvent; use std::sync::Mutex; use litellm_host::{ - interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest}, + interceptors::{ExecutionFacts, RawResponse, RequestContext, ResultSource, WireRequest}, lifecycle::CallEvent, }; use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages}; @@ -46,19 +46,23 @@ impl RecordingHost { Self::new(call, Box::new(Ok)) } - fn raw_responses(&self) -> Vec { + fn raw_heads(&self) -> Vec { self.events .lock() .unwrap() .iter() .filter_map(|event| match event { CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => { - Some(raw.body.clone()) + Some(raw.clone()) } _ => None, }) .collect() } + + fn raw_responses(&self) -> Vec { + self.raw_heads().into_iter().map(|raw| raw.body).collect() + } } impl RecordingHost { @@ -137,7 +141,7 @@ async fn rejected_results_are_not_delivered_or_cached( let response = if streaming { ResponseTemplate::new(200).set_body_raw( - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hi\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", "text/event-stream", ) } else { @@ -291,22 +295,23 @@ async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCa } #[rstest] -#[case::enable(false, json!(true), Some(true))] -#[case::disable(true, json!(false), Some(false))] -#[case::null(true, Value::Null, Some(false))] -#[case::invalid(false, json!("true"), None)] +#[case::enable(false, json!(true), true)] +#[case::disable(true, json!(false), false)] +#[case::null(true, Value::Null, false)] +#[case::truthy_string(false, json!("false"), true)] +#[case::falsy_integer(true, json!(0), false)] #[tokio::test] async fn response_mode_follows_the_intercepted_request( call: MessagesCall, traces: TraceCapture, #[case] original_stream: bool, #[case] rewritten_stream: Value, - #[case] expected_stream: Option, + #[case] expected_stream: bool, ) { use futures_util::TryStreamExt; let sse = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; - let response = if expected_stream == Some(true) { + let response = if expected_stream { ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream") } else { message_response() @@ -332,14 +337,14 @@ async fn response_mode_follows_the_intercepted_request( .await?; match output { MessagesCallResponse::Stream { chunks, .. } => { - assert_eq!(expected_stream, Some(true)); + assert!(expected_stream); assert_eq!( chunks.try_collect::>().await?.concat(), sse.as_bytes() ); } MessagesCallResponse::Complete(message) => { - assert_eq!(expected_stream, Some(false)); + assert!(!expected_stream); assert_eq!( message, MessagesBody::Message(serde_json::from_value(message_body()).unwrap()) @@ -352,18 +357,12 @@ async fn response_mode_follows_the_intercepted_request( let summaries = traces.summaries("litellm.route"); assert_eq!(summaries.len(), 1); - let Some(expected_stream) = expected_stream else { - assert!(result.unwrap_err().is_request()); - assert!(received(&upstream).await.is_empty()); - assert_eq!(summaries[0]["outcome"], "failure"); - return; - }; result.unwrap(); assert_eq!( only_request(&upstream).await.json()["stream"], rewritten_stream ); - assert_eq!(host.raw_responses().len(), usize::from(!expected_stream)); + assert_eq!(host.raw_responses().len(), 1); assert_eq!(summaries[0]["stream"], expected_stream); assert_eq!(summaries[0]["outcome"], "success"); } @@ -405,14 +404,9 @@ async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall) } #[rstest] -#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))] -#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))] #[tokio::test] -async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( - call: MessagesCall, - #[case] response: ResponseTemplate, -) { - let upstream = upstream([response]).await; +async fn no_raw_response_is_emitted_for_a_failure(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(500).set_body_string("boom")]).await; let host = RecordingHost::passthrough(authenticated( with_fields(call, json!({"stream": true})), upstream.uri(), @@ -424,6 +418,32 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure( assert!(host.raw_responses().is_empty()); } +#[rstest] +#[tokio::test] +async fn a_stream_emits_its_head_without_a_body_before_any_chunk(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(200) + .insert_header("request-id", "req_stream") + .set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream")]) + .await; + let host = RecordingHost::passthrough(authenticated( + with_fields(call, json!({"stream": true})), + upstream.uri(), + )); + + let _ = run_through(&host).await; + + let [raw] = <[RawResponse; 1]>::try_from(host.raw_heads()) + .unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len())); + assert_eq!(raw.body, ""); + let head = raw.head.expect("the stream's head is forwarded"); + assert_eq!(head.status, 200); + assert!( + head.headers + .iter() + .any(|(name, value)| name == "request-id" && value == "req_stream") + ); +} + /// Python logs `optional_params` after its requested-param gate, so a dropped param must /// not resurface in callbacks. #[rstest] @@ -456,6 +476,38 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages assert_eq!(optional_params, json!({"max_tokens": 16})); } +#[rstest] +#[tokio::test] +async fn the_request_context_keeps_the_nested_fields_additional_drop_params_strip( + call: MessagesCall, +) { + let upstream = upstream([message_response()]).await; + let tool = json!({"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}); + let host = RecordingHost::passthrough(authenticated( + MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + additional_drop_params: vec!["tools[*].input_examples".into()], + ..MessagesSettings::default() + }, + ..MessagesShaping::default() + }, + ..with_fields(call, json!({"tools": [tool.clone()]})) + }, + upstream.uri(), + )); + + run_through(&host).await.expect("messages call succeeds"); + + let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap()) + .unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len())); + assert_eq!(optional_params["tools"], json!([tool])); + assert_eq!( + only_request(&upstream).await.json()["tools"][0].get("input_examples"), + None + ); +} + #[rstest] #[tokio::test] async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) { @@ -534,3 +586,98 @@ async fn the_request_context_carries_the_requested_reasoning_not_its_translation ) ); } + +#[derive(Default)] +struct RecordingAnswer(Mutex>); + +impl litellm_inference_messages::HostAnswer for RecordingAnswer { + fn answer(&self, stream: bool) -> futures_util::future::BoxFuture<'_, Result<(), Error>> { + self.0.lock().unwrap().push(stream); + Box::pin(async { Ok(()) }) + } +} + +#[rstest] +#[case::unary(false)] +#[case::stream(true)] +#[tokio::test] +async fn an_intercepted_call_is_answered_by_the_host_before_the_mock_and_the_provider( + call: MessagesCall, + #[case] streaming: bool, +) { + let upstream = upstream([message_response()]).await; + let answer = Arc::new(RecordingAnswer::default()); + let route = messages_route(no_secrets()).with_host_answer(answer.clone()); + let mut request = authenticated( + with_fields(call, json!({"stream": streaming})), + upstream.uri(), + ); + request.shaping.settings.intercepted = true; + request.shaping.settings.mock_response = Some(object(json!({"id": "mocked"}))); + let host = RecordingHost::passthrough(request); + + let output = route + .execute(host.request().unwrap(), &host, None) + .await + .unwrap(); + + assert!(matches!( + output, + MessagesCallResponse::Complete(MessagesBody::Unchecked(body)) if body.is_empty() + )); + assert_eq!(*answer.0.lock().unwrap(), vec![streaming]); + assert!(received(&upstream).await.is_empty()); + let facts = host.facts.lock().unwrap(); + assert_eq!(facts.len(), 1); + assert_eq!(facts[0].source, ResultSource::Host); +} + +#[rstest] +#[tokio::test] +async fn a_call_no_interceptor_owns_never_asks_the_host(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let answer = Arc::new(RecordingAnswer::default()); + let route = messages_route(no_secrets()).with_host_answer(answer.clone()); + let host = RecordingHost::passthrough(authenticated(call, upstream.uri())); + + route + .execute(host.request().unwrap(), &host, None) + .await + .unwrap(); + + assert!(answer.0.lock().unwrap().is_empty()); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[derive(Default)] +struct DecliningAnswer(Mutex>); + +impl litellm_inference_messages::HostAnswer for DecliningAnswer { + fn answer(&self, _stream: bool) -> futures_util::future::BoxFuture<'_, Result<(), Error>> { + panic!("a declined interception never asks for an answer") + } + + fn intercept(&self, stream: bool) -> futures_util::future::BoxFuture<'_, Result> { + self.0.lock().unwrap().push(stream); + Box::pin(async { Ok(false) }) + } +} + +#[rstest] +#[tokio::test] +async fn an_interception_no_interceptor_takes_reaches_the_provider(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let answer = Arc::new(DecliningAnswer::default()); + let route = messages_route(no_secrets()).with_host_answer(answer.clone()); + let mut request = authenticated(call, upstream.uri()); + request.shaping.settings.intercepted = true; + let host = RecordingHost::passthrough(request); + + route + .execute(host.request().unwrap(), &host, None) + .await + .unwrap(); + + assert_eq!(*answer.0.lock().unwrap(), vec![false]); + assert_eq!(received(&upstream).await.len(), 1); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index fe203287766..1d515e87b90 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -19,12 +19,16 @@ use wiremock::ResponseTemplate; mod support; use support::*; +mod gateway; mod host; +mod pre_request; mod request; mod response; +mod retry; mod secrets; mod stream; mod tencent; +mod transport; const MODEL: &str = "claude-sonnet-4-5"; diff --git a/litellm-rust/crates/inference-messages/tests/messages/pre_request.rs b/litellm-rust/crates/inference-messages/tests/messages/pre_request.rs new file mode 100644 index 00000000000..13a632bdecd --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/messages/pre_request.rs @@ -0,0 +1,141 @@ +use futures_util::future::BoxFuture; +use litellm_inference_messages::PreRequestHooks; +use rstest::rstest; + +use super::*; + +struct Hooks { + rewrite: Box Result + Send + Sync>, + runs: Mutex, +} + +impl PreRequestHooks for Hooks { + fn run(&self, call: MessagesCall) -> BoxFuture<'_, Result> { + *self.runs.lock().unwrap() += 1; + let rewritten = (self.rewrite)(call); + Box::pin(async move { rewritten }) + } +} + +fn hooks( + rewrite: impl Fn(MessagesCall) -> Result + Send + Sync + 'static, +) -> Arc { + Arc::new(Hooks { + rewrite: Box::new(rewrite), + runs: Mutex::new(0), + }) +} + +async fn run_hooked(hooks: Arc, call: MessagesCall) -> Result { + let host = LocalMessagesHost::new(call); + let route = messages_route(Arc::new(RecordingSecrets::empty())).with_pre_request_hooks(hooks); + litellm_host_native::in_process::run_hosted( + route.machine(host.request()?, None), + host.runtime(), + ) + .await +} + +fn mocked(call: MessagesCall, mock: Map) -> MessagesCall { + MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + mock_response: Some(mock), + ..call.shaping.settings + }, + ..call.shaping + }, + ..call + } +} + +fn routed(call: MessagesCall, base: String) -> MessagesCall { + MessagesCall { + api_key: Some("sk-test".into()), + api_base: Some(base), + ..call + } +} + +#[rstest] +#[tokio::test] +async fn the_provider_receives_the_call_the_hooks_leave(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let hooks = hooks(|call| { + Ok(with_fields( + call, + json!({"system": "hooked", "max_tokens": 7}), + )) + }); + + run_hooked(hooks.clone(), routed(call, upstream.uri())) + .await + .expect("messages call succeeds"); + + let sent: Value = only_request(&upstream).await.body_json().unwrap(); + assert_eq!(sent["system"], "hooked"); + assert_eq!(sent["max_tokens"], 7); + assert_eq!(*hooks.runs.lock().unwrap(), 1); +} + +#[rstest] +#[case::hook_sets_the_mock(false, true, true)] +#[case::hook_keeps_the_callers_mock(true, false, true)] +#[case::hook_clears_the_callers_mock(true, true, false)] +#[tokio::test] +async fn the_mock_is_read_from_the_hooked_call( + call: MessagesCall, + #[case] called_with_mock: bool, + #[case] hook_rewrites_mock: bool, + #[case] mocked_response: bool, +) { + let upstream = upstream([message_response()]).await; + let mock = object(json!({"id": "mocked"})); + let call = routed(call, upstream.uri()); + let call = match called_with_mock { + true => mocked(call, mock.clone()), + false => call, + }; + let hooked_mock = (!called_with_mock).then(|| mock.clone()); + let hooks = hooks(move |call| { + Ok(match (hook_rewrites_mock, &hooked_mock) { + (false, _) => call, + (true, Some(mock)) => mocked(call, mock.clone()), + (true, None) => MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + mock_response: None, + ..call.shaping.settings + }, + ..call.shaping + }, + ..call + }, + }) + }); + + let output = run_hooked(hooks, call) + .await + .expect("messages call succeeds"); + + let mocked_output = matches!( + output, + MessagesOutput::Complete(MessagesBody::Unchecked(ref body)) if *body == mock + ); + assert_eq!(mocked_output, mocked_response); + assert_eq!(received(&upstream).await.is_empty(), mocked_response); +} + +#[rstest] +#[tokio::test] +async fn a_failing_hook_fails_the_call_before_the_provider(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let hooks = hooks(|_| Err(Error::InvalidProvider("refused".into()))); + + let Err(error) = run_hooked(hooks, routed(call, upstream.uri())).await else { + panic!("the hook's failure fails the call"); + }; + + assert!(matches!(error, Error::InvalidProvider(ref provider) if provider == "refused")); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index f978da5e995..7c2f3250acf 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -2,7 +2,9 @@ use litellm_host::{ interceptors::{ExecutionFacts, ResultSource}, lifecycle::ExecutionEvent, }; -use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body}; +use litellm_inference_messages::{ + MessagesBody, MessagesCallResponse, UnencodableHeader, messages_body, +}; use litellm_inference_testing::{ RecordingSecrets, http_config, no_secrets, provider_http, resources, }; @@ -153,6 +155,37 @@ async fn the_message_passes_through_losslessly(call: MessagesCall) { assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body); } +#[rstest] +#[case::without_stop_reason(json!({"stop_sequence": null}))] +#[case::without_stop_fields(json!({}))] +#[tokio::test] +async fn a_message_without_stop_reason_keeps_it_absent(call: MessagesCall, #[case] stop: Value) { + let mut upstream_body = json!({ + "id": "msg_3", "type": "message", "role": "assistant", "model": MODEL, + "content": [{"type": "text", "text": "hi"}], + "usage": {"input_tokens": 1, "output_tokens": 1} + }); + upstream_body + .as_object_mut() + .unwrap() + .extend(stop.as_object().unwrap().clone()); + let upstream = upstream([json_response(upstream_body.clone())]).await; + + let output = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .expect("messages call succeeds"); + + let MessagesOutput::Complete(MessagesBody::Unchecked(body)) = output else { + panic!("a message without stop_reason is returned as sent"); + }; + upstream_body["stop_sequence"] = Value::Null; + assert_eq!(Value::Object(body), upstream_body); +} + #[rstest] #[tokio::test] async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) { @@ -310,6 +343,88 @@ async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { assert_eq!(error.upstream_response(), None); } +fn with_settings(call: MessagesCall, settings: MessagesSettings) -> MessagesCall { + MessagesCall { + shaping: MessagesShaping { + settings, + ..call.shaping + }, + ..call + } +} + +#[rstest] +#[case::unary(false)] +#[case::stream(true)] +#[tokio::test] +async fn a_non_ascii_response_header_fails_the_call_when_python_reads_heads_as_ascii( + call: MessagesCall, + #[case] stream: bool, +) { + let upstream = upstream([message_response().insert_header("x-fz-unicode", "café")]).await; + let call = with_fields( + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + json!({"stream": stream}), + ); + let settings = MessagesSettings { + ascii_response_headers: true, + ..MessagesSettings::default() + }; + + let error = run(with_settings(call, settings)) + .await + .expect_err("the head cannot be encoded"); + + assert_eq!( + litellm_error::find::(&error), + Some(&UnencodableHeader::Received("café".as_bytes().to_vec())) + ); +} + +#[rstest] +#[tokio::test] +async fn a_non_ascii_response_header_is_served_when_python_keeps_raw_heads(call: MessagesCall) { + let upstream = upstream([message_response().insert_header("x-fz-unicode", "café")]).await; + + let message = run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(message.id, "msg_1"); +} + +#[rstest] +#[tokio::test] +async fn an_unencodable_request_header_fails_before_anything_is_sent(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let call = MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }; + let settings = MessagesSettings { + unencodable_request_header: Some(b"h\xed\xa0\x80".to_vec()), + ..MessagesSettings::default() + }; + + let error = run(with_settings(call, settings)) + .await + .expect_err("the header cannot be encoded"); + + assert_eq!( + litellm_error::find::(&error), + Some(&UnencodableHeader::Sent(b"h\xed\xa0\x80".to_vec())) + ); + assert!(received(&upstream).await.is_empty()); +} + #[rstest] #[tokio::test] async fn a_transport_error_does_not_echo_the_provider_url(call: MessagesCall) { @@ -365,7 +480,7 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes } #[rstest] -#[case::mistyped_param(json!({"model": MODEL, "messages": [], "max_tokens": "16"}))] +#[case::mistyped_param(json!({"model": MODEL, "messages": "hello", "max_tokens": 16}))] #[case::missing_messages(json!({"model": MODEL, "max_tokens": 16}))] fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) { let error = messages_body(object(raw)).expect_err("the body is rejected"); diff --git a/litellm-rust/crates/inference-messages/tests/messages/retry.rs b/litellm-rust/crates/inference-messages/tests/messages/retry.rs new file mode 100644 index 00000000000..e48482dc8b0 --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/messages/retry.rs @@ -0,0 +1,364 @@ +use futures_util::TryStreamExt; +use litellm_host::{ + interceptors::{RetriedRequest, RetryReason}, + lifecycle::{CallEvent, ExecutionEvent}, +}; +use litellm_inference_messages::MessagesCallResponse; +use litellm_inference_testing::RecordingSecrets; +use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; +use rstest::rstest; + +use super::*; + +const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block"; +const BOUND_ELSEWHERE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation."; + +fn rejection(status: u16, message: &str) -> ResponseTemplate { + status_response( + status, + json!({"type": "error", "error": {"type": "invalid_request_error", "message": message}}), + ) +} + +fn sse_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_raw( + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "text/event-stream", + ) +} + +fn replay(call: MessagesCall, api_base: String, stream: bool, thinking: Value) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + ..with_fields( + call, + json!({ + "max_tokens": 4096, + "stream": stream, + "thinking": thinking, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": "sig"}, + {"type": "text", "text": "hello"} + ]}, + {"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}, + {"role": "user", "content": "again"} + ] + }), + ) + } +} + +fn enabled() -> Value { + json!({"type": "enabled", "budget_tokens": 2048}) +} + +async fn execute( + host: &RecordingCall, +) -> Result<(MessagesCallResponse, Vec), Error> { + let route = messages_route(Arc::new(RecordingSecrets::empty())); + let response = route + .execute(host.request()?, host, Some(host.events.0.sender.clone())) + .await?; + match response { + MessagesCallResponse::Stream { head, chunks } => Ok(( + MessagesCallResponse::Stream { + head, + chunks: Box::pin(futures_util::stream::empty()), + }, + chunks.try_collect().await?, + )), + complete => Ok((complete, Vec::new())), + } +} + +fn observed_retries(host: &RecordingCall) -> Vec { + host.events + .0 + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { retry }) => { + Some(retry.clone()) + } + _ => None, + }) + .collect() +} + +#[rstest] +#[case::invalid_signature(INVALID_SIGNATURE)] +#[case::missing_signature( + "messages.1.content.0.thinking.signature.str: Input should be a valid string" +)] +#[case::empty_thinking("messages.1.content.0.thinking: each thinking block must contain thinking")] +#[case::unpinned_binding_mismatch(BOUND_ELSEWHERE)] +#[tokio::test] +async fn a_rejected_thinking_replay_is_sent_once_more_without_thinking( + call: MessagesCall, + #[case] message: &str, + #[values(false, true)] stream: bool, + #[values("anthropic", "azure_ai", "deepseek", "minimax")] provider: &str, +) { + let success = if stream { + sse_response() + } else { + message_response() + }; + let upstream = upstream([rejection(400, message), success]).await; + let host = RecordingCall::::new(replay( + MessagesCall { + custom_llm_provider: Some(provider.into()), + ..call + }, + upstream.uri(), + stream, + enabled(), + )); + + let (response, chunks) = execute(&host).await.expect("the retry succeeds"); + + let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap(); + assert_eq!(first.json()["thinking"], enabled()); + assert_eq!( + second.json()["messages"], + json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "hello"}]}, + {"role": "user", "content": "again"} + ]) + ); + let mut expected = first.json(); + expected.as_object_mut().unwrap().shift_remove("thinking"); + expected["messages"] = second.json()["messages"].clone(); + assert_eq!(second.json(), expected); + assert_eq!(first.header("x-hook"), second.header("x-hook")); + let expected_retry = RetriedRequest { + reasons: vec![RetryReason::ThinkingSignature], + body: second.json(), + }; + assert_eq!(*host.retried.lock().unwrap(), vec![expected_retry.clone()]); + assert_eq!(observed_retries(&host), vec![expected_retry]); + assert_eq!( + matches!(response, MessagesCallResponse::Stream { .. }), + stream + ); + assert_eq!(chunks.is_empty(), !stream); +} + +#[rstest] +#[tokio::test] +async fn a_second_rejection_is_the_error_the_caller_sees( + call: MessagesCall, + #[values(false, true)] stream: bool, +) { + let upstream = upstream([ + rejection(400, INVALID_SIGNATURE), + rejection( + 400, + "messages.0.content.0: Invalid `signature` in `thinking` block, again", + ), + ]) + .await; + let host = RecordingCall::::new(replay(call, upstream.uri(), stream, enabled())); + + let error = execute(&host).await.err().expect("the retry fails"); + + assert_eq!(received(&upstream).await.len(), 2); + let (status, body) = error.upstream_response().unwrap(); + assert_eq!(status, 400); + assert!(body.contains("again"), "{body}"); + assert_eq!(host.retried.lock().unwrap().len(), 1); +} + +#[rstest] +#[case::not_a_400(429, INVALID_SIGNATURE, json!({"type": "adaptive"}))] +#[case::unrelated_400(400, "max_tokens: 4096 > 64000", json!({"type": "adaptive"}))] +#[case::signature_without_thinking(400, "Invalid `signature` in `tool_use` block", json!({"type": "adaptive"}))] +#[case::caller_pinned_binding( + 400, + BOUND_ELSEWHERE, + json!({"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}) +)] +#[tokio::test] +async fn other_failures_reach_the_caller_without_a_retry( + call: MessagesCall, + #[case] status: u16, + #[case] message: &str, + #[case] thinking: Value, +) { + let upstream = upstream([rejection(status, message), message_response()]).await; + let adaptive = MessagesShaping { + capabilities: MessagesModelCapabilities { + supports_adaptive_thinking: true, + ..MessagesModelCapabilities::default() + }, + ..MessagesShaping::default() + }; + let host = RecordingCall::::new(replay( + MessagesCall { + shaping: adaptive, + ..call + }, + upstream.uri(), + false, + thinking.clone(), + )); + + let error = execute(&host).await.err().expect("the call fails"); + + let sent = only_request(&upstream).await; + assert_eq!(sent.json()["thinking"], thinking); + assert_eq!(error.upstream_response().unwrap().0, status); + assert!(host.retried.lock().unwrap().is_empty()); + assert!(observed_retries(&host).is_empty()); +} + +const STRICT_REJECTION: &str = "Too many strict tools (21). The maximum is 20."; + +fn strict_tools( + call: MessagesCall, + api_base: String, + stream: bool, + drop_params: bool, +) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + shaping: MessagesShaping { + settings: MessagesSettings { + drop_params, + drops_unsupported_params: drop_params, + ..MessagesSettings::default() + }, + ..MessagesShaping::default() + }, + ..with_fields( + call, + json!({ + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"name": "a", "input_schema": {"type": "object"}, "strict": true}, + {"name": "b", "input_schema": {"type": "object"}, "strict": false}, + {"name": "c", "input_schema": {"type": "object"}} + ] + }), + ) + } +} + +#[rstest] +#[tokio::test] +async fn a_rejected_strict_tool_is_sent_once_more_without_strict_under_drop_params( + call: MessagesCall, + #[values(false, true)] stream: bool, +) { + let success = if stream { + sse_response() + } else { + message_response() + }; + let upstream = upstream([rejection(400, STRICT_REJECTION), success]).await; + let host = RecordingCall::::new(strict_tools(call, upstream.uri(), stream, true)); + + let (response, _) = execute(&host).await.expect("the retry succeeds"); + + let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap(); + assert_eq!(first.json()["tools"][0]["strict"], json!(true)); + let mut expected = first.json(); + expected["tools"] = json!([ + {"name": "a", "input_schema": {"type": "object"}}, + {"name": "b", "input_schema": {"type": "object"}}, + {"name": "c", "input_schema": {"type": "object"}} + ]); + assert_eq!(second.json(), expected); + let expected_retry = RetriedRequest { + reasons: vec![RetryReason::StrictTools { dropped: 1 }], + body: second.json(), + }; + assert_eq!(*host.retried.lock().unwrap(), vec![expected_retry.clone()]); + assert_eq!(observed_retries(&host), vec![expected_retry]); + assert_eq!( + matches!(response, MessagesCallResponse::Stream { .. }), + stream + ); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_strict_tool_reaches_the_caller_without_drop_params(call: MessagesCall) { + let upstream = upstream([rejection(400, STRICT_REJECTION), message_response()]).await; + let host = RecordingCall::::new(strict_tools(call, upstream.uri(), false, false)); + + let error = execute(&host).await.err().expect("the call fails"); + + assert_eq!( + only_request(&upstream).await.json()["tools"][0]["strict"], + json!(true) + ); + assert_eq!(error.upstream_response().unwrap().0, 400); + assert!(observed_retries(&host).is_empty()); +} + +#[rstest] +#[tokio::test] +async fn strict_tools_and_a_thinking_replay_share_the_one_retry(call: MessagesCall) { + let upstream = upstream([ + rejection(400, STRICT_REJECTION), + rejection(400, INVALID_SIGNATURE), + message_response(), + ]) + .await; + let host = RecordingCall::::new(replay( + strict_tools(call, upstream.uri(), false, true), + upstream.uri(), + false, + enabled(), + )); + + let error = execute(&host) + .await + .err() + .expect("the second rejection is raised"); + + let [first, second] = <[_; 2]>::try_from(received(&upstream).await).unwrap(); + assert_eq!(first.json()["thinking"], enabled()); + assert_eq!(second.json()["thinking"], enabled()); + assert_eq!( + second.json()["tools"][0], + json!({"name": "a", "input_schema": {"type": "object"}}) + ); + assert_eq!(received(&upstream).await.len(), 2); + assert!(error.upstream_response().unwrap().1.contains("signature")); +} + +#[rstest] +#[tokio::test] +async fn a_truthy_global_drop_params_that_is_not_true_does_not_drop_strict(call: MessagesCall) { + let upstream = upstream([rejection(400, STRICT_REJECTION), message_response()]).await; + let strict = strict_tools(call, upstream.uri(), false, true); + let host = RecordingCall::::new(MessagesCall { + shaping: MessagesShaping { + settings: MessagesSettings { + drops_unsupported_params: false, + ..strict.shaping.settings.clone() + }, + ..strict.shaping.clone() + }, + ..strict + }); + + let error = execute(&host).await.err().expect("the call fails"); + + assert_eq!( + only_request(&upstream).await.json()["tools"][0]["strict"], + json!(true) + ); + assert_eq!(error.upstream_response().unwrap().0, 400); + assert!(observed_retries(&host).is_empty()); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/stream.rs b/litellm-rust/crates/inference-messages/tests/messages/stream.rs index 99ba08495aa..c9738fa2a0d 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/stream.rs @@ -451,7 +451,7 @@ async fn stalling_error_upstream() -> (String, JoinHandle<()>) { #[rstest] #[tokio::test] -async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) { +async fn a_stalled_stream_error_body_reads_as_empty(call: MessagesCall) { let (base, connection) = stalling_error_upstream().await; let host = RecordingStreamHost::new( MessagesCall { @@ -467,7 +467,7 @@ async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) { .expect_err("a stalled error body fails the call"); assert!(error.dispatched(), "{error:?}"); - assert_eq!(error.upstream_response(), None); + assert_eq!(error.upstream_response(), Some((429, ""))); assert!(host.seen.into_inner().unwrap().is_empty()); tokio::time::timeout(Duration::from_secs(5), connection) .await diff --git a/litellm-rust/crates/inference-messages/tests/messages/transport.rs b/litellm-rust/crates/inference-messages/tests/messages/transport.rs new file mode 100644 index 00000000000..f561b09e55f --- /dev/null +++ b/litellm-rust/crates/inference-messages/tests/messages/transport.rs @@ -0,0 +1,256 @@ +use std::collections::VecDeque; + +use bytes::Bytes; +use futures_util::{StreamExt, TryStreamExt, future::BoxFuture}; +use litellm_inference_messages::{ + Delivered, MessagesCallResponse, Transport, TransportBody, TransportRequest, TransportResponse, +}; +use litellm_inference_testing::RecordingSecrets; +use rstest::rstest; + +use super::*; + +const INVALID_SIGNATURE: &str = "messages.1.content.0: Invalid `signature` in `thinking` block"; +const UNREACHABLE_BASE: &str = "http://127.0.0.1:9"; + +#[derive(Debug, thiserror::Error)] +#[error("raised by the host client")] +struct Raised; + +enum Scripted { + Complete(Value), + Stream(&'static [&'static str]), + Rejected(u16, &'static str), + Failed, +} + +struct Sent { + url: String, + headers: Vec<(String, String)>, + body: Value, + stream: bool, +} + +#[derive(Default)] +struct ScriptedTransport { + script: Mutex>, + sent: Mutex>, +} + +impl ScriptedTransport { + fn new(script: impl IntoIterator) -> Arc { + Arc::new(Self { + script: Mutex::new(script.into_iter().collect()), + sent: Mutex::default(), + }) + } + + fn sent(&self) -> std::sync::MutexGuard<'_, Vec> { + self.sent.lock().unwrap() + } +} + +impl Transport for ScriptedTransport { + fn send( + &self, + request: TransportRequest, + ) -> BoxFuture<'static, Result> { + let TransportBody::Json(body) = request.body else { + panic!("an unsigned request hands the host its JSON body"); + }; + self.sent.lock().unwrap().push(Sent { + url: request.url, + headers: request.headers, + body, + stream: request.stream, + }); + let next = self + .script + .lock() + .unwrap() + .pop_front() + .expect("a scripted answer"); + Box::pin(async move { + let (status, delivered) = match next { + Scripted::Complete(body) => (200, Delivered::Complete(body.to_string())), + Scripted::Stream(frames) => ( + 200, + Delivered::Stream( + futures_util::stream::iter( + frames.iter().map(|frame| Ok(Bytes::from_static(frame.as_bytes()))), + ) + .boxed(), + ), + ), + Scripted::Rejected(status, message) => ( + status, + Delivered::Rejected { + text: json!({"type": "error", "error": {"type": "invalid_request_error", "message": message}}) + .to_string(), + error: Error::transport(Raised), + }, + ), + Scripted::Failed => return Err(Error::transport(Raised)), + }; + Ok(TransportResponse { + status, + headers: vec![("request-id".into(), "req_1".into())], + delivered, + }) + }) + } +} + +fn hosted_call(call: MessagesCall, fields: Value) -> MessagesCall { + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(UNREACHABLE_BASE.into()), + ..with_fields(call, fields) + } +} + +async fn execute( + transport: Arc, + call: MessagesCall, +) -> Result<(MessagesCallResponse, Vec), Error> { + let host = RecordingCall::::new(call); + let route = messages_route(Arc::new(RecordingSecrets::empty())).with_transport(transport); + match route + .execute(host.request()?, &host, Some(host.events.0.sender.clone())) + .await? + { + MessagesCallResponse::Stream { head, chunks } => Ok(( + MessagesCallResponse::Stream { + head, + chunks: Box::pin(futures_util::stream::empty()), + }, + chunks.try_collect().await?, + )), + complete => Ok((complete, Vec::new())), + } +} + +#[rstest] +#[tokio::test] +async fn the_host_transport_sends_the_wire_request_instead_of_the_native_client( + call: MessagesCall, +) { + let transport = ScriptedTransport::new([Scripted::Complete(message_body())]); + + let (response, _) = execute(transport.clone(), hosted_call(call, json!({}))) + .await + .expect("the host transport answers"); + + let MessagesCallResponse::Complete(MessagesBody::Message(message)) = response else { + panic!("a complete message"); + }; + assert_eq!(message.id, "msg_1"); + let sent = transport.sent(); + let [request] = sent.as_slice() else { + panic!("one request"); + }; + assert_eq!(request.url, format!("{UNREACHABLE_BASE}/v1/messages")); + assert!(!request.stream); + assert_eq!(request.body["model"], MODEL); + assert!( + request + .headers + .iter() + .any(|(name, value)| name == "x-api-key" && value == "sk-ant"), + "{:?}", + request.headers + ); +} + +#[rstest] +#[tokio::test] +async fn a_host_stream_reaches_the_caller_chunk_by_chunk(call: MessagesCall) { + let frames: &'static [&'static str] = &[ + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ]; + let transport = ScriptedTransport::new([Scripted::Stream(frames)]); + + let (response, chunks) = execute( + transport.clone(), + hosted_call(call, json!({"stream": true})), + ) + .await + .expect("the host stream opens"); + + let MessagesCallResponse::Stream { head, .. } = response else { + panic!("a stream"); + }; + assert_eq!( + head.headers, + vec![("request-id".to_string(), "req_1".to_string())] + ); + assert_eq!( + chunks, + frames + .iter() + .map(|frame| Bytes::from_static(frame.as_bytes())) + .collect::>() + ); + assert!(transport.sent()[0].stream); +} + +fn replayed_thinking() -> Value { + json!({ + "max_tokens": 4096, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": "sig"}, + {"type": "text", "text": "hello"} + ]}, + {"role": "user", "content": "again"} + ] + }) +} + +#[rstest] +#[tokio::test] +async fn a_recoverable_rejection_is_retried_through_the_host_transport(call: MessagesCall) { + let transport = ScriptedTransport::new([ + Scripted::Rejected(400, INVALID_SIGNATURE), + Scripted::Complete(message_body()), + ]); + + execute(transport.clone(), hosted_call(call, replayed_thinking())) + .await + .expect("the retry succeeds"); + + let sent = transport.sent(); + let [first, second] = sent.as_slice() else { + panic!("two requests"); + }; + assert!(first.body.get("thinking").is_some()); + assert!(second.body.get("thinking").is_none()); +} + +#[rstest] +#[case::final_rejection(vec![Scripted::Rejected(429, "slow down")])] +#[case::spent_retry(vec![ + Scripted::Rejected(400, INVALID_SIGNATURE), + Scripted::Rejected(400, INVALID_SIGNATURE), +])] +#[case::client_failure(vec![Scripted::Failed])] +#[tokio::test] +async fn the_caller_gets_the_error_the_host_transport_raised( + call: MessagesCall, + #[case] script: Vec, +) { + let attempts = script.len(); + let transport = ScriptedTransport::new(script); + + let error = execute(transport.clone(), hosted_call(call, replayed_thinking())) + .await + .err() + .expect("the call fails"); + + assert!(litellm_error::find::(&error).is_some(), "{error:?}"); + assert_eq!(error.upstream_response(), None); + assert_eq!(transport.sent().len(), attempts); +} diff --git a/litellm-rust/crates/inference-messages/tests/support/mod.rs b/litellm-rust/crates/inference-messages/tests/support/mod.rs index 3cecc865070..9115b6e89bb 100644 --- a/litellm-rust/crates/inference-messages/tests/support/mod.rs +++ b/litellm-rust/crates/inference-messages/tests/support/mod.rs @@ -117,6 +117,7 @@ pub struct RecordingCall { pub events: Arc, pub chunks: Mutex>, pub head: Mutex>, + pub retried: Mutex>, } #[derive(Default)] @@ -167,6 +168,7 @@ impl RecordingCall

{ events: Arc::new(CallEvents::default()), chunks: Mutex::new(Vec::new()), head: Mutex::new(None), + retried: Mutex::new(Vec::new()), } } } @@ -195,6 +197,14 @@ impl litellm_host::interceptors::Intercepto ) -> Result<(), P::Error> { Ok(()) } + + async fn provider_request_retried( + &self, + retry: litellm_host::interceptors::RetriedRequest, + ) -> Result<(), P::Error> { + self.retried.lock().unwrap().push(retry); + Ok(()) + } } impl

RecordingCall

diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index 9f59bb9a2a2..509704340fe 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -22,6 +22,7 @@ pub(crate) fn event_name(event: &CallEvent) -> &'static str { CallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "result_ready", CallEvent::Started { .. } => "started", CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }) => "response", + CallEvent::Execution(ExecutionEvent::ProviderRequestRetried { .. }) => "retried", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", CallEvent::Cancelled { .. } => "cancelled", diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs index c7a5d58ad3d..9da62e6ccf4 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/machine.rs @@ -66,6 +66,11 @@ async fn drive_until( reply.send(host.after_provider_response(raw).await); Ok(()) } + HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { reply, .. }) => { + ops.push("retried"); + reply.send(Ok(())); + Ok(()) + } }; if let Err(failure) = answer { break machine.interrupt(failure).await.map(completed); @@ -86,6 +91,7 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire), MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(Ok(())), + MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ProviderRequestRetried { reply, .. })) => reply.send(Ok(())), MachineStep::Suspended(HostRequest::Stream(stream)) => match stream { litellm_host::protocol::StreamDelivery::Open(head, _) => match head {}, litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {}, diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index bcbc98a209e..bc9be4825ea 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -81,3 +81,4 @@ - Cache lookup follows provider preparation, credential resolution, and the request interceptor - keys describe the effective provider URL, authenticated headers, and rewritten body - signed requests bypass caching until the signing identity has a stable cache representation +- A service that keys entries by the caller's arguments (`keyed_by_caller`, the Python host's legacy cache) is looked up before provider preparation, as Python's `@client` wrapper does, and a hit skips request building and the request interceptor diff --git a/litellm-rust/crates/inference/Cargo.toml b/litellm-rust/crates/inference/Cargo.toml index 807a3ea85d3..168b9a569ea 100644 --- a/litellm-rust/crates/inference/Cargo.toml +++ b/litellm-rust/crates/inference/Cargo.toml @@ -12,6 +12,7 @@ litellm-framer.workspace = true tokio-util = { version = "0.7", features = ["codec"] } litellm-secrets.workspace = true litellm-core-utils.workspace = true +litellm-error.workspace = true litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true @@ -29,7 +30,6 @@ time.workspace = true url.workspace = true [dev-dependencies] -litellm-error.workspace = true litellm-cache-memory.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-llms = { workspace = true, features = ["test-support"] } diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index d9b5b891015..fa9ceb6bd73 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -61,6 +61,10 @@ pub trait StreamCachable: Cachable { fn replay(data: Bytes) -> Option>; fn bytes(chunk: &Self::Chunk) -> &[u8]; + + fn storable_stream(text: &str) -> bool { + successful_stream(text, Self::TERMINAL_EVENT) + } } #[derive(Serialize, Deserialize)] @@ -238,12 +242,19 @@ impl CallCache

{ identity: &ProviderIdentity, wire: &WireRequest, ) -> Self { + Self::new( + cache, + policy, + CacheRequest::from_wire(identity.clone(), Some(wire)), + ) + } + + pub fn new(cache: Option<&ScopedCache>, policy: CachePolicy, request: CacheRequest) -> Self { let session = cache.and_then(|cache| { if !policy.enabled() { return None; } let options = cache.options(Some(policy)); - let request = CacheRequest::from_wire(identity.clone(), Some(wire)); Some(CacheSession { request: options.request( &cache.service.config().namespace, @@ -281,7 +292,7 @@ impl CallCache

{ P::Response: Serialize, { let Some(session) = self.session.filter(|session| { - *source == ResultSource::Provider && session.request.controls.writes() + !matches!(source, ResultSource::Cache { .. }) && session.request.controls.writes() }) else { return output; }; @@ -322,7 +333,7 @@ fn capture_stream( None => { if let Some(data) = captured && let Ok(text) = String::from_utf8(data) - && successful_stream(&text, P::TERMINAL_EVENT) + && P::storable_stream(&text) && let Ok(entry) = serde_json::to_value(ResponseEnvelope::new( P::SURFACE, CachedOutput::::Stream(text), diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index ef2635d24a1..39b259a4c97 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -1,6 +1,6 @@ use litellm_cache_response::CachePolicy; use litellm_host::{ - interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource}, + interceptors::{ExecutionFacts, Interceptors, RawResponse, ResultSource, RetriedRequest}, lifecycle::{CallEvent, ExecutionEvent}, observation::ObservationSender, }; @@ -57,4 +57,15 @@ impl<'a, I: Interceptors> CallContext<'a, I> { .await .map_err(Error::post_call) } + + pub async fn provider_request_retried(&self, retry: RetriedRequest) -> Result<(), Error> { + if let Some(observers) = &self.observers { + observers.emit(CallEvent::Execution( + ExecutionEvent::ProviderRequestRetried { + retry: retry.clone(), + }, + )); + } + self.interceptors.provider_request_retried(retry).await + } } diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 4135424d6ef..5f6e1b3f70d 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { Secret(#[from] SecretError), #[error("the host driver failed")] Host(#[source] MachineFault), + #[error("the host transport failed")] + Transport(#[source] litellm_error::Source), #[error("post-call hook failed")] PostCallHook(#[source] Arc), } @@ -47,6 +49,10 @@ impl Error { Self::PostCallHook(Arc::new(error)) } + pub fn transport(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self::Transport(litellm_error::Source::new(cause)) + } + /// The caller's request is what is wrong, as opposed to the environment, the wire, or /// the provider's answer. pub fn is_request(&self) -> bool { @@ -55,7 +61,11 @@ impl Error { Self::Llm(error) => error.is_request(), Self::Auth(error) => error.is_request(), Self::Http(error) => error.is_request(), - Self::Declined(_) | Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, + Self::Declined(_) + | Self::Secret(_) + | Self::Host(_) + | Self::Transport(_) + | Self::PostCallHook(_) => false, } } @@ -70,6 +80,7 @@ impl Error { | Self::Headers(_) | Self::Secret(_) | Self::Host(_) + | Self::Transport(_) | Self::PostCallHook(_) => false, } } @@ -85,6 +96,7 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) + | Self::Transport(_) | Self::PostCallHook(_) => false, } } @@ -100,6 +112,7 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) + | Self::Transport(_) | Self::PostCallHook(_) => None, } } @@ -121,7 +134,7 @@ impl Error { match self { Self::Llm(error) => error.dispatched(), Self::Http(error) => error.dispatched(), - Self::Host(_) | Self::PostCallHook(_) => true, + Self::Host(_) | Self::Transport(_) | Self::PostCallHook(_) => true, Self::InvalidProvider(_) | Self::UnresolvedProvider { .. } | Self::Declined(_) @@ -142,6 +155,7 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) + | Self::Transport(_) | Self::PostCallHook(_) => None, } } @@ -193,6 +207,7 @@ impl Error { | Self::Headers(_) | Self::Secret(_) | Self::Host(_) + | Self::Transport(_) | Self::PostCallHook(_) => None, } } diff --git a/litellm-rust/crates/inference/src/outbound.rs b/litellm-rust/crates/inference/src/outbound.rs index b0cc71d425d..87fae0bbe9e 100644 --- a/litellm-rust/crates/inference/src/outbound.rs +++ b/litellm-rust/crates/inference/src/outbound.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_http::outbound::{json, signed}; +use litellm_http::outbound::{json, json_bytes, signed}; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; @@ -34,3 +34,18 @@ pub fn outbound_request( Some(signer) => signed(request, &signer), } } + +pub fn raw_leaf_outbound_request( + authenticated: Authenticated, + url: &str, + body: &Value, + timeout: Option, +) -> Result { + let Authenticated { headers, signer } = authenticated; + let bytes = litellm_host::raw_json::to_vec(body).map_err(litellm_http::Error::request_body)?; + let request = json_bytes(url, &headers, bytes, timeout)?; + match signer { + None => Ok(request), + Some(signer) => signed(request, &signer), + } +} diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index d30d21e74e7..2e53b3adc5f 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -320,7 +320,7 @@ impl ThinkingConfig { #[macro_rules_attribute::apply(crate::wire_type)] pub struct MessagesRequest { pub model: String, - pub messages: Vec, + pub messages: Vec>, #[serde(flatten)] pub params: MessagesOptionalParams, } @@ -329,7 +329,7 @@ pub struct MessagesRequest { #[derive(Default)] pub struct MessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] - pub max_tokens: Option, + pub max_tokens: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub system: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -337,13 +337,21 @@ pub struct MessagesOptionalParams { #[serde(skip_serializing_if = "Option::is_none")] pub stop_sequences: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub stream: Option, + pub stream: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::recognized::exact_float" + )] + pub temperature: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "crate::recognized::exact_float" + )] + pub top_p: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub top_p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub top_k: Option, + pub top_k: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option>>, #[serde(skip_serializing_if = "Option::is_none")] @@ -374,6 +382,12 @@ pub struct MessagesOptionalParams { pub extra: Map, } +impl MessagesOptionalParams { + pub fn streams(&self) -> bool { + self.stream.as_ref().is_some_and(Recognized::is_truthy) + } +} + impl Message { pub fn blocks(&self) -> &[ContentBlock] { match &self.content { @@ -447,8 +461,8 @@ mod tests { assert_eq!( ( - request.params.max_tokens, - request.params.stream, + request.params.max_tokens.clone(), + request.params.stream.clone(), request .params .extra @@ -456,11 +470,39 @@ mod tests { .map(String::as_str) .collect::>(), ), - (Some(16_u64), Some(true), vec!["safeguards"]) + ( + Some(Recognized::Known(16_u64)), + Some(Recognized::Known(true)), + vec!["safeguards"] + ) ); assert_eq!(serde_json::to_value(request).unwrap(), body); } + #[rstest] + #[case::numeric_string(json!({"max_tokens": "1024", "temperature": "0.5", "top_p": "x", "top_k": "5"}))] + #[case::string_stream(json!({"stream": "true"}))] + #[case::integer_stream(json!({"stream": 0}))] + #[case::floats(json!({"max_tokens": 1024.0, "top_k": 1.5}))] + #[case::booleans(json!({"max_tokens": true, "temperature": false}))] + #[case::wide_integers(json!({"max_tokens": 18_446_744_073_709_551_615_u64, "top_k": 9_223_372_036_854_775_808_u64}))] + #[case::nested(json!({"temperature": [1], "top_p": {"p": 1}}))] + fn scalar_params_of_any_type_round_trip_unchanged(#[case] params: Value) { + let body = Value::Object( + [ + ("model".to_string(), json!("m")), + ( + "messages".to_string(), + json!([{"role": "user", "content": "hi"}]), + ), + ] + .into_iter() + .chain(params.as_object().unwrap().clone()) + .collect(), + ); + assert_eq!(round_trip::(&body), body); + } + #[test] fn text_constructor_serializes_as_a_text_block() { assert_eq!( diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index e3ad783c791..d1c839e4363 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -317,6 +317,11 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], + Self::Compact20260904 => &[ + BetaProvider::BedrockConverse, + BetaProvider::BedrockMantle, + BetaProvider::Databricks, + ], Self::AfkMode20260131 | Self::CePlugins20260901 | Self::CeUserManagement20260713 diff --git a/litellm-rust/crates/llms-types/src/recognized.rs b/litellm-rust/crates/llms-types/src/recognized.rs index 0545389afbc..2c389103aef 100644 --- a/litellm-rust/crates/llms-types/src/recognized.rs +++ b/litellm-rust/crates/llms-types/src/recognized.rs @@ -17,6 +17,25 @@ where as serde::Deserialize<'de>>::deserialize(deserializer).map(Some) } +/// `deserialize_with` for an optional float Python forwards as the caller sent it: an integer +/// stays `Unrecognized` with its exact digits, as Python sends, logs and prints it, and `null` +/// stays absent as before. +pub fn exact_float<'de, D>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = ::deserialize(deserializer)?; + if value.is_null() { + return Ok(None); + } + if matches!(&value, Value::Number(number) if !number.is_f64()) { + return Ok(Some(Recognized::Unrecognized(value))); + } + serde_json::from_value(value) + .map(Some) + .map_err(serde::de::Error::custom) +} + impl Recognized { pub fn known(&self) -> Option<&T> { match self { @@ -24,6 +43,60 @@ impl Recognized { Self::Unrecognized(_) => None, } } + + pub fn map_known(self, f: impl FnOnce(T) -> T) -> Self { + match self { + Self::Known(value) => Self::Known(f(value)), + unrecognized => unrecognized, + } + } +} + +impl Recognized { + pub fn is_truthy(&self) -> bool { + match self { + Self::Known(flag) => *flag, + Self::Unrecognized(other) => python_truthy(other), + } + } +} + +pub fn python_truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(flag) => *flag, + Value::Number(number) => number.as_f64().is_none_or(|number| number != 0.0), + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(fields) => !fields.is_empty(), + } +} + +pub fn map_known_runs( + items: Vec>, + mut f: impl FnMut(Vec) -> Vec, +) -> Vec> { + let mut mapped = Vec::with_capacity(items.len()); + let mut run = Vec::new(); + for item in items { + match item { + Recognized::Known(value) => run.push(value), + unrecognized => { + if !run.is_empty() { + mapped.extend( + f(std::mem::take(&mut run)) + .into_iter() + .map(Recognized::Known), + ); + } + mapped.push(unrecognized); + } + } + } + if !run.is_empty() { + mapped.extend(f(run).into_iter().map(Recognized::Known)); + } + mapped } #[cfg(test)] @@ -32,6 +105,82 @@ mod tests { use rstest::rstest; use serde_json::json; + #[rstest] + #[case::known_true(Recognized::Known(true), true)] + #[case::known_false(Recognized::Known(false), false)] + #[case::string_true(Recognized::Unrecognized(json!("true")), true)] + #[case::string_false(Recognized::Unrecognized(json!("false")), true)] + #[case::empty_string(Recognized::Unrecognized(json!("")), false)] + #[case::zero(Recognized::Unrecognized(json!(0)), false)] + #[case::negative_zero(Recognized::Unrecognized(json!(-0.0)), false)] + #[case::one(Recognized::Unrecognized(json!(1)), true)] + #[case::null(Recognized::Unrecognized(Value::Null), false)] + #[case::empty_list(Recognized::Unrecognized(json!([])), false)] + #[case::list(Recognized::Unrecognized(json!([0])), true)] + #[case::empty_object(Recognized::Unrecognized(json!({})), false)] + #[case::object(Recognized::Unrecognized(json!({"a": 0})), true)] + fn flags_are_read_by_python_truthiness(#[case] flag: Recognized, #[case] expected: bool) { + assert_eq!(flag.is_truthy(), expected); + } + + #[test] + fn runs_of_known_items_are_mapped_around_unrecognized_ones() { + let items = vec![ + Recognized::Known(1), + Recognized::Known(2), + Recognized::Unrecognized(json!("x")), + Recognized::Known(3), + ]; + let mut runs = Vec::new(); + let mapped = map_known_runs(items, |run: Vec| { + runs.push(run.clone()); + run.into_iter() + .filter(|item| *item != 2) + .map(|item| item * 10) + .collect() + }); + assert_eq!(runs, vec![vec![1, 2], vec![3]]); + assert_eq!( + mapped, + vec![ + Recognized::Known(10), + Recognized::Unrecognized(json!("x")), + Recognized::Known(30), + ] + ); + assert_eq!( + map_known_runs(Vec::>::new(), |run| run), + vec![] + ); + } + + #[derive(serde::Deserialize)] + struct Sampling { + #[serde(default, deserialize_with = "exact_float")] + temperature: Option>, + } + + #[rstest] + #[case::float(json!({"temperature": 0.5}), Some(Recognized::Known(0.5)))] + #[case::small_int(json!({"temperature": 2}), Some(Recognized::Unrecognized(json!(2))))] + #[case::zero(json!({"temperature": 0}), Some(Recognized::Unrecognized(json!(0))))] + #[case::integral_float(json!({"temperature": 1.0}), Some(Recognized::Known(1.0)))] + #[case::largest_exact_int(json!({"temperature": 9_007_199_254_740_992_i64}), Some(Recognized::Unrecognized(json!(9_007_199_254_740_992_i64))))] + #[case::inexact_int(json!({"temperature": 1_152_921_504_606_846_977_i64}), Some(Recognized::Unrecognized(json!(1_152_921_504_606_846_977_i64))))] + #[case::inexact_negative_int(json!({"temperature": -9_223_372_036_854_775_807_i64}), Some(Recognized::Unrecognized(json!(-9_223_372_036_854_775_807_i64))))] + #[case::inexact_unsigned_int(json!({"temperature": 18_446_744_073_709_551_615_u64}), Some(Recognized::Unrecognized(json!(18_446_744_073_709_551_615_u64))))] + #[case::string(json!({"temperature": "0.5"}), Some(Recognized::Unrecognized(json!("0.5"))))] + #[case::null(json!({"temperature": null}), None)] + #[case::absent(json!({}), None)] + fn integers_keep_their_digits(#[case] value: Value, #[case] expected: Option>) { + assert_eq!( + serde_json::from_value::(value) + .unwrap() + .temperature, + expected + ); + } + #[rstest] #[case::known(json!(7), Recognized::Known(7))] #[case::wrong_type(json!("7"), Recognized::Unrecognized(json!("7")))] diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 79ee17e8b97..d59cd1be1d2 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -23,9 +23,11 @@ litellm-http.workspace = true litellm-secrets.workspace = true litellm-python-compat.workspace = true litellm-router-types.workspace = true +litellm-tracing.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" +fancy-regex.workspace = true futures-util.workspace = true reqwest.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index 7842b478dce..4a3575c9bbc 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -1,5 +1,5 @@ use litellm_http::request::{with_header, without_headers}; -use litellm_llms_types::providers::anthropic::{BETA_HEADER, BetaProvider, BetaSet}; +use litellm_llms_types::providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet}; use crate::{anthropic::common_utils::existing_betas, base_llm::auth::Headers}; @@ -33,6 +33,16 @@ impl BetaPolicy { } } +pub fn accepted_or_unlisted(betas: BetaSet, provider: BetaProvider) -> BetaSet { + betas + .into_iter() + .filter_map(|beta| match beta { + AnthropicBeta::Other(_) => Some(beta), + known => known.on(provider), + }) + .collect() +} + #[cfg(test)] mod tests { use crate::anthropic::beta_headers::BetaPolicy; diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 0854820358b..cadc5711474 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,4 +1,7 @@ +use std::sync::LazyLock; + use crate::base_llm::messages::context::MessagesModelCapabilities; +use fancy_regex::Regex; use litellm_auth::CredentialPlacement; use litellm_secret_value::SecretValue; @@ -22,7 +25,7 @@ use serde_json::Value; use crate::{ anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, - base_llm::auth::{AuthScheme, Headers}, + base_llm::auth::{AuthScheme, Headers, ValidatedEnvironment}, }; pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; @@ -80,10 +83,21 @@ impl<'a> OauthToken<'a> { } pub fn into_auth(self) -> AuthScheme { - AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: SecretValue::new(self.0), - } + oauth_bearer(SecretValue::new(self.0)) + } +} + +fn oauth_bearer(secret: SecretValue) -> AuthScheme { + AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret, + } +} + +pub fn minted_oauth_environment(headers: Headers, token: SecretValue) -> ValidatedEnvironment { + ValidatedEnvironment { + headers: with_oauth_companions(headers, &[], BrowserAccess::Omitted), + auth: oauth_bearer(token), } } @@ -92,11 +106,23 @@ pub fn get_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, ) -> Option { - resolve_non_empty(api_key, env_lookup, &[ANTHROPIC_API_KEY_ENV]) + nonempty_secret(api_key.map(str::to_owned)) + .or_else(|| nonempty_secret(env_lookup(ANTHROPIC_API_KEY_ENV))) } pub fn get_auth_token(env_lookup: &dyn Fn(&str) -> Option) -> Option { - resolve_non_empty(None, env_lookup, &[ANTHROPIC_AUTH_TOKEN_ENV]) + nonempty_secret(env_lookup(ANTHROPIC_AUTH_TOKEN_ENV)) +} + +/// Python's `normalize_nonempty_secret_str`: `str.strip()` also strips the ASCII separators. +fn nonempty_secret(value: Option) -> Option { + value + .map(|value| { + value + .trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c)) + .to_owned() + }) + .filter(|value| !value.is_empty()) } /// Python's `AnthropicModelInfo.get_auth_header`, naming the credential instead of building @@ -232,14 +258,24 @@ pub fn forwarded_oauth_bearer(headers: &[(String, String)]) -> Option Headers { - merge_beta_headers( +#[derive(Clone, Copy, PartialEq, Eq)] +enum BrowserAccess { + Granted, + Omitted, +} + +fn with_oauth_companions(headers: Headers, dropped: &[&str], browser: BrowserAccess) -> Headers { + let headers = merge_beta_headers( without_headers(headers, dropped), BetaSet::from_iter([AnthropicBeta::Oauth20250420]), - ) - .into_iter() - .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) - .collect() + ); + match browser { + BrowserAccess::Granted => headers + .into_iter() + .chain([(DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string())]) + .collect(), + BrowserAccess::Omitted => headers, + } } pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str>) -> OauthHandling { @@ -247,13 +283,17 @@ pub fn optionally_handle_anthropic_oauth(headers: Headers, api_key: Option<&str> forwarded_oauth_bearer(&headers).map(|token| SecretValue::new(token.as_str())) { return OauthHandling::Bearer { - headers: with_oauth_companions(headers, &[API_KEY_HEADER, AUTHORIZATION]), + headers: with_oauth_companions( + headers, + &[API_KEY_HEADER, AUTHORIZATION], + BrowserAccess::Granted, + ), token, }; } if let Some(token) = api_key.and_then(OauthToken::parse) { return OauthHandling::Bearer { - headers: with_oauth_companions(headers, &[API_KEY_HEADER]), + headers: with_oauth_companions(headers, &[API_KEY_HEADER], BrowserAccess::Granted), token: SecretValue::new(token.as_str()), }; } @@ -278,12 +318,19 @@ pub fn has_advisor_tool(tools: Option<&[Recognized]>) -> bool { .any(|tool| matches!(tool, Recognized::Known(MessagesTool::Advisor { .. }))) } -pub fn requires_native_compaction_beta(compaction: Option<&Value>, messages: &[Message]) -> bool { +pub fn requires_native_compaction_beta( + compaction: Option<&Value>, + messages: &[Recognized], +) -> bool { compaction.is_some() - || messages.iter().flat_map(Message::blocks).any(|block| { - block.is_type(ContentBlockType::Compaction) - && block.signature.as_deref().is_some_and(|s| !s.is_empty()) - }) + || messages + .iter() + .filter_map(Recognized::known) + .flat_map(Message::blocks) + .any(|block| { + block.is_type(ContentBlockType::Compaction) + && block.signature.as_deref().is_some_and(|s| !s.is_empty()) + }) } fn is_blank(text: Option<&str>) -> bool { @@ -650,6 +697,113 @@ pub fn flatten_unencrypted_web_search_results(messages: Vec) -> Vec bool { + let lower = error_text.to_lowercase(); + if !lower.contains("thinking") { + return false; + } + (lower.contains("signature") && (lower.contains("invalid") || lower.contains("valid string"))) + || lower.contains("must contain thinking") +} + +pub fn is_anthropic_thinking_binding_mismatch_error(error_text: &str) -> bool { + error_text + .to_lowercase() + .contains("bound to a different conversation") +} + +pub fn anthropic_request_errors_on_thinking_binding_mismatch(request_body: &Value) -> bool { + request_body + .pointer("/thinking/block_binding/prefix_mismatch_behavior") + .and_then(Value::as_str) + == Some("error") +} + +fn is_thinking_value_block(block: &Value) -> bool { + matches!( + block.get("type").and_then(Value::as_str), + Some("thinking" | "redacted_thinking") + ) +} + +fn without_thinking_value_blocks(message: Value) -> Option { + let Value::Object(mut message) = message else { + return Some(message); + }; + if let Some(Value::Array(content)) = message.get_mut("content") { + content.retain(|block| !is_thinking_value_block(block)); + if content.is_empty() { + return None; + } + } + Some(Value::Object(message)) +} + +pub fn strip_thinking_blocks_from_anthropic_messages_request(request_body: Value) -> Value { + let Value::Object(mut body) = request_body else { + return request_body; + }; + if let Some(Value::Array(messages)) = body.get_mut("messages") { + *messages = std::mem::take(messages) + .into_iter() + .filter_map(without_thinking_value_blocks) + .collect(); + } + body.shift_remove("thinking"); + Value::Object(body) +} + +static CLAUDE_STRICT_TOOLS_REJECTION: LazyLock = LazyLock::new(|| { + Regex::new(concat!( + r"Too many strict tools \(\d+\)", + r"|Schemas contains too many (?:optional parameters|parameters with union types)", + r"|[Cc]ompiled grammar (?:is too large|size)", + r"|tools\.\d+\.custom(?:\.strict)?: ", + r"|Grammar compilation timed out", + )) + .expect("valid regex") +}); + +pub fn is_claude_strict_tools_rejection(status: u16, error_text: &str) -> bool { + status == 400 + && CLAUDE_STRICT_TOOLS_REJECTION + .is_match(error_text) + .unwrap_or(false) +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ToolsWithoutStrict { + pub tools: Vec, + pub dropped_count: usize, +} + +pub fn tools_without_strict(tools: Option<&Value>) -> Option { + let tools = tools?.as_array()?; + if !tools.iter().all(Value::is_object) { + return None; + } + let dropped_count = tools + .iter() + .filter(|tool| tool.get("strict") == Some(&Value::Bool(true))) + .count(); + if dropped_count == 0 { + return None; + } + Some(ToolsWithoutStrict { + tools: tools + .iter() + .cloned() + .map(|mut tool| { + if let Value::Object(fields) = &mut tool { + fields.shift_remove("strict"); + } + tool + }) + .collect(), + dropped_count, + }) +} + #[cfg(test)] mod tests { use crate::base_llm::messages::context::SupportedEffortTiers; @@ -668,7 +822,7 @@ mod tests { serde_json::from_value(value).unwrap() } - fn history(messages: Value) -> Vec { + fn history(messages: Value) -> Vec> { serde_json::from_value(messages).unwrap() } @@ -1636,6 +1790,11 @@ mod tests { &[("ANTHROPIC_AUTH_TOKEN", "env-token")], Some(("Authorization", "env-token")), )] + #[case::auth_token_stripped_of_whitespace( + None, + &[("ANTHROPIC_AUTH_TOKEN", " env-token \t\n\u{1f}")], + Some(("Authorization", "env-token")), + )] #[case::auth_token_when_the_env_key_is_blank( None, &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], @@ -1647,7 +1806,19 @@ mod tests { &[("ANTHROPIC_API_KEY", OAUTH_BEARER)], Some(("Authorization", OAUTH_TOKEN)), )] + #[case::env_key_stripped_of_whitespace( + None, + &[("ANTHROPIC_API_KEY", " sk-env \t\n\u{1f}")], + Some(("x-api-key", "sk-env")), + )] + #[case::api_key_param_stripped_of_separators(Some("\u{1c}sk-param "), &[], Some(("x-api-key", "sk-param")))] + #[case::separator_only_env_key_falls_to_the_auth_token( + None, + &[("ANTHROPIC_API_KEY", "\u{1c}\u{1f}"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + Some(("Authorization", "env-token")), + )] #[case::no_credentials(None, &[], None)] + #[case::separator_only_auth_token(None, &[("ANTHROPIC_AUTH_TOKEN", "\u{1c}\u{1f}")], None)] #[case::blank_everything(Some(""), &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], None)] fn auth_header_prefers_the_key_then_the_auth_token( #[case] api_key: Option<&str>, @@ -1974,4 +2145,166 @@ mod tests { fn the_messages_url_follows_python_ensure_api_route(#[case] api_base: &str, #[case] url: &str) { assert_eq!(complete_anthropic_url(Some(api_base), &|_| None), url); } + + fn shared_cases(table: &str) -> Vec<(String, String, bool)> { + let fixture: Value = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/anthropic_messages_http_error_retry.json" + ))) + .unwrap(); + fixture[table] + .as_array() + .unwrap() + .iter() + .map(|case| { + ( + case["name"].as_str().unwrap().to_string(), + case["error_text"].as_str().unwrap().to_string(), + case["expected"].as_bool().unwrap(), + ) + }) + .collect() + } + + #[rstest] + #[case::invalid_thinking_block( + "invalid_thinking_block_error", + is_anthropic_invalid_thinking_block_error + )] + #[case::thinking_binding_mismatch( + "thinking_binding_mismatch_error", + is_anthropic_thinking_binding_mismatch_error + )] + fn error_texts_match_the_cases_shared_with_python( + #[case] table: &str, + #[case] recognizes: fn(&str) -> bool, + ) { + let cases = shared_cases(table); + assert!(!cases.is_empty()); + let mismatched: Vec<&String> = cases + .iter() + .filter(|(_, error_text, expected)| recognizes(error_text) != *expected) + .map(|(name, _, _)| name) + .collect(); + assert_eq!(mismatched, Vec::<&String>::new()); + } + + #[rstest] + #[case::pinned(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "error"}}}), true)] + #[case::other_behavior(json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "drop"}}}), false)] + #[case::string_binding(json!({"thinking": {"block_binding": "strict"}}), false)] + #[case::no_thinking(json!({"messages": []}), false)] + #[case::not_an_object(json!(["thinking"]), false)] + fn a_pinned_binding_is_the_caller_asking_for_the_error( + #[case] body: Value, + #[case] expected: bool, + ) { + assert_eq!( + anthropic_request_errors_on_thinking_binding_mismatch(&body), + expected + ); + } + + #[rstest] + #[case::mixed_turn( + json!({"model": "m", "thinking": {"type": "enabled"}, "messages": [ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "t", "signature": "s"}, + {"type": "text", "text": "a"}, + {"type": "redacted_thinking", "data": "d"} + ]} + ]}), + json!({"model": "m", "messages": [{"role": "assistant", "content": [{"type": "text", "text": "a"}]}]}) + )] + #[case::thinking_only_and_empty_turns_are_dropped( + json!({"messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "redacted_thinking", "data": "d"}]}, + {"role": "assistant", "content": []}, + {"role": "user", "content": "again"} + ]}), + json!({"messages": [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}]}) + )] + #[case::non_object_entries_are_kept( + json!({"messages": ["raw", {"role": "user", "content": ["raw", {"type": "thinking"}]}]}), + json!({"messages": ["raw", {"role": "user", "content": ["raw"]}]}) + )] + #[case::messages_not_a_list( + json!({"messages": "hi", "thinking": {"type": "adaptive"}}), + json!({"messages": "hi"}) + )] + fn retrying_strips_thinking_from_the_replay_and_the_request( + #[case] body: Value, + #[case] expected: Value, + ) { + assert_eq!( + strip_thinking_blocks_from_anthropic_messages_request(body), + expected + ); + } + + #[test] + fn stripping_keeps_the_order_of_the_remaining_fields() { + let stripped = strip_thinking_blocks_from_anthropic_messages_request( + json!({"model": "m", "thinking": {}, "messages": [], "max_tokens": 1}), + ); + let keys: Vec<&str> = stripped + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect(); + assert_eq!(keys, ["model", "messages", "max_tokens"]); + } + + fn shared_table(table: &str) -> Vec { + let fixture: Value = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/anthropic_messages_http_error_retry.json" + ))) + .unwrap(); + fixture[table].as_array().unwrap().clone() + } + + #[test] + fn strict_tools_rejections_match_the_cases_shared_with_python() { + let cases = shared_table("claude_strict_tools_rejection"); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap(); + is_claude_strict_tools_rejection(status, case["error_text"].as_str().unwrap()) + != case["expected"].as_bool().unwrap() + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); + } + + #[test] + fn tools_without_strict_matches_the_cases_shared_with_python() { + let cases = shared_table("tools_without_strict"); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let expected = case["expected"] + .as_object() + .map(|expected| ToolsWithoutStrict { + tools: expected["tools"].as_array().unwrap().clone(), + dropped_count: usize::try_from(expected["dropped_count"].as_u64().unwrap()) + .unwrap(), + }); + tools_without_strict(Some(&case["tools"])) != expected + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); + } + + #[test] + fn absent_tools_have_no_strict_to_drop() { + assert_eq!(tools_without_strict(None), None); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 54f924ddec6..46a16e7d546 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -3,7 +3,7 @@ use litellm_llms_types::{ AdaptiveThinking, EnabledThinking, Message, MessagesOptionalParams, MessagesRequest, ThinkingConfig, ThinkingDisplay, }, - recognized::Recognized, + recognized::{Recognized, map_known_runs}, }; use serde_json::{Value, json}; @@ -13,14 +13,19 @@ use crate::{ flatten_unencrypted_web_search_results, sanitize_tool_use_ids, strip_empty_content_blocks, strip_provider_specific_fields, }, + base_llm::messages::context::MessagesShapeOptions, }; pub fn shape_anthropic_messages_request( request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { Ok(MessagesRequest { - messages: sanitize_anthropic_messages(request.messages), + messages: if options.presanitized { + map_known_runs(request.messages, strip_provider_specific_fields) + } else { + map_known_runs(request.messages, sanitize_anthropic_messages) + }, params: MessagesOptionalParams { metadata: request .params @@ -28,7 +33,10 @@ pub fn shape_anthropic_messages_request( .as_ref() .map(validate_anthropic_api_metadata) .transpose()?, - thinking: with_reasoning_auto_summary(request.params.thinking, reasoning_auto_summary), + thinking: with_reasoning_auto_summary( + request.params.thinking, + options.reasoning_auto_summary, + ), ..request.params }, ..request @@ -299,7 +307,10 @@ mod tests { "thinking": {"type": "enabled", "budget_tokens": 1024}, "safeguards": [{"type": "dangerous_tool_use"}] })), - true, + MessagesShapeOptions { + reasoning_auto_summary: true, + presanitized: false, + }, ) .unwrap(); assert_eq!( @@ -315,4 +326,27 @@ mod tests { }) ); } + + #[test] + fn presanitized_shaping_keeps_blocks_and_strips_only_provider_fields() { + let messages = json!([{"role": "user", "content": [ + {"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}}, + {"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok", "provider_specific_fields": {"x": 1}} + ]}]); + let sanitized = shape_anthropic_messages_request( + request(json!({"model": "m", "messages": messages})), + MessagesShapeOptions { + reasoning_auto_summary: false, + presanitized: true, + }, + ) + .unwrap(); + assert_eq!( + serde_json::to_value(sanitized).unwrap(), + json!({"model": "m", "messages": [{"role": "user", "content": [ + {"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}}, + {"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"} + ]}]}) + ); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs index 597a1d819d9..4fa24a3f0f5 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/mod.rs @@ -1,3 +1,4 @@ pub mod handler; +mod scalar; pub mod thinking; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/messages/scalar.rs b/litellm-rust/crates/llms/src/anthropic/messages/scalar.rs new file mode 100644 index 00000000000..db9aa01d177 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/messages/scalar.rs @@ -0,0 +1,113 @@ +use litellm_host::raw_json; +use litellm_llms_types::recognized::Recognized; +use litellm_python_compat::{ + Value as PythonValue, + json::{from_json, loads}, + repr::to_str, +}; +use serde_json::Value; + +pub(crate) fn python_value(value: &Value) -> PythonValue { + match value { + Value::String(text) => raw_json::decode(text) + .and_then(|raw| loads(raw).ok()) + .unwrap_or_else(|| PythonValue::Str(text.clone())), + Value::Array(items) => PythonValue::List(items.iter().map(python_value).collect()), + Value::Object(entries) => PythonValue::Dict( + entries + .iter() + .map(|(key, item)| (PythonValue::Str(key.clone()), python_value(item))) + .collect(), + ), + Value::Null | Value::Bool(_) | Value::Number(_) => from_json(value.clone()), + } +} + +pub(crate) fn python_number(value: &Value) -> Option { + match value { + Value::Bool(flag) => Some(f64::from(u8::from(*flag))), + Value::Number(number) => number.as_f64(), + Value::String(text) => match raw_json::decode(text)? { + "NaN" => Some(f64::NAN), + "Infinity" => Some(f64::INFINITY), + "-Infinity" => Some(f64::NEG_INFINITY), + raw => raw.parse().ok(), + }, + _ => None, + } +} + +pub(crate) fn is_one(value: &Recognized) -> bool { + match value { + Recognized::Known(number) => *number == 1.0, + Recognized::Unrecognized(other) => python_number(other) == Some(1.0), + } +} + +pub(crate) fn python_str(value: &Recognized) -> String { + match value { + Recognized::Known(known) => serde_json::to_value(known) + .map(|known| to_str(&from_json(known))) + .unwrap_or_default(), + Recognized::Unrecognized(other) => unrecognized_str(other), + } +} + +fn unrecognized_str(value: &Value) -> String { + match value { + Value::Null => "None".into(), + Value::Bool(true) => "True".into(), + Value::Bool(false) => "False".into(), + Value::String(text) => match raw_json::decode(text) { + Some("NaN") => "nan".into(), + Some("Infinity") => "inf".into(), + Some("-Infinity") => "-inf".into(), + Some(raw) => raw.trim_matches('"').into(), + None => text.clone(), + }, + other => other.to_string(), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::int(json!(3), Some(3.0))] + #[case::float(json!(1.5), Some(1.5))] + #[case::bool(json!(true), Some(1.0))] + #[case::wide(json!(raw_json::encode("18446744073709551616")), Some(18_446_744_073_709_551_616.0))] + #[case::negative_infinity(json!(raw_json::encode("-Infinity")), Some(f64::NEG_INFINITY))] + #[case::string(json!("1"), None)] + #[case::marked_string(json!(raw_json::encode("\"1\"")), None)] + #[case::null(Value::Null, None)] + fn numbers_compare_as_python_compares_them( + #[case] value: Value, + #[case] expected: Option, + ) { + assert_eq!(python_number(&value), expected); + } + + #[test] + fn nan_is_a_number_unequal_to_everything() { + let nan = json!(raw_json::encode("NaN")); + assert!(python_number(&nan).is_some_and(f64::is_nan)); + assert!(!is_one(&Recognized::Unrecognized(nan))); + } + + #[rstest] + #[case::known(Recognized::Known(0.5), "0.5")] + #[case::known_wide(Recognized::Known(1e16), "1e+16")] + #[case::int(Recognized::Unrecognized(json!(0)), "0")] + #[case::bool(Recognized::Unrecognized(json!(false)), "False")] + #[case::string(Recognized::Unrecognized(json!("0.5")), "0.5")] + #[case::nan(Recognized::Unrecognized(json!(raw_json::encode("NaN"))), "nan")] + #[case::wide(Recognized::Unrecognized(json!(raw_json::encode("-1180591620717411303424"))), "-1180591620717411303424")] + fn values_print_as_python_prints_them(#[case] value: Recognized, #[case] expected: &str) { + assert_eq!(python_str(&value), expected); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index 0ddbc915154..c6f7c1679ef 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -1,17 +1,20 @@ +use litellm_host::raw_json; use litellm_llms_types::{ formats::{ chat_completions::ReasoningEffort, messages::{ - AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams, - MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay, + AdaptiveThinking, DisabledThinking, EffortLevel, EnabledThinking, + MessagesOptionalParams, MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay, }, }, recognized::Recognized, }; -use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy}; +use litellm_python_compat::{ + Value as PythonValue, json::from_json, repr::repr, reprlib::reprlib, truthy::truthy, +}; use serde_json::Value; -use strum::VariantArray; +use super::scalar::{is_one, python_number, python_str, python_value}; use crate::base_llm::messages::context::{ MessagesModelCapabilities, ThinkingBudgets, ThinkingContext, }; @@ -56,13 +59,61 @@ fn unmapped_effort(effort: &Value) -> Error { Error::InvalidRequest(crate::ErrorDetail::InvalidChoice { field: "reasoning effort", actual: repr(&from_json(effort.clone())), - choices: ReasoningEffort::VARIANTS - .iter() - .map(|effort| <&'static str>::from(*effort)) - .collect(), + choices: vec!["minimal", "low", "medium", "high", "xhigh", "max", "none"], }) } +fn invalid_adaptive_effort(effort: &Value) -> Error { + Error::InvalidRequest(crate::ErrorDetail::InvalidReasoningEffort { + actual: repr(&from_json(effort.clone())), + }) +} + +const NON_STRING_EFFORT_REASON: &str = "it is not a string. Send reasoning_effort as a level name such as 'high' to set a reasoning level."; +const NON_NUMERIC_BUDGET_REASON: &str = + "it is not a number. Send budget_tokens as an integer to set a reasoning level."; + +fn warn_dropped_reasoning( + param: &str, + value: &PythonValue, + model: &str, + reason: &str, + drop_params_gated: bool, +) { + let gate = if drop_params_gated { + " (drop_params=True)" + } else { + "" + }; + litellm_tracing::warn!( + "Dropping {param}={} for model={}{gate} because {reason}", + reprlib(value), + repr(&PythonValue::Str(model.to_owned())) + ); +} + +fn warn_non_string_reasoning_effort(effort: &Value, model: &str) { + let value = python_value(effort); + let requested = match &value { + PythonValue::Str(_) => return, + PythonValue::Dict(entries) => entries + .iter() + .find(|(key, _)| *key == PythonValue::Str("effort".to_owned())) + .map(|(_, level)| level), + other => Some(other), + }; + if matches!(requested, None | Some(PythonValue::None)) { + return; + } + warn_dropped_reasoning( + "reasoning_effort", + &value, + model, + NON_STRING_EFFORT_REASON, + false, + ); +} + fn unsupported_effort(level: EffortLevel, model: &str) -> Error { Error::InvalidRequest(crate::ErrorDetail::UnsupportedValue { field: "effort", @@ -82,11 +133,56 @@ fn output_effort(effort: ReasoningEffort) -> Option { } } -fn fit_budget_to_max_tokens(budget_tokens: u64, max_tokens: Option) -> Option { - let Some(max_tokens) = max_tokens else { - return Some(budget_tokens); +/// `cap_thinking_budget_to_max_tokens`, whose first comparison (`max_tokens <= 1024`) raises +/// `TypeError` for a `max_tokens` Python cannot order against an int, such as a string. +fn fit_budget_to_max_tokens( + budget_tokens: u64, + max_tokens: Option<&Recognized>, +) -> Result>, Error> { + let max_tokens = match max_tokens { + None => return Ok(Some(Recognized::Known(budget_tokens))), + Some(Recognized::Known(max_tokens)) => *max_tokens, + Some(Recognized::Unrecognized(other)) => { + return fit_budget_to_python_number(budget_tokens, other); + } }; - (max_tokens > ANTHROPIC_MIN_THINKING_BUDGET_TOKENS).then(|| budget_tokens.min(max_tokens - 1)) + Ok((max_tokens > ANTHROPIC_MIN_THINKING_BUDGET_TOKENS) + .then(|| Recognized::Known(budget_tokens.min(max_tokens - 1)))) +} + +fn enabled_thinking(budget_tokens: Recognized) -> ThinkingConfig { + ThinkingConfig::Enabled(EnabledThinking { + budget_tokens: Some(budget_tokens), + ..EnabledThinking::default() + }) +} + +fn fit_budget_to_python_number( + budget_tokens: u64, + max_tokens: &Value, +) -> Result>, Error> { + let Some(max_tokens) = python_number(max_tokens) else { + if max_tokens.is_null() { + return Ok(Some(Recognized::Known(budget_tokens))); + } + return Err(Error::InvalidRequest( + crate::ErrorDetail::UnorderableTypes { + op: "<=", + left: python_type_name(max_tokens), + right: "int", + }, + )); + }; + if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS as f64 { + return Ok(None); + } + if (budget_tokens as f64) < max_tokens { + return Ok(Some(Recognized::Known(budget_tokens))); + } + Ok(Some(Recognized::Unrecognized( + serde_json::Number::from_f64(max_tokens - 1.0) + .map_or_else(|| Value::String(raw_json::encode("NaN")), Value::Number), + ))) } fn known_thinking(request: &MessagesRequest) -> Option<&ThinkingConfig> { @@ -161,10 +257,17 @@ fn translate_reasoning_effort( }; let effort = match reasoning_effort { Recognized::Known(effort) => effort, - Recognized::Unrecognized(value @ Value::String(_)) => { - return Err(unmapped_effort(&value)); + Recognized::Unrecognized(value) if matches!(python_value(&value), PythonValue::Str(_)) => { + return Err(if context.capabilities.supports_adaptive_thinking { + invalid_adaptive_effort(&value) + } else { + unmapped_effort(&value) + }); + } + Recognized::Unrecognized(value) => { + warn_non_string_reasoning_effort(&value, &request.model); + return Ok(request); } - Recognized::Unrecognized(_) => return Ok(request), }; let (Some(level), Some(budget)) = ( output_effort(effort), @@ -203,10 +306,20 @@ fn translate_reasoning_effort( ..request }); } - let Some(budget) = fit_budget_to_max_tokens(budget, request.params.max_tokens) else { + let Some(budget) = fit_budget_to_max_tokens(budget, request.params.max_tokens.as_ref())? else { + litellm_tracing::warn!( + "Dropping `thinking` mapped from reasoning_effort={} for model={}: max_tokens={} is too small to fit the minimum thinking budget.", + <&'static str>::from(effort), + request.model, + request + .params + .max_tokens + .as_ref() + .map_or_else(|| "None".to_owned(), python_str) + ); return Ok(request); }; - let enabled = ThinkingConfig::enabled(budget); + let enabled = enabled_thinking(budget); Ok(MessagesRequest { params: MessagesOptionalParams { thinking: Some( @@ -340,6 +453,15 @@ fn translate_legacy_thinking_for_adaptive_model( crate::ErrorDetail::NonNumericThinkingBudget { actual }, )); } + if let (None, Some(Recognized::Unrecognized(dropped))) = (budget, &enabled.budget_tokens) { + warn_dropped_reasoning( + "thinking.budget_tokens", + &python_value(dropped), + &request.model, + NON_NUMERIC_BUDGET_REASON, + true, + ); + } let thinking = ThinkingConfig::Adaptive(AdaptiveThinking { display: kept_display(enabled.display.as_ref(), context), extra: carried_thinking_keys(&enabled.extra, context), @@ -396,11 +518,20 @@ fn translate_adaptive_effort_for_non_adaptive_model( } else { None }; + let thinking = match budget { + Some(budget) => fit_budget_to_max_tokens(budget, request.params.max_tokens.as_ref())? + .map(|budget| Recognized::Known(enabled_thinking(budget))), + None => None, + }; + if thinking.is_none() { + litellm_tracing::warn!( + "Dropping adaptive `thinking`/`output_config.effort` for model={}: the model does not support extended thinking, or max_tokens is too small to fit the minimum thinking budget.", + request.model + ); + } Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: budget - .and_then(|budget| fit_budget_to_max_tokens(budget, request.params.max_tokens)) - .map(|budget| Recognized::Known(ThinkingConfig::enabled(budget))), + thinking, output_config: without_effort(request.params.output_config), ..request.params }, @@ -418,7 +549,8 @@ fn drop_incompatible_temperature_for_thinking( let pinned = request .params .temperature - .is_some_and(|temperature| temperature != 1.0); + .as_ref() + .is_some_and(|temperature| !is_one(temperature)); let thinking_enabled = matches!(known_thinking(&request), Some(ThinkingConfig::Enabled(_))); let effort_enabled = known_effort(&request).is_some(); if !pinned || !(thinking_enabled || effort_enabled) { @@ -451,7 +583,7 @@ mod tests { use super::*; use crate::base_llm::messages::context::SupportedEffortTiers; - const EFFORT_CHOICES: &str = "'none', 'minimal', 'low', 'medium', 'high', 'xhigh', 'max'"; + const EFFORT_CHOICES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"; fn request(fields: Value) -> MessagesRequest { let mut body = serde_json::json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); @@ -819,17 +951,17 @@ mod tests { #[case::invalid_on_adaptive_model( opus_4_7(), serde_json::json!({"max_tokens": 1024, "reasoning_effort": "invalid"}), - format!("Unmapped reasoning effort: 'invalid'. Must be one of: {EFFORT_CHOICES}.") + format!("Invalid reasoning_effort: 'invalid'. Must be one of: {EFFORT_CHOICES}") )] #[case::disabled_on_adaptive_model( opus_4_7(), serde_json::json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), - format!("Unmapped reasoning effort: 'disabled'. Must be one of: {EFFORT_CHOICES}.") + format!("Invalid reasoning_effort: 'disabled'. Must be one of: {EFFORT_CHOICES}") )] #[case::empty_on_adaptive_model( opus_4_7(), serde_json::json!({"max_tokens": 1024, "reasoning_effort": ""}), - format!("Unmapped reasoning effort: ''. Must be one of: {EFFORT_CHOICES}.") + format!("Invalid reasoning_effort: ''. Must be one of: {EFFORT_CHOICES}") )] #[case::xhigh_without_xhigh_tier_on_4_6( sonnet_4_6(), @@ -1434,6 +1566,34 @@ mod tests { ); } + #[rstest] + #[case::reasoning_effort(json!({"max_tokens": "8192", "reasoning_effort": "high"}), "str")] + #[case::caller_thinking( + json!({"max_tokens": "8192", "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 2048}}), + "str" + )] + #[case::adaptive_interface( + json!({"max_tokens": "8192", "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}), + "str" + )] + #[case::list(json!({"max_tokens": [8192], "reasoning_effort": "high"}), "list")] + fn an_unorderable_max_tokens_raises_python_type_error_while_capping_the_budget( + haiku_4_5: MessagesModelCapabilities, + #[case] fields: Value, + #[case] left: &'static str, + ) { + assert_eq!( + translate(haiku_4_5, fields), + Err(Error::InvalidRequest( + crate::ErrorDetail::UnorderableTypes { + op: "<=", + left, + right: "int", + } + )) + ); + } + #[rstest] #[case::summarized("summarized", false, Some("summarized"))] #[case::updates_kept("updates", true, Some("updates"))] diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 18789e5b2be..bbb9b4e4c17 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,29 +1,45 @@ use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::{RetriedRequest, RetryReason}; +use litellm_http::request::{with_header, without_headers}; use litellm_llms_types::{ formats::messages::{ - ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, + ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, MessageContent, MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay, }, - providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, - recognized::Recognized, + providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet, DEFAULT_HEADERS}, + recognized::{Recognized, map_known_runs}, }; use litellm_router_types::LitellmParams; use serde_json::{Map, Value, json}; -use super::{handler::shape_anthropic_messages_request, thinking::translate_thinking}; -use crate::base_llm::messages::context::MessagesTransformContext; +use super::{ + handler::shape_anthropic_messages_request, + scalar::{is_one, python_str, python_value}, + thinking::translate_thinking, +}; +use crate::anthropic::beta_headers::accepted_or_unlisted; +use crate::base_llm::messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + tools::ToolCustomField, +}; use crate::{ Error, anthropic::common_utils::{ ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV, - ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, get_auth_header, - has_advisor_tool, has_anthropic_credential, is_tool_search_used, merge_beta_headers, - optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks, - strip_encrypted_reasoning_blocks, + ANTHROPIC_BASE_URL_ENV, OauthHandling, + anthropic_request_errors_on_thinking_binding_mismatch, complete_anthropic_url, + existing_betas, get_auth_header, has_advisor_tool, has_anthropic_credential, + is_anthropic_invalid_thinking_block_error, is_anthropic_thinking_binding_mismatch_error, + is_claude_strict_tools_rejection, is_tool_search_used, merge_beta_headers, + minted_oauth_environment, optionally_handle_anthropic_oauth, + requires_native_compaction_beta, strip_advisor_blocks, strip_encrypted_reasoning_blocks, + strip_thinking_blocks_from_anthropic_messages_request, tools_without_strict, }, base_llm::{ auth::AuthScheme, - messages::transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, + messages::transformation::{ + BaseMessagesConfig, FederatedEnvironment, Headers, ValidatedEnvironment, + }, }, }; @@ -35,9 +51,9 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( @@ -59,6 +75,12 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { transform_messages_request(request, &context.keeping_display_updates()) } + fn tool_custom_field(&self, model: &str) -> ToolCustomField { + ToolCustomField::Rewrite { + hoist_defer_loading: is_claude_model(model), + } + } + fn secret_names(&self) -> &'static [&'static str] { &[ ANTHROPIC_API_KEY_ENV, @@ -79,31 +101,22 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { _litellm_params: &LitellmParams, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - let headers = match optionally_handle_anthropic_oauth(headers, api_key) { - OauthHandling::Bearer { headers, token } => { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: token, - }, - }); - } - OauthHandling::Untouched(headers) => headers, - }; - if has_anthropic_credential(&headers) { - return Ok(ValidatedEnvironment { - headers, - auth: AuthScheme::Forwarded, - }); + match anthropic_environment(headers, api_key, env_lookup) { + FederatedEnvironment::Validated(validated) => Ok(validated), + FederatedEnvironment::Unauthenticated { missing, .. } => Err(missing), } - let auth = get_auth_header(api_key, env_lookup).ok_or(Error::Auth( - litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: ANTHROPIC_API_KEY_ENV, - }, - ))?; - Ok(ValidatedEnvironment { headers, auth }) + } + + fn validate_federated_environment( + &self, + headers: Headers, + api_key: Option<&str>, + _model: &str, + _optional_params: &MessagesOptionalParams, + _litellm_params: &LitellmParams, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(anthropic_environment(headers, api_key, env_lookup)) } fn default_headers(&self) -> &'static [(&'static str, &'static str)] { @@ -116,6 +129,157 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { provider_feature_betas(request, BetaProvider::Anthropic), ) } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + + fn headers_for_url(&self, headers: Headers, url: &str) -> Headers { + if !is_foundry_url(url) { + return headers; + } + let betas = existing_betas(&headers); + if betas.is_empty() { + return headers; + } + let kept = accepted_or_unlisted(betas, BetaProvider::AzureAi); + if kept.is_empty() { + return without_headers(headers, &[BETA_HEADER]); + } + with_header(headers, BETA_HEADER, kept.to_string()) + } +} + +const FOUNDRY_HOST_SUFFIXES: [&str; 3] = [ + "services.ai.azure.com", + "openai.azure.com", + "cognitiveservices.azure.com", +]; + +fn is_foundry_url(url: &str) -> bool { + let absolute = if url.contains("://") { + url.to_string() + } else { + format!("https://{url}") + }; + url::Url::parse(&absolute) + .ok() + .and_then(|parsed| parsed.host_str().map(str::to_ascii_lowercase)) + .is_some_and(|host| { + FOUNDRY_HOST_SUFFIXES + .iter() + .any(|suffix| host == *suffix || host.ends_with(&format!(".{suffix}"))) + }) +} + +pub(crate) fn is_claude_model(model: &str) -> bool { + model.to_lowercase().contains("claude") +} + +pub fn should_retry_anthropic_messages_on_http_error( + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, +) -> bool { + if status != 400 { + return false; + } + let drops_rejected_strict_tools = context.drops_unsupported_params + && is_claude_strict_tools_rejection(status, error_text) + && tools_without_strict(request_body.get("tools")).is_some(); + if drops_rejected_strict_tools { + return true; + } + let caller_asked_for_this_error = + anthropic_request_errors_on_thinking_binding_mismatch(request_body) + && is_anthropic_thinking_binding_mismatch_error(error_text); + !caller_asked_for_this_error && is_anthropic_invalid_thinking_block_error(error_text) +} + +pub fn transform_anthropic_messages_request_on_http_error( + status: u16, + error_text: &str, + request_body: Value, +) -> RetriedRequest { + let mut reasons = Vec::new(); + let mut body = if status == 400 && is_anthropic_invalid_thinking_block_error(error_text) { + reasons.push(RetryReason::ThinkingSignature); + strip_thinking_blocks_from_anthropic_messages_request(request_body) + } else { + request_body + }; + let without_strict = if is_claude_strict_tools_rejection(status, error_text) { + tools_without_strict(body.get("tools")) + } else { + None + }; + if let (Some(without_strict), Value::Object(fields)) = (without_strict, &mut body) { + reasons.push(RetryReason::StrictTools { + dropped: without_strict.dropped_count, + }); + fields.insert("tools".to_owned(), Value::Array(without_strict.tools)); + } + RetriedRequest { reasons, body } +} + +pub fn retry_anthropic_messages_after_http_error( + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, +) -> Option { + should_retry_anthropic_messages_on_http_error(status, error_text, request_body, context).then( + || { + transform_anthropic_messages_request_on_http_error( + status, + error_text, + request_body.clone(), + ) + }, + ) +} + +fn anthropic_environment( + headers: Headers, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> FederatedEnvironment { + let headers = match optionally_handle_anthropic_oauth(headers, api_key) { + OauthHandling::Bearer { headers, token } => { + return FederatedEnvironment::Validated(ValidatedEnvironment { + headers, + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: token, + }, + }); + } + OauthHandling::Untouched(headers) => headers, + }; + if has_anthropic_credential(&headers) { + return FederatedEnvironment::Validated(ValidatedEnvironment { + headers, + auth: AuthScheme::Forwarded, + }); + } + match get_auth_header(api_key, env_lookup) { + Some(auth) => FederatedEnvironment::Validated(ValidatedEnvironment { headers, auth }), + None => FederatedEnvironment::Unauthenticated { + headers, + missing: Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }), + federate: minted_oauth_environment, + }, + } } pub(crate) fn transform_messages_request( @@ -125,6 +289,7 @@ pub(crate) fn transform_messages_request( if request.params.max_tokens.is_none() { return Err(Error::MissingField("max_tokens")); } + let request = without_unstructured_output_format(request, context); let request = drop_unsupported_params(request, context)?; let request = translate_thinking(request, &context.thinking)?; let context_management = request @@ -132,13 +297,20 @@ pub(crate) fn transform_messages_request( .context_management .clone() .map(map_openai_context_management_to_anthropic); - let messages = if has_advisor_tool(request.params.tools.as_deref()) { + let messages = if has_advisor_tool(request.params.tools.as_deref()) + || continues_message_thread(&request.params) + { request.messages } else { - strip_advisor_blocks(request.messages) + map_known_runs(request.messages, strip_advisor_blocks) + }; + let messages = if context.drops_unsupported_params { + drop_non_string_tool_result_texts(messages) + } else { + messages }; Ok(MessagesRequest { - messages: strip_encrypted_reasoning_blocks(messages), + messages: map_known_runs(messages, strip_encrypted_reasoning_blocks), params: MessagesOptionalParams { context_management, extra: without_null_safeguards(request.params.extra), @@ -148,6 +320,84 @@ pub(crate) fn transform_messages_request( }) } +fn without_unstructured_output_format( + request: MessagesRequest, + context: &MessagesTransformContext, +) -> MessagesRequest { + if !context.drops_unsupported_params + || request + .params + .output_format + .as_ref() + .is_none_or(Value::is_object) + { + return request; + } + MessagesRequest { + params: MessagesOptionalParams { + output_format: None, + ..request.params + }, + ..request + } +} + +fn is_python_str(value: &Value) -> bool { + match value { + Value::String(text) => { + litellm_host::raw_json::decode(text).is_none_or(|raw| raw.starts_with('"')) + } + _ => false, + } +} + +fn is_non_string_text(part: &Value) -> bool { + part.get("type").and_then(Value::as_str) == Some("text") + && part.get("text").is_some_and(|text| !is_python_str(text)) +} + +fn without_non_string_texts(content: &mut Value) { + if let Value::Array(parts) = content { + parts.retain(|part| !is_non_string_text(part)); + } +} + +fn drop_non_string_tool_result_texts( + messages: Vec>, +) -> Vec> { + messages + .into_iter() + .map(|message| match message { + Recognized::Known(mut message) => { + if message.role == "user" + && let MessageContent::Blocks(blocks) = &mut message.content + { + blocks + .iter_mut() + .filter(|block| block.is_type(ContentBlockType::ToolResult)) + .filter_map(|block| block.content.as_mut()) + .for_each(without_non_string_texts); + } + Recognized::Known(message) + } + Recognized::Unrecognized(mut message) => { + if message.get("role").and_then(Value::as_str) == Some("user") + && let Some(Value::Array(blocks)) = message.get_mut("content") + { + blocks + .iter_mut() + .filter(|block| { + block.get("type").and_then(Value::as_str) == Some("tool_result") + }) + .filter_map(|block| block.get_mut("content")) + .for_each(without_non_string_texts); + } + Recognized::Unrecognized(message) + } + }) + .collect() +} + fn without_null_safeguards(extra: Map) -> Map { extra .into_iter() @@ -193,22 +443,45 @@ fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet { .collect() } +fn continues_message_thread(params: &MessagesOptionalParams) -> bool { + let Some(Value::Object(thread)) = params.extra.get("thread") else { + return false; + }; + thread.get("type").and_then(Value::as_str) == Some("continue") + && thread + .get("previous_message_id") + .is_none_or(Value::is_string) +} + +fn is_threaded(params: &MessagesOptionalParams) -> bool { + params + .extra + .get("thread") + .is_some_and(|thread| !thread.is_null()) +} + fn feature_betas(request: &MessagesRequest) -> BetaSet { let params = &request.params; let tools = params.tools.as_deref(); + let scanned_messages: &[Recognized] = if continues_message_thread(params) { + &[] + } else { + &request.messages + }; [ - requires_native_compaction_beta(params.compaction.as_ref(), &request.messages) + requires_native_compaction_beta(params.compaction.as_ref(), scanned_messages) .then_some(AnthropicBeta::Compact20260904), uses_structured_output(params).then_some(AnthropicBeta::StructuredOutputs20251113), (params.speed == Some(Recognized::Known(Speed::Fast))) .then_some(AnthropicBeta::FastMode20260201), - messages_carry_output_config(&request.messages) + messages_carry_output_config(scanned_messages) .then_some(AnthropicBeta::PerTurnControl20260701), + is_threaded(params).then_some(AnthropicBeta::MessageThreads20260812), uses_thinking_display_updates(params) .then_some(AnthropicBeta::ThinkingDisplayUpdates20260818), uses_thinking_block_binding(params) .then_some(AnthropicBeta::ThinkingBindingControls20260801), - messages_carry_tool_changes(&request.messages) + messages_carry_tool_changes(scanned_messages) .then_some(AnthropicBeta::MidConversationToolChanges20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), @@ -254,9 +527,10 @@ fn carries_safeguards(params: &MessagesOptionalParams) -> bool { .is_some_and(|value| !value.is_null()) } -fn messages_carry_output_config(messages: &[Message]) -> bool { +fn messages_carry_output_config(messages: &[Recognized]) -> bool { messages .iter() + .filter_map(Recognized::known) .any(|message| message.extra.contains_key("output_config")) } @@ -301,9 +575,10 @@ fn is_tool_reference_change(block: &ContentBlock) -> bool { == Some("tool_reference") } -fn messages_carry_tool_changes(messages: &[Message]) -> bool { +fn messages_carry_tool_changes(messages: &[Recognized]) -> bool { messages .iter() + .filter_map(Recognized::known) .filter(|message| message.role == "system") .flat_map(Message::blocks) .any(is_tool_reference_change) @@ -334,7 +609,7 @@ pub fn drop_unsupported_params( let params = request.params; let speed = match ¶ms.speed { Some(speed) if !capabilities.supports_speed => { - reject("speed", format!("'{}'", speed_text(speed)), "")?; + reject("speed", speed_repr(speed), "")?; None } _ => params.speed.clone(), @@ -346,21 +621,21 @@ pub fn drop_unsupported_params( }); } let temperature = match params.temperature { - Some(temperature) if temperature != 1.0 => { + Some(temperature) if !is_one(&temperature) => { reject( "temperature", - json!(temperature).to_string(), + python_str(&temperature), "Only temperature=1 is supported. ", )?; None } temperature => temperature, }; - if let Some(top_p) = params.top_p { - reject("top_p", json!(top_p).to_string(), "")?; + if let Some(top_p) = ¶ms.top_p { + reject("top_p", python_str(top_p), "")?; } - if let Some(top_k) = params.top_k { - reject("top_k", json!(top_k).to_string(), "")?; + if let Some(top_k) = ¶ms.top_k { + reject("top_k", python_str(top_k), "")?; } Ok(MessagesRequest { params: MessagesOptionalParams { @@ -374,11 +649,10 @@ pub fn drop_unsupported_params( }) } -fn speed_text(speed: &Recognized) -> String { +fn speed_repr(speed: &Recognized) -> String { match speed { - Recognized::Known(speed) => <&'static str>::from(*speed).to_string(), - Recognized::Unrecognized(Value::String(text)) => text.clone(), - Recognized::Unrecognized(other) => other.to_string(), + Recognized::Known(speed) => format!("'{}'", <&'static str>::from(*speed)), + Recognized::Unrecognized(other) => litellm_python_compat::repr::repr(&python_value(other)), } } @@ -437,6 +711,7 @@ mod tests { use super::*; use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX; + use litellm_http::request::header_values; type Env = &'static [(&'static str, &'static str)]; @@ -679,6 +954,21 @@ mod tests { json!({"speed": "fast"}), "claude does not support speed='fast'. To drop unsupported params, set `litellm.drop_params = True`." )] + #[case::integer_temperature_keeps_no_decimal( + sampling_removed(), + json!({"temperature": 0}), + "claude does not support temperature=0. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::integer_top_p( + sampling_removed(), + json!({"top_p": 1}), + "claude does not support top_p=1. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::non_string_speed_is_reprd( + unmapped(), + json!({"speed": 2}), + "claude does not support speed=2. To drop unsupported params, set `litellm.drop_params = True`." + )] #[case::speed_ahead_of_sampling_params( sampling_removed(), json!({"speed": "fast", "temperature": 0.5}), @@ -696,6 +986,28 @@ mod tests { assert_eq!(detail.to_string(), message); } + #[rstest] + #[case::string_dropped(json!("json"), true, None)] + #[case::string_kept_without_drop_params(json!("json"), false, Some(json!("json")))] + #[case::object_kept(json!({"type": "json_schema"}), true, Some(json!({"type": "json_schema"})))] + fn a_non_object_output_format_is_dropped_under_drop_params( + #[case] output_format: Value, + #[case] drops_unsupported_params: bool, + #[case] expected: Option, + ) { + let context = MessagesTransformContext { + drops_unsupported_params, + ..MessagesTransformContext::with_lookup(unmapped(), false, &no_env) + }; + let transformed = ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request(json!({"output_format": output_format})), + &context, + ) + .unwrap(); + assert_eq!(transformed.params.output_format, expected); + } + #[rstest] #[case::compaction_threshold( json!([{"type": "compaction", "compact_threshold": 200000}]), @@ -802,6 +1114,40 @@ mod tests { ); } + #[rstest] + #[case::continue_turn(json!({"type": "continue", "previous_message_id": "msg_1"}), false)] + #[case::continue_without_a_previous_message_id(json!({"type": "continue"}), false)] + #[case::continue_with_extra_keys(json!({"type": "continue", "previous_message_id": "msg_1", "note": 1}), false)] + #[case::create_turn(json!({"type": "create"}), true)] + #[case::continue_with_a_null_previous_message_id(json!({"type": "continue", "previous_message_id": null}), true)] + #[case::continue_with_a_numeric_previous_message_id(json!({"type": "continue", "previous_message_id": 1}), true)] + #[case::not_an_object(json!("continue"), true)] + fn advisor_history_is_kept_on_a_thread_continue_turn( + #[case] thread: Value, + #[case] stripped: bool, + unmapped: MessagesModelCapabilities, + ) { + let messages = if stripped { + json!([ + {"role": "user", "content": "Build a worker pool."}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "text", "text": "Here is the implementation."} + ]} + ]) + } else { + advisor_history() + }; + assert_eq!( + transform( + json!({"messages": advisor_history(), "thread": thread.clone()}), + unmapped, + false + ), + Ok(body(json!({"messages": messages, "thread": thread}))) + ); + } + #[rstest] fn bridge_minted_reasoning_is_stripped_from_the_wire(unmapped: MessagesModelCapabilities) { let messages = json!([ @@ -864,6 +1210,7 @@ mod tests { keeps_display_updates: false, }, drop_params: true, + drops_unsupported_params: false, } ); return; @@ -1165,6 +1512,55 @@ mod tests { ]}), &[] )] + #[case::thread_create(json!({"thread": {"type": "create"}}), &["message-threads-2026-08-12"])] + #[case::thread_continue( + json!({"thread": {"type": "continue", "previous_message_id": "msg_1"}}), + &["message-threads-2026-08-12"] + )] + #[case::thread_of_any_shape(json!({"thread": "opaque"}), &["message-threads-2026-08-12"])] + #[case::null_thread(json!({"thread": null}), &[])] + #[case::thread_create_still_scans_history( + json!({ + "thread": {"type": "create"}, + "messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, + {"role": "user", "content": "hi", "output_config": {"effort": "low"}}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + }), + &[ + "compact-2026-09-04", + "message-threads-2026-08-12", + "mid-conversation-tool-changes-2026-07-01", + "per-turn-control-2026-07-01", + ] + )] + #[case::thread_continue_skips_history( + json!({ + "thread": {"type": "continue", "previous_message_id": "msg_1"}, + "messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, + {"role": "user", "content": "hi", "output_config": {"effort": "low"}}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + ], + }), + &["message-threads-2026-08-12"] + )] + #[case::thread_continue_keeps_param_betas( + json!({ + "thread": {"type": "continue", "previous_message_id": "msg_1"}, + "compaction": {}, + "speed": "fast", + }), + &["compact-2026-09-04", "fast-mode-2026-02-01", "message-threads-2026-08-12"] + )] + #[case::malformed_continue_scans_history( + json!({ + "thread": {"type": "continue", "previous_message_id": 7}, + "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], + }), + &["message-threads-2026-08-12", "per-turn-control-2026-07-01"] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } @@ -1352,4 +1748,158 @@ mod tests { .collect(); assert_eq!(undeclared, Vec::<&String>::new()); } + + const FOUNDRY: &str = "https://team.services.ai.azure.com/anthropic/v1/messages"; + + #[rstest] + #[case::drops_a_beta_foundry_lacks( + FOUNDRY, + "advisor-tool-2026-03-01,context-1m-2025-08-07", + Some("context-1m-2025-08-07") + )] + #[case::forwards_unknown_betas_sorted( + FOUNDRY, + "zzz,files-api-2025-04-14,aaa", + Some("aaa,files-api-2025-04-14,zzz") + )] + #[case::removes_the_header_when_nothing_is_left(FOUNDRY, "advisor-tool-2026-03-01", None)] + #[case::cognitive_services_host( + "team.cognitiveservices.azure.com/anthropic", + "advisor-tool-2026-03-01", + None + )] + #[case::anthropic_host_is_untouched( + "https://api.anthropic.com/v1/messages", + "advisor-tool-2026-03-01", + Some("advisor-tool-2026-03-01") + )] + #[case::lookalike_host_is_untouched( + "https://notservices.ai.azure.com.example/v1/messages", + "advisor-tool-2026-03-01", + Some("advisor-tool-2026-03-01") + )] + fn foundry_urls_filter_betas_through_the_azure_ai_row( + #[case] url: &str, + #[case] betas: &str, + #[case] expected: Option<&str>, + ) { + let filtered = ANTHROPIC_MESSAGES_CONFIG + .headers_for_url(headers(&[("x-a", "1"), (BETA_HEADER, betas)]), url); + assert_eq!( + header_values(&filtered, BETA_HEADER).collect::>(), + expected.into_iter().collect::>() + ); + assert_eq!(header_values(&filtered, "x-a").collect::>(), ["1"]); + } + + #[rstest] + fn foundry_url_without_betas_keeps_the_headers() { + let sent = headers(&[("x-a", "1")]); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.headers_for_url(sent.clone(), FOUNDRY), + sent + ); + } + + #[test] + fn the_http_error_retry_matches_the_cases_shared_with_python() { + let fixture: Value = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/anthropic_messages_http_error_retry.json" + ))) + .unwrap(); + let table = &fixture["http_error_retry"]; + let cases = table["cases"].as_array().unwrap(); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let mut body = table["request"].clone(); + body["thinking"] = case["thinking"].clone(); + let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap(); + let error_text = case["error_text"].as_str().unwrap(); + let retried = should_retry_anthropic_messages_on_http_error( + status, + error_text, + &body, + &MessagesTransformContext::default(), + ); + let stripped = case["stripped"].as_bool().unwrap(); + let transformed = transform_anthropic_messages_request_on_http_error( + status, + error_text, + body.clone(), + ); + let expected = if stripped { + RetriedRequest { + reasons: vec![RetryReason::ThinkingSignature], + body: table["stripped_request"].clone(), + } + } else { + RetriedRequest { + reasons: Vec::new(), + body, + } + }; + retried != case["retried"].as_bool().unwrap() || transformed != expected + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); + } + + #[test] + fn the_strict_tools_retry_matches_the_cases_shared_with_python() { + let fixture: Value = serde_json::from_str(include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/anthropic_messages_http_error_retry.json" + ))) + .unwrap(); + let replay = &fixture["http_error_retry"]; + let table = &fixture["strict_tools_retry"]; + let cases = table["cases"].as_array().unwrap(); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let mut body = replay["request"].clone(); + body["thinking"] = case["thinking"].clone(); + body["tools"] = table["tools"].clone(); + let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap(); + let error_text = case["error_text"].as_str().unwrap(); + let context = MessagesTransformContext { + drops_unsupported_params: case["drops_unsupported_params"].as_bool().unwrap(), + ..MessagesTransformContext::default() + }; + let mut reasons = Vec::new(); + let mut expected = if case["stripped"].as_bool().unwrap() { + reasons.push(RetryReason::ThinkingSignature); + replay["stripped_request"].clone() + } else { + body.clone() + }; + expected["tools"] = if case["strict_dropped"].as_bool().unwrap() { + reasons.push(RetryReason::StrictTools { + dropped: usize::try_from(table["dropped"].as_u64().unwrap()).unwrap(), + }); + table["tools_without_strict"].clone() + } else { + table["tools"].clone() + }; + let retried = should_retry_anthropic_messages_on_http_error( + status, error_text, &body, &context, + ); + let transformed = + transform_anthropic_messages_request_on_http_error(status, error_text, body); + retried != case["retried"].as_bool().unwrap() + || transformed + != RetriedRequest { + reasons, + body: expected, + } + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); + } } diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index d13cfa81a94..fb50d229bc7 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,3 +1,5 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; +use litellm_host::interceptors::RetriedRequest; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, @@ -5,6 +7,7 @@ use litellm_llms_types::{ }; use litellm_router_types::LitellmParams; use litellm_secret_value::SecretValue; +use serde_json::Value; use crate::{ Error, @@ -21,7 +24,7 @@ use crate::{ base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, normalization::{normalize_system_role_messages, strip_cache_control_scope}, transformation::BaseMessagesConfig, }, @@ -36,12 +39,28 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = AzureAnthropicMessagesConfig; impl BaseMessagesConfig for AzureAnthropicMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index 3d8b0f41d8c..e5593d3b216 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -23,6 +23,12 @@ impl SupportedEffortTiers { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MessagesShapeOptions { + pub reasoning_auto_summary: bool, + pub presanitized: bool, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct MessagesModelCapabilities { #[serde(default)] @@ -123,6 +129,7 @@ pub struct ThinkingContext { pub struct MessagesTransformContext { pub thinking: ThinkingContext, pub drop_params: bool, + pub drops_unsupported_params: bool, } impl MessagesTransformContext { pub fn new(capabilities: MessagesModelCapabilities, drop_params: bool) -> Self { @@ -142,6 +149,7 @@ impl MessagesTransformContext { keeps_display_updates: false, }, drop_params, + drops_unsupported_params: false, } } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs index e9591daa9da..5c1563c7a23 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs @@ -2,4 +2,5 @@ pub mod context; pub mod mid_conversation_system; pub mod normalization; pub mod streaming; +pub mod tools; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs index 7e24da2b5eb..ee2f4dcaeef 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/normalization.rs @@ -1,6 +1,9 @@ -use litellm_llms_types::formats::messages::{ - CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams, MessagesRequest, - SystemPrompt, +use litellm_llms_types::{ + formats::messages::{ + CacheControl, ContentBlock, Message, MessageContent, MessagesOptionalParams, + MessagesRequest, SystemPrompt, + }, + recognized::{Recognized, map_known_runs}, }; use crate::{ @@ -22,15 +25,22 @@ pub fn normalize_system_role_messages( let leading_count = request .messages .iter() - .take_while(|message| is_system_role_message(message)) + .take_while(|message| message.known().is_some_and(is_system_role_message)) .count(); let mut messages = request.messages.into_iter(); - let hoisted: Vec = messages.by_ref().take(leading_count).collect(); - let remaining: Vec = messages.collect(); + let hoisted: Vec = messages + .by_ref() + .take(leading_count) + .filter_map(|message| match message { + Recognized::Known(message) => Some(message), + Recognized::Unrecognized(_) => None, + }) + .collect(); + let remaining: Vec> = messages.collect(); let remaining = if supports_mid_conversation_system { remaining } else { - convert_mid_conversation_system_turns(remaining) + map_known_runs(remaining, convert_mid_conversation_system_turns) }; let system = if hoisted.is_empty() { request.params.system @@ -94,7 +104,7 @@ pub fn strip_cache_control_scope(request: MessagesRequest) -> MessagesRequest { messages: request .messages .into_iter() - .map(strip_scope_from_message) + .map(|message| message.map_known(strip_scope_from_message)) .collect(), params: MessagesOptionalParams { system: request.params.system.map(strip_scope_from_system), diff --git a/litellm-rust/crates/llms/src/base_llm/messages/tools.rs b/litellm-rust/crates/llms/src/base_llm/messages/tools.rs new file mode 100644 index 00000000000..49113872c56 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/messages/tools.rs @@ -0,0 +1,194 @@ +use litellm_llms_types::{ + formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, + recognized::Recognized, +}; +use serde_json::{Map, Value}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ToolCustomField { + Keep, + Rewrite { hoist_defer_loading: bool }, +} + +impl ToolCustomField { + pub fn apply(self, request: MessagesRequest, drop_params: bool) -> MessagesRequest { + match self { + Self::Keep => request, + Self::Rewrite { + hoist_defer_loading, + } => tools_without_custom_field(request, hoist_defer_loading, drop_params), + } + } +} + +fn tools_without_custom_field( + request: MessagesRequest, + hoist: bool, + drop_custom: bool, +) -> MessagesRequest { + let tools = request.params.tools.map(|tools| { + tools + .into_iter() + .map(|tool| tool_without_custom_field(tool, hoist, drop_custom)) + .collect() + }); + MessagesRequest { + params: MessagesOptionalParams { + tools, + ..request.params + }, + ..request + } +} + +fn tool_without_custom_field( + tool: Recognized, + hoist: bool, + drop_custom: bool, +) -> Recognized { + let Ok(Value::Object(fields)) = serde_json::to_value(&tool) else { + return tool; + }; + let Some(custom) = fields.get("custom").cloned() else { + return tool; + }; + let defer_loading = hoist + .then(|| hoistable_defer_loading(&fields, &custom)) + .flatten(); + let without_custom = fields + .iter() + .filter(|(key, _)| *key != "custom") + .map(|(key, value)| (key.clone(), value.clone())); + let rebuilt: Map = match defer_loading { + None if drop_custom => without_custom.collect(), + None => return tool, + Some(defer_loading) => { + let unhoisted: Map = custom + .as_object() + .into_iter() + .flatten() + .filter(|(key, _)| *key != "defer_loading") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let kept_custom = (!unhoisted.is_empty() && !drop_custom) + .then(|| ("custom".to_string(), Value::Object(unhoisted))); + without_custom + .chain([("defer_loading".to_string(), Value::Bool(defer_loading))]) + .chain(kept_custom) + .collect() + } + }; + serde_json::from_value(Value::Object(rebuilt)).unwrap_or(tool) +} + +fn hoistable_defer_loading(tool: &Map, custom: &Value) -> Option { + if tool.contains_key("defer_loading") { + return None; + } + custom.as_object()?.get("defer_loading")?.as_bool() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + const WEATHER: &str = "get_weather"; + + fn tool(fields: Value) -> Value { + merged( + json!({"name": WEATHER, "input_schema": {"type": "object"}}), + fields, + ) + } + + fn tools_after(policy: ToolCustomField, tools: Value, drop_custom: bool) -> Value { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": [tools] + })) + .unwrap(); + serde_json::to_value(policy.apply(request, drop_custom)).unwrap()["tools"][0].clone() + } + + fn merged(base: Value, fields: Value) -> Value { + let mut merged = base.as_object().unwrap().clone(); + merged.extend(fields.as_object().unwrap().clone()); + Value::Object(merged) + } + + const HOIST: ToolCustomField = ToolCustomField::Rewrite { + hoist_defer_loading: true, + }; + const NO_HOIST: ToolCustomField = ToolCustomField::Rewrite { + hoist_defer_loading: false, + }; + + #[rstest] + #[case::hoists_defer_loading( + HOIST, + json!({"custom": {"defer_loading": true}}), + false, + json!({"defer_loading": true}) + )] + #[case::keeps_other_custom_keys( + HOIST, + json!({"custom": {"defer_loading": false, "note": "kept"}}), + false, + json!({"defer_loading": false, "custom": {"note": "kept"}}) + )] + #[case::drops_other_custom_keys_with_drop_params( + HOIST, + json!({"custom": {"defer_loading": true, "note": "dropped"}}), + true, + json!({"defer_loading": true}) + )] + #[case::keeps_an_opaque_custom( + HOIST, + json!({"custom": "opaque"}), + false, + json!({"custom": "opaque"}) + )] + #[case::drops_an_opaque_custom_with_drop_params( + HOIST, + json!({"custom": "opaque"}), + true, + json!({}) + )] + #[case::own_defer_loading_wins( + HOIST, + json!({"defer_loading": false, "custom": {"defer_loading": true}}), + false, + json!({"defer_loading": false, "custom": {"defer_loading": true}}) + )] + #[case::not_hoisted_off_claude_hosts( + NO_HOIST, + json!({"custom": {"defer_loading": true}}), + false, + json!({"custom": {"defer_loading": true}}) + )] + fn tool_custom_field_follows_the_python_handler( + #[case] policy: ToolCustomField, + #[case] fields: Value, + #[case] drop_custom: bool, + #[case] expected: Value, + ) { + assert_eq!( + tools_after(policy, tool(fields), drop_custom), + tool(expected) + ); + } + + #[rstest] + fn keep_leaves_the_custom_field() { + let custom = tool(json!({"custom": {"defer_loading": true}})); + assert_eq!( + tools_after(ToolCustomField::Keep, custom.clone(), true), + custom + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 744ce89afdb..13c021c2aec 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -1,19 +1,33 @@ +use litellm_host::interceptors::RetriedRequest; use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use serde_json::Value; -use super::context::MessagesTransformContext; +use super::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + tools::ToolCustomField, +}; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{Error, base_llm::messages::streaming::StreamDecoder}; +pub enum FederatedEnvironment { + Validated(ValidatedEnvironment), + Unauthenticated { + headers: Headers, + missing: Error, + federate: fn(Headers, SecretValue) -> ValidatedEnvironment, + }, +} + pub trait BaseMessagesConfig: Sync { fn shape_request( &self, request: MessagesRequest, - _reasoning_auto_summary: bool, + _options: MessagesShapeOptions, ) -> Result { Ok(request) } @@ -35,6 +49,10 @@ pub trait BaseMessagesConfig: Sync { Ok(request) } + fn tool_custom_field(&self, _model: &str) -> ToolCustomField { + ToolCustomField::Keep + } + fn transform_anthropic_messages_response( &self, _model: &str, @@ -62,6 +80,26 @@ pub trait BaseMessagesConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn validate_federated_environment( + &self, + headers: Headers, + api_key: Option<&str>, + model: &str, + optional_params: &MessagesOptionalParams, + litellm_params: &LitellmParams, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + self.validate_environment( + headers, + api_key, + model, + optional_params, + litellm_params, + env_lookup, + ) + .map(FederatedEnvironment::Validated) + } + /// `None` relays the upstream bytes untouched, which is right for every host that already /// speaks Anthropic SSE. A host on another wire returns the decoder that lifts its frames /// into Anthropic stream events, and the route re-encodes those as Anthropic SSE. @@ -77,6 +115,10 @@ pub trait BaseMessagesConfig: Sync { headers } + fn headers_for_url(&self, headers: Headers, _url: &str) -> Headers { + headers + } + fn request_body( &self, request: MessagesRequest, @@ -91,6 +133,16 @@ pub trait BaseMessagesConfig: Sync { fn wire_body(&self, body: Value) -> Value { body } + + fn retry_after_http_error( + &self, + _status: u16, + _error_text: &str, + _request_body: &Value, + _context: &MessagesTransformContext, + ) -> Option { + None + } } #[cfg(test)] @@ -149,9 +201,10 @@ mod tests { } #[rstest] - #[case::disabled(false)] - #[case::enabled(true)] - fn default_shaping_preserves_provider_policy_inputs(#[case] reasoning_auto_summary: bool) { + fn default_shaping_preserves_provider_policy_inputs( + #[values(false, true)] reasoning_auto_summary: bool, + #[values(false, true)] presanitized: bool, + ) { let request: MessagesRequest = serde_json::from_value(serde_json::json!({ "model": "test-model", "metadata": {"user_id": 7, "extra": "keep"}, @@ -160,7 +213,13 @@ mod tests { })) .unwrap(); assert_eq!( - DefaultsConfig.shape_request(request.clone(), reasoning_auto_summary), + DefaultsConfig.shape_request( + request.clone(), + MessagesShapeOptions { + reasoning_auto_summary, + presanitized + } + ), Ok(request) ); assert_eq!( diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 5b89345a1c5..4bf30b56f00 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,3 +1,5 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; +use litellm_host::interceptors::RetriedRequest; use std::sync::LazyLock; use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec}; @@ -21,7 +23,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, bedrock::common_utils::validated_region, }; @@ -50,12 +55,28 @@ pub const BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG: BedrockClaudePlatformMessages BedrockClaudePlatformMessagesConfig; impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 1d380c917f6..fa037974342 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -1,9 +1,11 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_core_utils::decline::Decline; +use litellm_host::interceptors::RetriedRequest; use std::convert::Infallible; use crate::{ anthropic::messages::handler::shape_anthropic_messages_request, - base_llm::messages::context::MessagesTransformContext, + base_llm::messages::context::{MessagesShapeOptions, MessagesTransformContext}, }; use futures_util::StreamExt; use litellm_auth::AwsParams; @@ -103,12 +105,22 @@ fn invoke_url( } impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index f139b93e259..0e67ea1a139 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -1,3 +1,5 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; +use litellm_host::interceptors::RetriedRequest; use std::sync::LazyLock; use litellm_auth::{AwsParams, CredentialPlacement}; @@ -17,6 +19,7 @@ use serde_json::Value; use crate::{ Error, anthropic::{ + beta_headers::accepted_or_unlisted, common_utils::{existing_betas, merge_beta_headers}, messages::{ handler::shape_anthropic_messages_request, @@ -26,7 +29,7 @@ use crate::{ base_llm::{ auth::AuthScheme, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, }, }, @@ -56,6 +59,16 @@ pub const BEDROCK_MANTLE_MESSAGES_CONFIG: BedrockMantleAnthropicMessagesConfig = BedrockMantleAnthropicMessagesConfig; impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn validate_environment( &self, headers: Headers, @@ -129,9 +142,9 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn wire_body(&self, body: Value) -> Value { @@ -191,14 +204,10 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { .filter_map(Value::as_str) .filter_map(|beta| beta.parse::().ok()) .collect(); - let caller: BetaSet = existing_betas(&headers) - .union(body_betas) - .into_iter() - .filter_map(|beta| match beta { - AnthropicBeta::Other(_) => Some(beta), - _ => beta.on(BetaProvider::BedrockMantle), - }) - .collect(); + let caller = accepted_or_unlisted( + existing_betas(&headers).union(body_betas), + BetaProvider::BedrockMantle, + ); let safeguards: BetaSet = request .params .extra diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 0dbcea01439..ee2aec9efb9 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -1,4 +1,6 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::RetriedRequest; use litellm_secret_value::SecretValue; use litellm_core_utils::settings::resolve_non_empty; @@ -21,7 +23,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -42,12 +47,28 @@ pub const DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG: DeepSeekAnthropicMessagesConfig = DeepSeekAnthropicMessagesConfig; impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index d5944484522..d4f577e09a9 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -1,4 +1,6 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::RetriedRequest; use litellm_secret_value::SecretValue; use litellm_core_utils::settings::resolve_non_empty; @@ -33,6 +35,22 @@ pub struct EdenAIAnthropicMessagesConfig; pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig; impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn validate_environment( &self, headers: Headers, diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 2975ceee3cb..77ee1f1036f 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -168,6 +168,10 @@ pub enum ErrorDetail { actual: String, choices: Vec<&'static str>, }, + #[error( + "Invalid reasoning_effort: {actual}. Must be one of: 'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'" + )] + InvalidReasoningEffort { actual: String }, #[error("{field}='{value}' is not supported by this model. Got model: {model}")] UnsupportedValue { field: &'static str, @@ -187,6 +191,14 @@ pub enum ErrorDetail { "thinking.budget_tokens must be a number, got {actual}. To drop it, set `litellm.drop_params=True`, or `drop_params: true` under `litellm_settings` on the proxy" )] NonNumericThinkingBudget { actual: &'static str }, + /// Python's comparison of two operands it cannot order, which raises `TypeError` with this + /// text before any request is sent. + #[error("'{op}' not supported between instances of '{left}' and '{right}'")] + UnorderableTypes { + op: &'static str, + left: &'static str, + right: &'static str, + }, #[error( "Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params" )] diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index dafe5fbd16b..2eaf182e657 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -1,5 +1,8 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::RetriedRequest; use litellm_secret_value::SecretValue; +use serde_json::Value; use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ @@ -19,7 +22,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -31,12 +37,28 @@ pub struct MinimaxMessagesConfig; pub const MINIMAX_MESSAGES_CONFIG: MinimaxMessagesConfig = MinimaxMessagesConfig; impl BaseMessagesConfig for MinimaxMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index eeb3d1cb7f8..dc09f4e4149 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -1,5 +1,8 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::RetriedRequest; use litellm_secret_value::SecretValue; +use serde_json::Value; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{has_bearer_auth, without_headers}; @@ -31,6 +34,16 @@ pub const OPENROUTER_MESSAGES_CONFIG: OpenRouterAnthropicMessagesConfig = OpenRouterAnthropicMessagesConfig; impl BaseMessagesConfig for OpenRouterAnthropicMessagesConfig { + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn validate_environment( &self, headers: Headers, diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index 621a09f4fb0..b3ee28b7474 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -1,5 +1,8 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::CredentialPlacement; +use litellm_host::interceptors::RetriedRequest; use litellm_secret_value::SecretValue; +use serde_json::Value; use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ @@ -19,7 +22,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -32,12 +38,28 @@ pub struct TencentAnthropicMessagesConfig; pub const TENCENT_MESSAGES_CONFIG: TencentAnthropicMessagesConfig = TencentAnthropicMessagesConfig; impl BaseMessagesConfig for TencentAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 038617c2e10..674b9951d3e 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -1,8 +1,10 @@ +use crate::anthropic::messages::transformation::retry_anthropic_messages_after_http_error; use litellm_auth::VertexParams; use litellm_auth_gcp::{ VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, secret_names, }; +use litellm_host::interceptors::RetriedRequest; use litellm_http::request::{has_header, header_value, without_headers}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool, OutputConfig}, @@ -24,7 +26,7 @@ use crate::{ base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, normalization::{normalize_system_role_messages, strip_cache_control_scope}, transformation::BaseMessagesConfig, }, @@ -47,12 +49,28 @@ pub const VERTEX_ANTHROPIC_MESSAGES_CONFIG: VertexAiPartnerModelsAnthropicMessag VertexAiPartnerModelsAnthropicMessagesConfig; impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + + fn retry_after_http_error( + &self, + status: u16, + error_text: &str, + request_body: &Value, + context: &MessagesTransformContext, + ) -> Option { + retry_anthropic_messages_after_http_error(status, error_text, request_body, context) + } + fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( @@ -700,6 +718,7 @@ mod tests { ..ThinkingContext::default() }, drop_params: false, + drops_unsupported_params: false, }, ); assert_eq!(body["system"], json!("be terse")); @@ -717,6 +736,7 @@ mod tests { ..ThinkingContext::default() }, drop_params: false, + drops_unsupported_params: false, } } diff --git a/litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json b/litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json new file mode 100644 index 00000000000..843a1689056 --- /dev/null +++ b/litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json @@ -0,0 +1,623 @@ +{ + "invalid_thinking_block_error": [ + { + "name": "invalid_signature", + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "expected": true + }, + { + "name": "invalid_signature_in_an_error_envelope", + "error_text": "{\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"messages.3.content.3: Invalid `signature` in `thinking` block\"},\"request_id\":\"req_011Ca2EtQDxp7x6RGUY2jVn9\"}", + "expected": true + }, + { + "name": "signature_not_a_string", + "error_text": "{\"message\":\"messages.2.content.0.thinking.signature.str: Input should be a valid string\"}", + "expected": true + }, + { + "name": "signature_not_a_string_bare", + "error_text": "messages.4.content.1.thinking.signature.str: Input should be a valid string", + "expected": true + }, + { + "name": "empty_thinking", + "error_text": "messages.1.content.0.thinking: each thinking block must contain thinking", + "expected": true + }, + { + "name": "bound_to_a_different_conversation", + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.", + "expected": true + }, + { + "name": "uppercase", + "error_text": "MESSAGES.1: INVALID SIGNATURE IN THINKING BLOCK", + "expected": true + }, + { + "name": "signature_without_a_verdict", + "error_text": "thinking signature is malformed", + "expected": false + }, + { + "name": "invalid_without_thinking", + "error_text": "Invalid `signature` in `tool_use` block", + "expected": false + }, + { + "name": "model_not_found", + "error_text": "invalid_request_error: model not found", + "expected": false + }, + { + "name": "rate_limited", + "error_text": "rate limit exceeded", + "expected": false + }, + { + "name": "unrelated", + "error_text": "max_tokens: 4096 > 64000", + "expected": false + }, + { + "name": "empty", + "error_text": "", + "expected": false + } + ], + "thinking_binding_mismatch_error": [ + { + "name": "bound_elsewhere", + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.", + "expected": true + }, + { + "name": "mixed_case", + "error_text": "The block is Bound To A Different Conversation.", + "expected": true + }, + { + "name": "plain_signature", + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "expected": false + } + ], + "http_error_retry": { + "request": { + "model": "claude-sonnet-4-5", + "max_tokens": 4096, + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "plan", + "signature": "sig" + }, + { + "type": "text", + "text": "hello" + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "redacted_thinking", + "data": "opaque" + } + ] + }, + { + "role": "user", + "content": "again" + } + ] + }, + "stripped_request": { + "model": "claude-sonnet-4-5", + "max_tokens": 4096, + "messages": [ + { + "role": "user", + "content": "hi" + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "hello" + } + ] + }, + { + "role": "user", + "content": "again" + } + ] + }, + "cases": [ + { + "name": "invalid_signature", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": true, + "stripped": true + }, + { + "name": "missing_signature", + "status": 400, + "error_text": "messages.1.content.0.thinking.signature.str: Input should be a valid string", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": true, + "stripped": true + }, + { + "name": "empty_thinking", + "status": 400, + "error_text": "messages.1.content.0.thinking: each thinking block must contain thinking", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": true, + "stripped": true + }, + { + "name": "unpinned_binding_mismatch", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.", + "thinking": { + "type": "adaptive" + }, + "retried": true, + "stripped": true + }, + { + "name": "caller_pinned_binding", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.", + "thinking": { + "type": "adaptive", + "block_binding": { + "prefix_mismatch_behavior": "error" + } + }, + "retried": false, + "stripped": true + }, + { + "name": "caller_pinned_binding_other_signature_error", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "thinking": { + "type": "adaptive", + "block_binding": { + "prefix_mismatch_behavior": "error" + } + }, + "retried": true, + "stripped": true + }, + { + "name": "not_a_400", + "status": 429, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": false, + "stripped": false + }, + { + "name": "unrelated_400", + "status": 400, + "error_text": "max_tokens: 4096 > 64000", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": false, + "stripped": false + }, + { + "name": "signature_without_thinking", + "status": 400, + "error_text": "Invalid `signature` in `tool_use` block", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "retried": false, + "stripped": false + } + ] + }, + "claude_strict_tools_rejection": [ + { + "name": "too_many_strict_tools", + "status": 400, + "error_text": "Too many strict tools (21). The maximum number of strict tools is 20.", + "expected": true + }, + { + "name": "too_many_strict_tools_in_an_error_envelope", + "status": 400, + "error_text": "{\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Too many strict tools (24). The maximum is 20.\"}}", + "expected": true + }, + { + "name": "non_ascii_digits", + "status": 400, + "error_text": "Too many strict tools (٢١)", + "expected": true + }, + { + "name": "optional_parameters", + "status": 400, + "error_text": "Schemas contains too many optional parameters (30)", + "expected": true + }, + { + "name": "union_types", + "status": 400, + "error_text": "Schemas contains too many parameters with union types", + "expected": true + }, + { + "name": "compiled_grammar_too_large", + "status": 400, + "error_text": "The compiled grammar is too large", + "expected": true + }, + { + "name": "compiled_grammar_size", + "status": 400, + "error_text": "Compiled grammar size exceeds the limit", + "expected": true + }, + { + "name": "custom_strict", + "status": 400, + "error_text": "tools.3.custom.strict: Extra inputs are not permitted", + "expected": true + }, + { + "name": "custom", + "status": 400, + "error_text": "tools.0.custom: Extra inputs are not permitted", + "expected": true + }, + { + "name": "grammar_compilation_timed_out", + "status": 400, + "error_text": "Grammar compilation timed out", + "expected": true + }, + { + "name": "no_count", + "status": 400, + "error_text": "tools.0.strict: too many strict tools", + "expected": false + }, + { + "name": "empty_count", + "status": 400, + "error_text": "Too many strict tools ()", + "expected": false + }, + { + "name": "lowercase_grammar_timeout", + "status": 400, + "error_text": "grammar compilation timed out", + "expected": false + }, + { + "name": "custom_without_separator", + "status": 400, + "error_text": "tools.0.custom.strict", + "expected": false + }, + { + "name": "uppercase_schemas", + "status": 400, + "error_text": "SCHEMAS CONTAINS TOO MANY OPTIONAL PARAMETERS", + "expected": false + }, + { + "name": "not_a_400", + "status": 429, + "error_text": "Too many strict tools (21). The maximum number of strict tools is 20.", + "expected": false + }, + { + "name": "empty", + "status": 400, + "error_text": "", + "expected": false + } + ], + "tools_without_strict": [ + { + "name": "strict_dropped_from_every_tool", + "tools": [ + { + "name": "a", + "strict": true, + "input_schema": {} + }, + { + "name": "b", + "strict": false + }, + { + "name": "c" + } + ], + "expected": { + "tools": [ + { + "name": "a", + "input_schema": {} + }, + { + "name": "b" + }, + { + "name": "c" + } + ], + "dropped_count": 1 + } + }, + { + "name": "only_a_true_strict_counts", + "tools": [ + { + "name": "a", + "strict": 1 + }, + { + "name": "b", + "strict": "true" + } + ], + "expected": null + }, + { + "name": "no_strict", + "tools": [ + { + "name": "a" + } + ], + "expected": null + }, + { + "name": "empty", + "tools": [], + "expected": null + }, + { + "name": "non_object_tool", + "tools": [ + { + "name": "a", + "strict": true + }, + "b" + ], + "expected": null + }, + { + "name": "not_a_list", + "tools": { + "name": "a", + "strict": true + }, + "expected": null + }, + { + "name": "null", + "tools": null, + "expected": null + } + ], + "strict_tools_retry": { + "tools": [ + { + "name": "a", + "input_schema": { + "type": "object" + }, + "strict": true + }, + { + "name": "b", + "input_schema": { + "type": "object" + }, + "strict": false + }, + { + "name": "c", + "input_schema": { + "type": "object" + } + } + ], + "tools_without_strict": [ + { + "name": "a", + "input_schema": { + "type": "object" + } + }, + { + "name": "b", + "input_schema": { + "type": "object" + } + }, + { + "name": "c", + "input_schema": { + "type": "object" + } + } + ], + "dropped": 1, + "cases": [ + { + "name": "strict_rejection_under_drop_params", + "status": 400, + "error_text": "Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": true, + "retried": true, + "stripped": false, + "strict_dropped": true + }, + { + "name": "strict_rejection_without_drop_params", + "status": 400, + "error_text": "Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": false, + "retried": false, + "stripped": false, + "strict_dropped": true + }, + { + "name": "compiled_grammar_under_drop_params", + "status": 400, + "error_text": "The compiled grammar is too large", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": true, + "retried": true, + "stripped": false, + "strict_dropped": true + }, + { + "name": "strict_text_on_a_429", + "status": 429, + "error_text": "Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": true, + "retried": false, + "stripped": false, + "strict_dropped": false + }, + { + "name": "unrelated_400_under_drop_params", + "status": 400, + "error_text": "max_tokens: 4096 > 64000", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": true, + "retried": false, + "stripped": false, + "strict_dropped": false + }, + { + "name": "thinking_rejection_keeps_strict", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": true, + "retried": true, + "stripped": true, + "strict_dropped": false + }, + { + "name": "thinking_rejection_without_drop_params_keeps_strict", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": false, + "retried": true, + "stripped": true, + "strict_dropped": false + }, + { + "name": "one_error_naming_both_without_drop_params", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "enabled", + "budget_tokens": 2048 + }, + "drops_unsupported_params": false, + "retried": true, + "stripped": true, + "strict_dropped": true + }, + { + "name": "pinned_binding_mismatch_with_a_strict_rejection_under_drop_params", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation. Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "adaptive", + "block_binding": { + "prefix_mismatch_behavior": "error" + } + }, + "drops_unsupported_params": true, + "retried": true, + "stripped": true, + "strict_dropped": true + }, + { + "name": "pinned_binding_mismatch_with_a_strict_rejection_without_drop_params", + "status": 400, + "error_text": "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation. Too many strict tools (21). The maximum number of strict tools is 20.", + "thinking": { + "type": "adaptive", + "block_binding": { + "prefix_mismatch_behavior": "error" + } + }, + "drops_unsupported_params": false, + "retried": false, + "stripped": true, + "strict_dropped": true + } + ] + } +} diff --git a/litellm-rust/crates/llms/tests/messages_http_error_retry.rs b/litellm-rust/crates/llms/tests/messages_http_error_retry.rs new file mode 100644 index 00000000000..ac8ceac4ac6 --- /dev/null +++ b/litellm-rust/crates/llms/tests/messages_http_error_retry.rs @@ -0,0 +1,138 @@ +use litellm_host::interceptors::{RetriedRequest, RetryReason}; +use litellm_llms::{ + anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, + azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, + base_llm::messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + bedrock::{ + claude_platform::messages::transformation::BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG, + messages::invoke_transformations::anthropic_claude_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + }, + bedrock_mantle::messages::transformation::BEDROCK_MANTLE_MESSAGES_CONFIG, + deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, + edenai::messages::transformation::EDENAI_MESSAGES_CONFIG, + minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG, + openrouter::messages::transformation::OPENROUTER_MESSAGES_CONFIG, + tencent::messages::transformation::TENCENT_MESSAGES_CONFIG, + vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG, +}; +use rstest::rstest; +use serde_json::Value; + +#[rstest] +#[case::anthropic(&ANTHROPIC_MESSAGES_CONFIG)] +#[case::azure_ai(&AZURE_ANTHROPIC_MESSAGES_CONFIG)] +#[case::bedrock_invoke(&BEDROCK_ANTHROPIC_MESSAGES_CONFIG)] +#[case::bedrock_claude_platform(&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG)] +#[case::bedrock_mantle(&BEDROCK_MANTLE_MESSAGES_CONFIG)] +#[case::deepseek(&DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG)] +#[case::edenai(&EDENAI_MESSAGES_CONFIG)] +#[case::minimax(&MINIMAX_MESSAGES_CONFIG)] +#[case::openrouter(&OPENROUTER_MESSAGES_CONFIG)] +#[case::tencent(&TENCENT_MESSAGES_CONFIG)] +#[case::vertex_ai(&VERTEX_ANTHROPIC_MESSAGES_CONFIG)] +fn every_messages_provider_retries_like_the_python_base_config( + #[case] config: &dyn BaseMessagesConfig, +) { + let fixture: Value = serde_json::from_str(include_str!( + "fixtures/anthropic_messages_http_error_retry.json" + )) + .unwrap(); + let table = &fixture["http_error_retry"]; + let cases = table["cases"].as_array().unwrap(); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let mut body = table["request"].clone(); + body["thinking"] = case["thinking"].clone(); + let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap(); + let expected = case["retried"].as_bool().unwrap().then(|| { + if case["stripped"].as_bool().unwrap() { + RetriedRequest { + reasons: vec![RetryReason::ThinkingSignature], + body: table["stripped_request"].clone(), + } + } else { + RetriedRequest { + reasons: Vec::new(), + body: body.clone(), + } + } + }); + config.retry_after_http_error( + status, + case["error_text"].as_str().unwrap(), + &body, + &MessagesTransformContext::default(), + ) != expected + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); +} + +#[rstest] +#[case::anthropic(&ANTHROPIC_MESSAGES_CONFIG)] +#[case::azure_ai(&AZURE_ANTHROPIC_MESSAGES_CONFIG)] +#[case::bedrock_invoke(&BEDROCK_ANTHROPIC_MESSAGES_CONFIG)] +#[case::bedrock_claude_platform(&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG)] +#[case::bedrock_mantle(&BEDROCK_MANTLE_MESSAGES_CONFIG)] +#[case::deepseek(&DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG)] +#[case::edenai(&EDENAI_MESSAGES_CONFIG)] +#[case::minimax(&MINIMAX_MESSAGES_CONFIG)] +#[case::openrouter(&OPENROUTER_MESSAGES_CONFIG)] +#[case::tencent(&TENCENT_MESSAGES_CONFIG)] +#[case::vertex_ai(&VERTEX_ANTHROPIC_MESSAGES_CONFIG)] +fn every_messages_provider_drops_rejected_strict_tools_like_the_python_base_config( + #[case] config: &dyn BaseMessagesConfig, +) { + let fixture: Value = serde_json::from_str(include_str!( + "fixtures/anthropic_messages_http_error_retry.json" + )) + .unwrap(); + let replay = &fixture["http_error_retry"]; + let table = &fixture["strict_tools_retry"]; + let cases = table["cases"].as_array().unwrap(); + assert!(!cases.is_empty()); + let mismatched: Vec<&str> = cases + .iter() + .filter(|case| { + let mut body = replay["request"].clone(); + body["thinking"] = case["thinking"].clone(); + body["tools"] = table["tools"].clone(); + let status = u16::try_from(case["status"].as_u64().unwrap()).unwrap(); + let expected = case["retried"].as_bool().unwrap().then(|| { + let mut reasons = Vec::new(); + let mut expected = if case["stripped"].as_bool().unwrap() { + reasons.push(RetryReason::ThinkingSignature); + replay["stripped_request"].clone() + } else { + body.clone() + }; + expected["tools"] = if case["strict_dropped"].as_bool().unwrap() { + reasons.push(RetryReason::StrictTools { + dropped: usize::try_from(table["dropped"].as_u64().unwrap()).unwrap(), + }); + table["tools_without_strict"].clone() + } else { + table["tools"].clone() + }; + RetriedRequest { + reasons, + body: expected, + } + }); + config.retry_after_http_error( + status, + case["error_text"].as_str().unwrap(), + &body, + &MessagesTransformContext { + drops_unsupported_params: case["drops_unsupported_params"].as_bool().unwrap(), + ..MessagesTransformContext::default() + }, + ) != expected + }) + .map(|case| case["name"].as_str().unwrap()) + .collect(); + assert_eq!(mismatched, Vec::<&str>::new()); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index e79562b8762..a95229c5bf4 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -2,6 +2,6 @@ This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition +`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission. A route's protocol carries cache operations through any host call type that converts from `CacheCall` `python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index fe9e06b3d23..959c211b1d3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -2,7 +2,7 @@ mod python; mod selection; pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled}; +pub(crate) use selection::{Selection, configure, ensure_cache_disabled}; use pyo3::{exceptions::PyRuntimeError, prelude::*}; diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 3b516251f43..9aba11a3376 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -1,6 +1,6 @@ # Python cache delegation -This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` calls the Python cache's sync or async API and delivers the result back to Rust. +This directory lets Rust inference use a selected Python cache. `service.rs` implements the injected Rust response-cache service and yields typed cache operations. `host.rs` delivers each operation's result back to Rust. Asynchronous lookups and stream stores go through `litellm.rust_bridge.host.cache`, which reads with the `LLMCachingHandler` the `@client` wrapper uses, for the legacy function the route names, and writes with its dual cache. `configure` disables the cache for synchronous calls. Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index a8cbd0ce607..208f40a324b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -1,6 +1,6 @@ use litellm_cache::Error; use litellm_host::protocol::Reply; -use litellm_host_python::{from_py, to_py}; +use litellm_host_python::{json_from_py, to_py}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -18,7 +18,9 @@ enum Pending { pub(crate) struct PythonCache { cache: Option>, + function: Option>, arguments: Option>, + read_handler: Option>, pending: Option, asynchronous: bool, deferred: bool, @@ -28,19 +30,27 @@ impl PythonCache { pub fn new(asynchronous: bool) -> Self { Self { cache: None, + function: None, arguments: None, + read_handler: None, pending: None, asynchronous, deferred: false, } } + pub(in crate::cache) fn asynchronous(&self) -> bool { + self.asynchronous + } + pub(in crate::cache) fn bind( &mut self, cache: Bound<'_, PyAny>, + function: &Bound<'_, PyAny>, arguments: &Bound<'_, PyDict>, ) { self.cache = Some(cache.unbind()); + self.function = Some(function.clone().unbind()); self.arguments = Some(arguments.clone().unbind()); } @@ -59,7 +69,7 @@ impl PythonCache { self.pending = Some(Pending::Lookup(reply)); ( if self.asynchronous { - "async_get_cache" + "lookup" } else { "get_cache" }, @@ -77,7 +87,7 @@ impl PythonCache { self.pending = Some(Pending::Store(reply)); ( if self.asynchronous { - "async_add_cache" + "store_stream" } else { "add_cache" }, @@ -86,9 +96,20 @@ impl PythonCache { } }; let result = match result { + Some(value) if self.asynchronous => py + .import(CACHE_HOST_MODULE)? + .getattr(method)? + .call((cache, value), Some(&arguments)), Some(value) => cache .bind(py) .call_method(method, (value,), Some(&arguments)), + None if self.asynchronous => py.import(CACHE_HOST_MODULE)?.getattr(method)?.call( + ( + cache, + self.function.as_ref().ok_or(HandleError::Unconfigured)?, + ), + Some(&arguments), + ), None => cache.bind(py).call_method(method, (), Some(&arguments)), } .map(Bound::unbind); @@ -98,7 +119,12 @@ impl PythonCache { self.resume(py, result) } - pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { + pub fn store_response( + &mut self, + py: Python<'_>, + response: &Bound<'_, PyAny>, + verbatim: bool, + ) -> PyResult<()> { if !std::mem::take(&mut self.deferred) { return Ok(()); } @@ -106,11 +132,30 @@ impl PythonCache { return Err(HandleError::Unconfigured.into()); }; py.import(CACHE_HOST_MODULE)? - .getattr("store_response")? + .getattr(if verbatim { + "store_verbatim" + } else { + "store_response" + })? .call((cache, response), Some(&arguments.bind(py).copy()?))?; Ok(()) } + pub fn wrap_stream(&self, py: Python<'_>, stream: Py) -> PyResult> { + let (Some(cache), Some(function), Some(arguments)) = + (&self.cache, &self.function, &self.arguments) + else { + return Ok(stream); + }; + py.import(CACHE_HOST_MODULE)? + .getattr("wrap_stream")? + .call( + (cache, function, stream, &self.read_handler), + Some(&arguments.bind(py).copy()?), + ) + .map(Bound::unbind) + } + pub fn resume( &mut self, py: Python<'_>, @@ -125,11 +170,14 @@ impl PythonCache { }; match self.pending.take() { Some(Pending::Lookup(reply)) => { + if self.asynchronous { + self.read_handler = self.caching_handler(py).ok().flatten(); + } let value = result.map_err(Error::unavailable).and_then(|value| { if value.bind(py).is_none() { Ok(None) } else { - from_py(value.bind(py)) + json_from_py(value.bind(py)) .map(Some) .map_err(Error::invalid_entry) } @@ -144,15 +192,30 @@ impl PythonCache { Ok(None) } + fn caching_handler(&self, py: Python<'_>) -> PyResult>> { + let Some(arguments) = self.arguments.as_ref() else { + return Ok(None); + }; + let handler = py + .import(CACHE_HOST_MODULE)? + .getattr("caching_handler")? + .call((), Some(arguments.bind(py)))?; + Ok((!handler.is_none()).then(|| handler.unbind())) + } + pub fn close(&mut self) { self.pending = None; + self.read_handler = None; self.deferred = false; self.cache = None; + self.function = None; self.arguments = None; } pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { visit.call(&self.cache)?; - visit.call(&self.arguments) + visit.call(&self.function)?; + visit.call(&self.arguments)?; + visit.call(&self.read_handler) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index 018c8647043..bd54accc0dc 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -55,11 +55,12 @@ struct PythonCacheService { config: ResponseCacheConfig, } -pub(in crate::cache) fn service>( +pub(in crate::cache) fn service( services: HostServices

, namespace: String, ) -> std::sync::Arc where + P::HostCall: From, P::Error: From + std::error::Error, { std::sync::Arc::new(PythonCacheService { @@ -71,14 +72,19 @@ where }) } -impl> ResponseCacheService for PythonCacheService

+impl ResponseCacheService for PythonCacheService

where + P::HostCall: From, P::Error: From + std::error::Error, { fn config(&self) -> &ResponseCacheConfig { &self.config } + fn keyed_by_caller(&self) -> bool { + true + } + fn lookup<'a>( &'a self, _: &'a ResponseCacheRequest, @@ -86,7 +92,7 @@ where ) -> Pin, Error>> + Send + 'a>> { Box::pin(async move { self.services - .call(|reply| CacheCall::Lookup { reply }) + .call(|reply| CacheCall::Lookup { reply }.into()) .await .map_err(Error::unavailable)?? .map(from_python) @@ -104,10 +110,13 @@ where let value = to_python(value)?; let unary = value.get(STREAM_EVENTS_KEY).is_none(); self.services - .call(|reply| CacheCall::Store { - value, - unary, - reply, + .call(|reply| { + CacheCall::Store { + value, + unary, + reply, + } + .into() }) .await .map_err(Error::unavailable)? diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index 78c1606ea40..6e4d222a1d0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -6,17 +6,6 @@ use litellm_host::{ }; use pyo3::{prelude::*, types::PyDict}; -pub(crate) struct Cached

(std::marker::PhantomData

); - -impl Protocol for Cached

{ - type Request = (P::Request, Selection); - type Response = P::Response; - type Error = P::Error; - type HostCall = python::CacheCall; - type Chunk = P::Chunk; - type StreamHead = P::StreamHead; -} - enum Backend { Disabled, Python { namespace: String }, @@ -28,11 +17,12 @@ pub(crate) struct Selection { } impl Selection { - pub(crate) fn attach>( + pub(crate) fn attach( self, services: HostServices

, ) -> (Option, CacheOptions) where + P::HostCall: From, P::Error: From + std::error::Error, { let service = match self.backend { @@ -51,15 +41,21 @@ fn selected_cache<'py>( kwargs: &Bound<'py, PyDict>, call_type: &str, ) -> PyResult>> { + if kwargs + .get_item("caching")? + .is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false))) + { + return Ok(None); + } + configured_cache(py, call_type) +} + +fn configured_cache<'py>(py: Python<'py>, call_type: &str) -> PyResult>> { let configured = py .import(CACHE_HOST_MODULE)? .getattr("configured_cache")? .call0()?; - if configured.is_none() - || kwargs - .get_item("caching")? - .is_some_and(|value| value.is(pyo3::types::PyBool::new(py, false))) - { + if configured.is_none() { return Ok(None); } let supported = configured.getattr("supported_call_types")?; @@ -90,9 +86,14 @@ pub(crate) fn configure( python: &mut python::PythonCache, py: Python<'_>, arguments: &Bound<'_, PyDict>, - call_type: &str, + function: &Bound<'_, PyAny>, ) -> PyResult { - let selected = selected_cache(py, arguments, call_type)?; + let call_type: String = function.getattr("__name__")?.extract()?; + let selected = if python.asynchronous() { + configured_cache(py, &call_type)? + } else { + None + }; let Some(cache) = selected else { return Ok(Selection { backend: Backend::Disabled, @@ -102,6 +103,9 @@ pub(crate) fn configure( let enabled = cache .call_method("should_use_cache", (), Some(arguments))? .extract::()?; + let skips_read = arguments + .get_item("caching")? + .is_some_and(|value| !value.is_none() && !value.is(pyo3::types::PyBool::new(py, true))); let controls = arguments .get_item("cache")? .filter(|value| !value.is_none()); @@ -117,7 +121,7 @@ pub(crate) fn configure( }; let options = CacheOptions { policy: CachePolicy { - no_cache: boolean("no-cache")?, + no_cache: boolean("no-cache")? || skips_read, no_store: boolean("no-store")?, ..CachePolicy::default() }, @@ -129,7 +133,7 @@ pub(crate) fn configure( .map(|value| value.extract()) .transpose()? .unwrap_or_default(); - python.bind(cache, arguments); + python.bind(cache, function, arguments); Ok(Selection { backend: if enabled { Backend::Python { namespace } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index db5696d52ce..29c26ed30a3 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -3,7 +3,7 @@ use std::error::Error as StdError; use litellm_error::Report; use litellm_inference::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, }; use reqwest::header::HeaderMap; @@ -36,6 +36,37 @@ const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; /// Python refuses the request only once its logger knows the resolved model and provider (the /// provider config's request transform), so the failure logs them. const MESSAGES_LATE_REJECTION: &str = "messages_late_rejection"; +/// Python's request transform raises a builtin exception (`TypeError`) that its handler does not +/// map, so the caller gets the native exception as it is. +const MESSAGES_PYTHON_RAISED: &str = "messages_python_raised"; +/// The text of a refusal Python's provider config raises from its request transform, which its +/// HTTP handler maps as a provider error. +const MESSAGES_PROVIDER_REJECTION: &str = "messages_provider_rejection"; + +fn provider_rejection(rejection: &(dyn StdError + 'static)) -> Option { + match rejection.downcast_ref::() { + Some(litellm_llms::Error::MissingField("max_tokens")) => { + return Some("max_tokens is required for Anthropic /v1/messages API".to_owned()); + } + Some(_) => return None, + None => {} + } + match rejection.downcast_ref::() { + Some( + detail @ litellm_llms::ErrorDetail::InvalidChoice { + field: "reasoning effort", + .. + }, + ) => Some(format!("litellm.BadRequestError: {detail}")), + Some( + detail @ (litellm_llms::ErrorDetail::InvalidReasoningEffort { .. } + | litellm_llms::ErrorDetail::UnsupportedValue { + field: "effort", .. + }), + ) => Some(detail.to_string()), + _ => None, + } +} pub(crate) fn upstream_error( py: Python<'_>, @@ -90,6 +121,9 @@ pub(crate) fn route_failure( if let Some(raised) = crate::secrets::python_error(py, &error) { return Ok(raised); } + if let Some(raised) = crate::routes::messages::workload_identity::python_error(py, &error) { + return Ok(raised); + } if let Some((status, body)) = error.upstream_response() { let headers = error .upstream_headers() @@ -104,6 +138,18 @@ pub(crate) fn route_failure( return public(upstream_error(py, status, body, &headers)?); } let native = match (marker, error.rejected_request()) { + (RequestMarker::Messages, Some(rejection)) + if matches!( + rejection.downcast_ref::(), + Some(litellm_llms::ErrorDetail::UnorderableTypes { .. }) + ) => + { + let native = PyTypeError::new_err(rejection.to_string()); + let value = native.value(py); + value.setattr(MESSAGES_LATE_REJECTION, true)?; + value.setattr(MESSAGES_PYTHON_RAISED, true)?; + native + } (RequestMarker::Messages, Some(rejection)) => { let native = PyValueError::new_err(Report::new(rejection).to_string()); let value = native.value(py); @@ -123,6 +169,9 @@ pub(crate) fn route_failure( ) { value.setattr(MESSAGES_LATE_REJECTION, true)?; } + if let Some(text) = provider_rejection(rejection) { + value.setattr(MESSAGES_PROVIDER_REJECTION, text)?; + } native } (RequestMarker::Messages, None) => { @@ -324,6 +373,72 @@ mod tests { }); } + #[rstest::rstest] + #[case::missing_max_tokens( + litellm_llms::Error::MissingField("max_tokens").into(), + Some("max_tokens is required for Anthropic /v1/messages API") + )] + #[case::unmapped_effort( + Error::invalid_request(litellm_llms::ErrorDetail::InvalidChoice { + field: "reasoning effort", + actual: "'ultra'".into(), + choices: vec!["low", "none"], + }), + Some("litellm.BadRequestError: Unmapped reasoning effort: 'ultra'. Must be one of: 'low', 'none'.") + )] + #[case::other_missing_field(litellm_llms::Error::MissingField("model").into(), None)] + #[case::unsupported_param( + Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter { + model: "claude".into(), + param: "top_k".into(), + value: "5".into(), + hint: String::new(), + }), + None + )] + fn provider_config_rejections_carry_python_text( + #[case] error: Error, + #[case] expected: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap(); + let text = native + .value(py) + .getattr_opt(MESSAGES_PROVIDER_REJECTION) + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(text.as_deref(), expected); + }); + } + + #[test] + fn an_unorderable_comparison_raises_python_type_error_as_it_is() { + Python::initialize(); + Python::attach(|py| { + let error = Error::invalid_request(litellm_llms::ErrorDetail::UnorderableTypes { + op: "<=", + left: "str", + right: "int", + }); + let native = route_failure(py, error, RequestMarker::Messages, Ok).unwrap(); + assert!(native.is_instance_of::(py)); + let value = native.value(py); + assert_eq!( + value.str().unwrap().to_string(), + "'<=' not supported between instances of 'str' and 'int'" + ); + assert!(value.getattr_opt(MESSAGES_PYTHON_RAISED).unwrap().is_some()); + assert!( + value + .getattr_opt(MESSAGES_LATE_REJECTION) + .unwrap() + .is_some() + ); + assert!(value.getattr_opt(MESSAGES_REQUEST_ERROR).unwrap().is_none()); + }); + } + #[rstest::rstest] #[case::rejected_request( Error::invalid_request(litellm_llms::ErrorDetail::UnsupportedParameter { diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 57c58c3f7cd..e7cc8e35172 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -101,15 +101,22 @@ pub(crate) fn call_config( kwargs: &Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult { - let settings = HttpSettings::from_layers([ + Ok(resolved(py, call_settings(py, kwargs, asynchronous)?)) +} + +fn call_settings( + py: Python<'_>, + kwargs: &Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult { + Ok(HttpSettings::from_layers([ for_call(call_ssl_verify(kwargs)?, asynchronous), HttpSettingsLayer::from_environment(&ProcessEnvironment), match PythonSettings::Http.read_or_unset(py)? { Some(snapshot) => configured(&snapshot)?, None => HttpSettingsLayer::default(), }, - ]); - Ok(resolved(py, settings)) + ])) } pub(crate) fn settings_config( @@ -132,10 +139,31 @@ pub(crate) fn settings_config( Ok(resolved(py, settings)) } +#[pyfunction] +#[pyo3(signature = (ssl_security_level=None, ssl_ecdh_curve=None))] +pub(crate) fn unsupported_tls_settings( + ssl_security_level: Option, + ssl_ecdh_curve: Option, +) -> Vec { + let settings = HttpSettings::from_layers([HttpSettingsLayer { + ssl_security_level, + ssl_ecdh_curve, + ..HttpSettingsLayer::default() + }]); + Resolution::from(&settings) + .unsupported + .iter() + .map(ToString::to_string) + .collect() +} + +fn resolution(settings: HttpSettings) -> Resolution { + Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists())) +} + /// The client config `settings` resolve to, warning once per process about each setting the client cannot apply fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig { - let resolution = - Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists())); + let resolution = resolution(settings); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}")); } @@ -151,6 +179,15 @@ pub(crate) fn provider_client( Ok(pool().client(&config, ClientVariant::Provider)) } +pub(crate) fn messages_client( + py: Python<'_>, + kwargs: &Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let config = resolution(call_settings(py, kwargs, asynchronous)?).config; + Ok(pool().client(&config, ClientVariant::Undecoded)) +} + pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult { let config = call_config(py, &PyDict::new(py), true)?; pool().client(&config, variant).map_err(client_error) diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7edf13035c1..45ea95560e1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -35,6 +35,8 @@ mod _native { #[pymodule_export] use crate::errors::{RustBridgeDeclined, RustUpstreamError}; #[pymodule_export] + use crate::http::unsupported_tls_settings; + #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] use crate::model_cost::{ @@ -191,6 +193,7 @@ mod tests { "gil_stats", "process_state_started", "reserve_process_for_forking", + "unsupported_tls_settings", ]; #[cfg(feature = "huggingface")] expected.push("HuggingFaceEncoding"); diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index dec8ff0b0ce..927ed741669 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -4,7 +4,7 @@ use std::{ }; use litellm_auth::InputSource; -use litellm_host_python::{from_py, from_py_argument, to_py}; +use litellm_host_python::{from_py, from_py_argument, json_to_py, to_py}; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::{Map, Value}; @@ -143,6 +143,19 @@ pub(crate) fn public_response( .map(Bound::unbind) } +pub(crate) fn public_json_response( + py: Python<'_>, + module: &str, + response: &(impl Serialize + ?Sized), +) -> PyResult> { + let factory = py.import(module)?.getattr("response")?; + let document = + serde_json::to_value(response).map_err(|error| PyValueError::new_err(error.to_string()))?; + factory + .call1((json_to_py(py, &document)?,)) + .map(Bound::unbind) +} + struct RequestFieldSources<'py> { body: Option>, credentials: Option>, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/agentic.rs b/litellm-rust/crates/python-bridge/src/routes/messages/agentic.rs new file mode 100644 index 00000000000..f6652d82050 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/agentic.rs @@ -0,0 +1,172 @@ +use std::sync::{Arc, Mutex, PoisonError}; + +use litellm_host::{ + hooks::CallHooks, + interceptors::{RequestContext, WireRequest}, + lifecycle::Timing, +}; +use litellm_host_python::{HookStep, PythonOwned, PythonRuntime, json_to_py, to_py}; +use pyo3::{ + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::PyDict, +}; + +use super::host::ROUTE_HOST_MODULE; + +/// What the provider was sent, which Python's agentic hooks read back as the handler's own +/// model, provider and request params. +struct Sent { + request: Py, + hooked: bool, + body: Py, + model: String, + provider: String, + requested: Py, +} + +/// Python's Messages handler hands the provider's answer to the callbacks overriding +/// `async_should_run_agentic_loop` before `@client` sees it, and wraps a stream so they run +/// once it ends. This hook asks the route host to do both at the same points. +pub(super) type HookedRequest = Arc>>>; + +pub(super) struct AgenticLoop { + asynchronous: bool, + request: Option>, + hooked: HookedRequest, + decided: bool, + sent: Option, +} + +impl AgenticLoop { + pub(super) fn new(asynchronous: bool, hooked: HookedRequest) -> Self { + Self { + asynchronous, + request: None, + hooked, + decided: false, + sent: None, + } + } + + fn hooked_request(&self, py: Python<'_>) -> Option> { + self.hooked + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .map(|request| request.clone_ref(py)) + } + + fn host_call<'py>( + &self, + py: Python<'py>, + name: &str, + subject: &Bound<'py, PyAny>, + sent: &Sent, + ) -> PyResult> { + py.import(ROUTE_HOST_MODULE)?.getattr(name)?.call1(( + subject, + sent.request.bind(py), + sent.body.bind(py), + &sent.model, + &sent.provider, + sent.requested.bind(py), + sent.hooked, + )) + } + + fn resume_response( + &mut self, + _: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + result.map(HookStep::Ready) + } +} + +impl CallHooks for AgenticLoop { + fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + if self.asynchronous { + self.request = Some(arguments.clone_ref(py)); + } + Ok(()) + } + + fn before_provider_request( + &mut self, + py: Python<'_>, + wire: Box, + context: &RequestContext, + ) -> PyResult>> { + if self.decided { + return Ok(HookStep::Ready(wire)); + } + let Some(caller) = &self.request else { + return Ok(HookStep::Ready(wire)); + }; + self.decided = true; + let hooked = self.hooked_request(py); + let hooked_flag = hooked.is_some(); + let request = hooked.unwrap_or_else(|| caller.clone_ref(py)); + let module = py.import(ROUTE_HOST_MODULE)?; + if module + .getattr("runs_agentic_hooks")? + .call1((request.bind(py),))? + .extract::()? + { + module + .getattr("record_agentic_loop")? + .call1((request.bind(py), &context.custom_llm_provider))?; + self.sent = Some(Sent { + request, + hooked: hooked_flag, + body: to_py(py, &wire.body)?, + model: context.model.clone(), + provider: context.custom_llm_provider.clone(), + requested: json_to_py(py, &context.optional_params)?.unbind(), + }); + } + Ok(HookStep::Ready(wire)) + } + + fn transform_response( + &mut self, + py: Python<'_>, + response: Py, + _: Timing, + ) -> PyResult>> { + let Some(sent) = &self.sent else { + return Ok(HookStep::Ready(response)); + }; + let awaitable = self.host_call(py, "agentic_response", response.bind(py), sent)?; + Ok(HookStep::Await(awaitable.unbind(), Self::resume_response)) + } + + fn on_stream_open(&mut self, py: Python<'_>, head: &Py) -> PyResult<()> { + if let Some(sent) = &self.sent { + self.host_call(py, "agentic_stream", head.bind(py), sent)?; + } + Ok(()) + } +} + +impl PythonOwned for AgenticLoop { + fn close(&mut self, _: Python<'_>) { + self.request = None; + self.sent = None; + self.hooked + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.request)?; + if let Some(sent) = &self.sent { + visit.call(&sent.request)?; + visit.call(&sent.body)?; + visit.call(&sent.requested)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/answer.rs b/litellm-rust/crates/python-bridge/src/routes/messages/answer.rs new file mode 100644 index 00000000000..80cf2002287 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/answer.rs @@ -0,0 +1,29 @@ +use futures_util::future::BoxFuture; +use litellm_host::machine::HostServices; +use litellm_inference_messages::{Error, HostAnswer}; + +use super::host::{MessagesHostCall, MessagesHosted}; + +pub(super) struct PythonAnswer(HostServices); + +impl PythonAnswer { + pub(super) fn new(services: HostServices) -> Self { + Self(services) + } +} + +impl HostAnswer for PythonAnswer { + fn answer(&self, stream: bool) -> BoxFuture<'_, Result<(), Error>> { + Box::pin( + self.0 + .call(move |reply| MessagesHostCall::Answer { stream, reply }), + ) + } + + fn intercept(&self, stream: bool) -> BoxFuture<'_, Result> { + Box::pin( + self.0 + .call(move |reply| MessagesHostCall::Intercept { stream, reply }), + ) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index c975ac339c2..67a96cac32d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -1,34 +1,121 @@ -use crate::cache::{CacheCall, Cached, PythonCache, Selection}; +use crate::cache::{CacheCall, PythonCache, Selection}; use litellm_host_python::{PythonHostCalls, PythonOwned}; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + use bytes::Bytes; -use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; +use litellm_host::protocol::Reply; +use litellm_host_python::{ + InvokeError, PythonBinding, effective_py_args, from_py, json_from_py, to_py, +}; use litellm_inference_messages::{ - Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, + DecodingError, Error, MessagesCall, MessagesSettings, MessagesShaping, UnencodableHeader, + litellm_params, messages_body, route::{Messages, MessagesStreamHead}, }; use litellm_llms::base_llm::messages::context::MessagesModelCapabilities; use litellm_llms_types::headers::ProviderSpecificHeaders; use litellm_router_types::LitellmParams; use pyo3::{ - exceptions::PyException, + exceptions::{PyException, PyUnicodeEncodeError}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::{PyBytes, PyDict}, + sync::PyOnceLock, + types::{PyBytes, PyDict, PyList, PyMapping, PyString, PyTuple}, }; use serde_json::{Map, Value}; +use super::agentic::HookedRequest; +use super::safeguards::{SafeguardsCall, SafeguardsHost}; +use super::transport::{HostedTransport, TransportCall, python_failure}; use crate::{ errors::{RequestMarker, route_failure}, marshal::{ - ArgumentError, optional_timeout, project_optional_fields, public_response, + ArgumentError, optional_timeout, project_optional_fields, public_json_response, python_timeout_seconds, }, python_settings::PythonSettings, routes::RequestView, }; -const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; +type MergedHeaders = (Option>, Option>); + +fn ascii_headers<'py>( + headers: &Bound<'py, PyAny>, + unencodable: &mut Option>, +) -> PyResult> { + let Ok(headers) = headers.cast::() else { + return Ok(headers.clone()); + }; + let py = headers.py(); + let kept = PyDict::new(py); + for (name, value) in headers.iter() { + if value.is_instance_of::() + && let Err(error) = value.call_method1("encode", ("ascii",)) + { + if !error.is_instance_of::(py) { + return Err(error); + } + if unencodable.is_none() { + let sent = value.call_method1("encode", ("utf-8", "surrogatepass"))?; + *unencodable = Some(sent.cast::()?.as_bytes().to_vec()); + } + continue; + } + kept.set_item(name, value)?; + } + Ok(kept.into_any()) +} + +const PROVIDER_SPECIFIC_HEADERS_MODULE: &str = + "litellm.litellm_core_utils.get_provider_specific_headers"; + +fn ascii_scoped_headers<'py>(value: &Bound<'py, PyAny>) -> PyResult> { + let ascii_entry = |entry: Bound<'py, PyAny>| -> PyResult> { + let Ok(entry) = entry.cast::() else { + return Ok(entry); + }; + let Some(headers) = entry.get_item("extra_headers")? else { + return Ok(entry.clone().into_any()); + }; + let kept = entry.copy()?; + kept.set_item("extra_headers", ascii_headers(&headers, &mut None)?)?; + Ok(kept.into_any()) + }; + if let Ok(entries) = value.cast::() { + return Ok(PyList::new( + value.py(), + entries + .iter() + .map(ascii_entry) + .collect::>>()?, + )? + .into_any()); + } + if let Ok(entries) = value.cast::() { + return Ok(PyTuple::new( + value.py(), + entries + .iter() + .map(ascii_entry) + .collect::>>()?, + )? + .into_any()); + } + ascii_entry(value.clone()) +} + +fn json_messages(messages: &Bound<'_, PyAny>) -> PyResult { + if messages.is_instance_of::() || messages.is_instance_of::() { + return json_from_py(messages); + } + json_from_py(&messages.py().get_type::().call1((messages,))?) +} + +pub(super) const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; const MODEL_CAPABILITIES_MODULE: &str = "litellm.rust_bridge.host.model_capabilities"; fn merge_headers( @@ -41,6 +128,19 @@ fn merge_headers( (!merged.is_empty()).then_some(merged) } +const FEDERATING_PROVIDER: &str = "anthropic"; + +fn workload_identity_fields(py: Python<'_>) -> PyResult<&'static [String]> { + static FIELDS: PyOnceLock> = PyOnceLock::new(); + FIELDS + .get_or_try_init(py, || { + py.import(ROUTE_HOST_MODULE)? + .getattr("WORKLOAD_IDENTITY_FIELDS")? + .extract() + }) + .map(Vec::as_slice) +} + fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult>> { PythonSettings::ProviderDefaults .read_or_unset(py)? @@ -50,11 +150,17 @@ fn module_global<'py>(py: Python<'py>, name: &str) -> PyResult( +fn project_litellm_params<'a, 'py>( + names: impl IntoIterator, argument: impl Fn(&str) -> PyResult>>, global: impl Fn(&str) -> PyResult>>, ) -> PyResult> { - let fields = project_optional_fields(LitellmParams::fields(), &argument)?; + let fields = project_optional_fields( + LitellmParams::fields() + .map(|name| -> &'a str { name }) + .chain(names), + &argument, + )?; let globals = LitellmParams::specs() .filter(|spec| !spec.wire.iter().any(|name| fields.contains_key(*name))) .filter_map(|spec| spec.module_global); @@ -62,36 +168,279 @@ fn project_litellm_params<'py>( Ok(litellm_params(fields.into_iter().chain(folded).collect())) } +pub(super) struct MessagesHosted; + +impl litellm_host::protocol::Protocol for MessagesHosted { + type Request = (MessagesCall, Selection, bool); + type Response = ::Response; + type Error = Error; + type HostCall = MessagesHostCall; + type Chunk = Bytes; + type StreamHead = MessagesStreamHead; +} + +pub(super) enum MessagesHostCall { + Answer { stream: bool, reply: Reply<()> }, + Intercept { stream: bool, reply: Reply }, + Cache(CacheCall), + PreRequest(Reply>), + Safeguards(SafeguardsCall), + Transport(TransportCall), +} + +impl From for MessagesHostCall { + fn from(call: CacheCall) -> Self { + Self::Cache(call) + } +} + +impl From for MessagesHostCall { + fn from(call: TransportCall) -> Self { + Self::Transport(call) + } +} + +enum Answered { + Stream, + Whole { verbatim: bool }, +} + +#[derive(serde::Deserialize)] +struct HostSettings { + #[serde(flatten)] + settings: MessagesSettings, + model: String, + stream_cost_model: Option, + extra_headers: Option>, +} + /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesPythonHost { view: RequestView, cache: PythonCache, + transport: HostedTransport, + safeguards: SafeguardsHost, + classifies: Arc, + awaiting_safeguards: bool, + stream_cost_model: Option, + asynchronous: bool, + pre_request: Option>>, + hooked: bool, + hooked_request: HookedRequest, + requested_stream: bool, + arguments: Option>, + interceptor: Option>, + pending_answer: Option>, + pending_intercept: Option>, + pending_short_circuit: Option<(bool, Reply)>, + answer: Option>, + answered: Option, } impl MessagesPythonHost { - pub(super) fn new(view: RequestView, asynchronous: bool) -> Self { + pub(super) fn new( + view: RequestView, + asynchronous: bool, + hooked_request: HookedRequest, + classifies: Arc, + ) -> Self { Self { view, cache: PythonCache::new(asynchronous), + transport: HostedTransport::default(), + safeguards: SafeguardsHost::default(), + classifies, + awaiting_safeguards: false, + stream_cost_model: None, + asynchronous, + pre_request: None, + hooked: false, + hooked_request, + requested_stream: false, + arguments: None, + interceptor: None, + pending_answer: None, + pending_intercept: None, + pending_short_circuit: None, + answer: None, + answered: None, } } + fn interceptor(&self, py: Python<'_>) -> PyResult>> { + if !self.asynchronous { + return Ok(None); + } + let interceptor = py + .import(ROUTE_HOST_MODULE)? + .getattr("interceptor")? + .call1((self.view.request(py)?,))?; + Ok((!interceptor.is_none()).then(|| interceptor.unbind())) + } + + fn may_intercept(&self, py: Python<'_>) -> PyResult { + if !self.asynchronous { + return Ok(false); + } + let module = py.import(ROUTE_HOST_MODULE)?; + Ok(module.getattr("has_interceptors")?.call0()?.extract()? + || module.getattr("short_circuits")?.call0()?.extract()?) + } + + fn begin_intercept( + &mut self, + py: Python<'_>, + stream: bool, + reply: Reply, + ) -> PyResult>> { + let module = py.import(ROUTE_HOST_MODULE)?; + if !module + .getattr("short_circuits")? + .call0()? + .extract::()? + { + return self.begin_interceptor(py, stream, reply); + } + let pending = module.getattr("short_circuit")?.call1(( + self.view.request(py)?, + stream, + self.hooked, + ))?; + self.pending_short_circuit = Some((stream, reply)); + Ok(Some(pending.unbind())) + } + + fn resume_short_circuit( + &mut self, + py: Python<'_>, + stream: bool, + reply: Reply, + answer: Py, + ) -> PyResult>> { + if answer.is_none(py) { + return self.begin_interceptor(py, stream, reply); + } + self.answer = Some(answer); + self.answered = Some(match stream { + true => Answered::Stream, + false => Answered::Whole { verbatim: false }, + }); + reply.send(true); + Ok(None) + } + + fn begin_interceptor( + &mut self, + py: Python<'_>, + stream: bool, + reply: Reply, + ) -> PyResult>> { + self.interceptor = self.interceptor(py)?; + if self.interceptor.is_none() { + reply.send(false); + return Ok(None); + } + let pending = self.start_answer(py, stream)?; + self.pending_intercept = Some(reply); + Ok(Some(pending)) + } + + fn begin_answer( + &mut self, + py: Python<'_>, + stream: bool, + reply: Reply<()>, + ) -> PyResult>> { + let pending = self.start_answer(py, stream)?; + self.pending_answer = Some(reply); + Ok(Some(pending)) + } + + fn start_answer(&mut self, py: Python<'_>, stream: bool) -> PyResult> { + let module = py.import(ROUTE_HOST_MODULE)?; + let request = self.view.request(py)?; + let pending = match &self.interceptor { + Some(interceptor) => { + module + .getattr("answer")? + .call1((interceptor, request, stream, self.hooked))? + } + None => module + .getattr("gateway")? + .call1((request, self.asynchronous, self.hooked))?, + }; + if pending.is_none() { + return Err(litellm_host_python::Error::MissingState.into()); + } + self.answered = Some(match (stream, &self.interceptor) { + (true, _) => Answered::Stream, + (false, Some(_)) => Answered::Whole { verbatim: false }, + (false, None) => Answered::Whole { verbatim: true }, + }); + Ok(pending.unbind()) + } + + fn host_settings( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult { + from_py(&py.import(ROUTE_HOST_MODULE)?.getattr("settings")?.call1(( + arguments, + self.view.request(py)?, + self.asynchronous, + ))?) + } + + fn body_params( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult> { + let request = effective_py_args(self.view.request(py)?, arguments)?; + let params = py + .import(ROUTE_HOST_MODULE)? + .getattr("body_params")? + .call1((request, self.asynchronous, self.hooked))? + .cast_into::()?; + params + .items()? + .iter() + .map(|item| { + let (name, value): (Bound<'_, PyAny>, Bound<'_, PyAny>) = item.extract()?; + let name: String = name.extract()?; + let value = match name.as_str() { + "messages" => json_messages(&value)?, + _ => json_from_py(&value)?, + }; + Ok((name, value)) + }) + .collect() + } + fn projection( &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, + settings: HostSettings, ) -> PyResult> { let argument = |name: &str| self.view.argument(py, arguments, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; - let model = string("model")?.ok_or(ArgumentError::Required("model"))?; + let model = settings.model; let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; - let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; + let mut fields = self.body_params(py, arguments)?; + let prepared = fields.remove("messages"); + let presanitized = prepared.is_some(); + let messages = match prepared { + Some(messages) => messages, + None => json_messages(&messages)?, + }; let body = [ ("model".to_string(), Value::String(model.clone())), - ("messages".to_string(), from_py(&messages)?), + ("messages".to_string(), messages), ] .into_iter() .chain(fields) @@ -112,10 +461,35 @@ impl MessagesPythonHost { .call1((&model, api_base.as_deref()))? .extract()?, }; - let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; - let extra_headers = self.merged_headers(py, arguments)?; - let provider_specific_header = self.provider_specific_header(py, arguments)?; - let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; + let (extra_headers, mut unencodable_request_header) = match settings.extra_headers { + Some(headers) => (Some(headers), None), + None => self.merged_headers(py, arguments)?, + }; + let provider_specific_header = self.provider_specific_header( + py, + arguments, + custom_llm_provider.as_deref(), + &mut unencodable_request_header, + )?; + let shaping = MessagesShaping { + capabilities: self.capabilities(py, &model, custom_llm_provider.as_deref())?, + settings: MessagesSettings { + whole_sse_frames: settings.stream_cost_model.is_some(), + presanitized, + intercepted: self.may_intercept(py)?, + unencodable_request_header, + ..settings.settings + }, + }; + let federation_fields: &[String] = match custom_llm_provider.as_deref() { + Some(FEDERATING_PROVIDER) => workload_identity_fields(py)?, + _ => &[], + }; + let litellm_params = project_litellm_params( + federation_fields.iter().map(String::as_str), + argument, + |name| module_global(py, name), + )?; Ok(messages_body(body).and_then(|body| { Ok(MessagesCall { body, @@ -135,49 +509,55 @@ impl MessagesPythonHost { &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - ) -> PyResult>> { - let mapping = |name: &str| -> PyResult>> { + ) -> PyResult { + let mut unencodable = None; + let mut mapping = |name: &str| -> PyResult>> { self.view .argument(py, arguments, name)? - .map(|value| from_py(&value)) + .map(|value| from_py(&ascii_headers(&value, &mut unencodable)?)) .transpose() }; - Ok(merge_headers( - mapping("headers")?, - mapping("extra_headers")?, - )) + let headers = mapping("headers")?; + let extra_headers = mapping("extra_headers")?; + Ok((merge_headers(headers, extra_headers), unencodable)) } fn provider_specific_header( &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, + provider: Option<&str>, + unencodable: &mut Option>, ) -> PyResult> { - self.view + let Some(value) = self + .view .argument(py, arguments, "provider_specific_header")? - .map(|value| from_py(&value)) - .transpose() + else { + return Ok(None); + }; + if unencodable.is_none() + && let Some(provider) = provider + && let Ok(scoped) = py + .import(PROVIDER_SPECIFIC_HEADERS_MODULE)? + .getattr("ProviderSpecificHeaderUtils")? + .call_method1("get_provider_specific_headers", (&value, provider, true)) + { + ascii_headers(&scoped, unencodable)?; + } + from_py(&ascii_scoped_headers(&value)?).map(Some) } - fn shaping( + fn capabilities( &self, py: Python<'_>, model: &str, custom_llm_provider: Option<&str>, - arguments: &Bound<'_, PyDict>, - ) -> PyResult { - let module = py.import(ROUTE_HOST_MODULE)?; - let capabilities: MessagesModelCapabilities = from_py( + ) -> PyResult { + from_py( &py.import(MODEL_CAPABILITIES_MODULE)? .getattr("anthropic_model_capabilities")? .call1((model, custom_llm_provider))?, - )?; - let settings: MessagesSettings = - from_py(&module.getattr("settings")?.call1((arguments,))?)?; - Ok(MessagesShaping { - capabilities, - settings, - }) + ) } fn provider(&self, py: Python<'_>) -> String { @@ -190,6 +570,14 @@ impl MessagesPythonHost { .unwrap_or_else(|| "anthropic".into()) } + fn map_transport_failure(&self, py: Python<'_>, error: PyErr) -> PyResult { + let mapped = py + .import(ROUTE_HOST_MODULE)? + .getattr("map_transport_failure")? + .call1((error.value(py), self.view.request(py)?, self.provider(py)))?; + Ok(PyErr::from_value(mapped)) + } + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { if !error.is_instance_of::(py) { return error; @@ -197,7 +585,14 @@ impl MessagesPythonHost { let mapped = py .import(ROUTE_HOST_MODULE) .and_then(|module| module.getattr("map_failure")) - .and_then(|map| map.call1((error.value(py), self.view.request(py)?, self.provider(py)))) + .and_then(|map| { + map.call1(( + error.value(py), + self.view.request(py)?, + self.provider(py), + self.asynchronous, + )) + }) .and_then(|mapped| { mapped .extract::>() @@ -210,25 +605,109 @@ impl MessagesPythonHost { } } +impl MessagesPythonHost { + fn decode_call( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> Result> { + let settings = self + .host_settings(py, arguments) + .map_err(InvokeError::Python)?; + self.stream_cost_model = settings.stream_cost_model.clone(); + self.projection(py, arguments, settings) + .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? + .map_err(InvokeError::Native) + } + + fn begin_pre_request( + &mut self, + py: Python<'_>, + reply: Reply>, + ) -> PyResult> { + let hooks = py + .import(ROUTE_HOST_MODULE)? + .getattr("pre_request_hooks")? + .call1((self.view.request(py)?,))?; + self.pre_request = Some(reply); + Ok(hooks.unbind()) + } + + fn resume_pre_request( + &mut self, + py: Python<'_>, + reply: Reply>, + result: PyResult>, + ) -> Result<(), InvokeError> { + let hooked = result?; + let arguments = hooked + .bind(py) + .cast::() + .map_err(|error| InvokeError::Python(error.into()))?; + self.view + .resolve(py, arguments) + .map_err(InvokeError::Python)?; + self.hooked = true; + let effective = self + .view + .request(py) + .and_then(|request| effective_py_args(request, arguments)) + .map_err(InvokeError::Python)?; + self.safeguards + .rebind(py, &effective) + .map_err(InvokeError::Python)?; + *self + .hooked_request + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(arguments.clone().unbind()); + match self.decode_call(py, arguments) { + Ok(call) => reply.send(Ok(call)), + Err(InvokeError::Native(error)) => reply.send(Err(error)), + Err(error) => return Err(error), + } + Ok(()) + } +} + impl PythonBinding for MessagesPythonHost { - type Protocol = Cached; + type Protocol = MessagesHosted; type Failure = PyErr; fn decode_request( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, - ) -> Result<(MessagesCall, Selection), InvokeError> { + ) -> Result<(MessagesCall, Selection, bool), InvokeError> { self.view .resolve(py, arguments) .map_err(InvokeError::Python)?; - let selection = - crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages") - .map_err(InvokeError::Python)?; - self.projection(py, arguments) - .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? - .map_err(InvokeError::Native) - .map(|request| (request, selection)) + let function = py + .import(ROUTE_HOST_MODULE) + .and_then(|module| module.call_method0("cached_function")) + .map_err(InvokeError::Python)?; + let selection = crate::cache::configure(&mut self.cache, py, arguments, &function) + .map_err(InvokeError::Python)?; + let effective = self + .view + .request(py) + .and_then(|request| effective_py_args(request, arguments)) + .map_err(InvokeError::Python)?; + let classifies = self + .safeguards + .configure(py, &effective) + .map_err(InvokeError::Python)?; + self.classifies.store(classifies, Ordering::Release); + self.requested_stream = arguments + .get_item("stream") + .map_err(InvokeError::Python)? + .is_some_and(|stream| stream.is(pyo3::types::PyBool::new(py, true))); + self.arguments = Some(arguments.clone().unbind()); + let request = self.decode_call(py, arguments)?; + let hosted = self + .transport + .select(py, self.view.request(py).map_err(InvokeError::Python)?) + .map_err(InvokeError::Python)?; + Ok((request, selection, hosted)) } fn encode_response( @@ -236,11 +715,23 @@ impl PythonBinding for MessagesPythonHost { py: Python<'_>, response: litellm_inference_messages::MessagesBody, ) -> PyResult> { - public_response(py, ROUTE_HOST_MODULE, &response) + if let Some(answer) = self.answer.take() { + return Ok(answer); + } + let public = public_json_response(py, ROUTE_HOST_MODULE, &response)?; + self.safeguards.published(py, public) } - fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> { - self.cache.store_response(py, response) + fn deliver_response(&mut self, py: Python<'_>, response: Py) -> PyResult> { + match self.answered.take() { + Some(Answered::Stream) => return self.cache.wrap_stream(py, response), + Some(Answered::Whole { verbatim }) => { + self.cache.store_response(py, response.bind(py), verbatim)?; + } + None if self.requested_stream => {} + None => self.cache.store_response(py, response.bind(py), false)?, + } + Ok(response) } fn encode_stream_head( @@ -248,17 +739,63 @@ impl PythonBinding for MessagesPythonHost { py: Python<'_>, head: MessagesStreamHead, ) -> PyResult> { - py.import(ROUTE_HOST_MODULE)? + let module = py.import(ROUTE_HOST_MODULE)?; + let head = module .getattr("stream_hidden_params")? - .call1((to_py(py, &head.headers)?,)) - .map(Bound::unbind) + .call1((to_py(py, &head.headers)?,))?; + if let Some(arguments) = self + .arguments + .as_ref() + .filter(|_| self.asynchronous && !self.requested_stream) + { + module + .getattr("unrequested_stream")? + .call1((&head, arguments))?; + } + Ok(head.unbind()) } fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { - Ok(PyBytes::new(py, &chunk).into_any().unbind()) + let frames = PyBytes::new(py, &chunk).into_any(); + let Some(model) = &self.stream_cost_model else { + return Ok(frames.unbind()); + }; + py.import(ROUTE_HOST_MODULE)? + .getattr("inject_stream_cost")? + .call1((frames, model, self.view.request(py)?)) + .map(Bound::unbind) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { + if let Some((raised, raw)) = python_failure(py, &error) { + return if raw { + Ok(raised) + } else { + self.map_transport_failure(py, raised) + }; + } + if let Some(unencodable) = litellm_error::find::(&error) { + let (value, received) = match unencodable { + UnencodableHeader::Sent(value) => (value, false), + UnencodableHeader::Received(value) => (value, true), + }; + let raised = py + .import(ROUTE_HOST_MODULE)? + .getattr("unencodable_header")? + .call1((PyBytes::new(py, value), received))?; + return self.map_transport_failure(py, PyErr::from_value(raised)); + } + if let Some(undecodable) = litellm_error::find::(&error) { + let raised = py + .import(ROUTE_HOST_MODULE)? + .getattr("decoding_failure")? + .call1(( + undecodable.to_string(), + self.view.request(py)?, + self.provider(py), + ))?; + return Ok(PyErr::from_value(raised)); + } route_failure(py, error, RequestMarker::Messages, |native| { Ok(self.map_failure(py, native)) }) @@ -269,24 +806,35 @@ impl PythonBinding for MessagesPythonHost { } } -impl PythonHostCalls> for MessagesPythonHost { +impl PythonHostCalls for MessagesPythonHost { fn handle_host_call( &mut self, py: Python<'_>, - op: CacheCall, + op: MessagesHostCall, ) -> Result<(), InvokeError> { - self.cache - .begin(py, op) - .map(|_| ()) - .map_err(InvokeError::Python) + self.begin_host_call(py, op).map(|_| ()) } fn begin_host_call( &mut self, py: Python<'_>, - op: CacheCall, + op: MessagesHostCall, ) -> Result>, InvokeError> { - self.cache.begin(py, op).map_err(InvokeError::Python) + self.awaiting_safeguards = matches!(op, MessagesHostCall::Safeguards(_)); + match op { + MessagesHostCall::Cache(op) => self.cache.begin(py, op), + MessagesHostCall::Transport(op) => { + let request = self.view.request(py).map_err(InvokeError::Python)?; + self.transport.begin(py, op, request) + } + MessagesHostCall::PreRequest(reply) => self.begin_pre_request(py, reply).map(Some), + MessagesHostCall::Safeguards(op) => self.safeguards.begin(py, op).map(Some), + MessagesHostCall::Answer { stream, reply } => self.begin_answer(py, stream, reply), + MessagesHostCall::Intercept { stream, reply } => { + self.begin_intercept(py, stream, reply) + } + } + .map_err(InvokeError::Python) } fn resume_host_call( @@ -294,6 +842,37 @@ impl PythonHostCalls> for MessagesPythonHost { py: Python<'_>, result: PyResult>, ) -> Result>, InvokeError> { + if let Some(reply) = self.pre_request.take() { + return self.resume_pre_request(py, reply, result).map(|()| None); + } + if let Some((stream, reply)) = self.pending_short_circuit.take() { + return self + .resume_short_circuit(py, stream, reply, result?) + .map_err(InvokeError::Python); + } + if let Some(reply) = self.pending_answer.take() { + self.answer = Some(result?); + reply.send(()); + return Ok(None); + } + if let Some(reply) = self.pending_intercept.take() { + self.answer = Some(result?); + reply.send(true); + return Ok(None); + } + if std::mem::take(&mut self.awaiting_safeguards) { + return self + .safeguards + .resume(py, result) + .map(|()| None) + .map_err(InvokeError::Python); + } + if self.transport.pending() { + return self + .transport + .resume(py, result) + .map_err(InvokeError::Python); + } self.cache.resume(py, result).map_err(InvokeError::Python) } } @@ -302,10 +881,25 @@ impl PythonOwned for MessagesPythonHost { fn close(&mut self, _: Python<'_>) { self.view.close(); self.cache.close(); + self.transport.close(); + self.pre_request = None; + self.safeguards.close(); + self.interceptor = None; + self.pending_answer = None; + self.pending_intercept = None; + self.pending_short_circuit = None; + self.answer = None; + self.answered = None; + self.arguments = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { self.view.traverse(visit)?; - self.cache.traverse(visit) + self.cache.traverse(visit)?; + self.transport.traverse(visit)?; + self.safeguards.traverse(visit)?; + visit.call(&self.interceptor)?; + visit.call(&self.arguments)?; + visit.call(&self.answer) } } @@ -416,7 +1010,40 @@ mod tests { assert_projection(kwargs, globals, expected); } + #[rstest] + #[case::named_fields_are_projected( + json!({"model": "m", "anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60}), + Ok(json!({"model": "m", "anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60})), + )] + #[case::unnamed_fields_are_not( + json!({"model": "m", "anthropic_unrelated": "x"}), + Ok(json!({"model": "m"})), + )] + fn workload_identity_fields_are_projected_by_the_names_the_host_lists( + #[case] kwargs: Value, + #[case] expected: Result, + ) { + assert_projection_of( + &[ + "anthropic_federation_rule_id", + "anthropic_issuer_ttl_seconds", + ], + kwargs, + json!({}), + expected, + ); + } + fn assert_projection(kwargs: Value, globals: Value, expected: Result) { + assert_projection_of(&[], kwargs, globals, expected); + } + + fn assert_projection_of( + names: &[&str], + kwargs: Value, + globals: Value, + expected: Result, + ) { Python::initialize(); Python::attach(|py| { let dict = |value: &Value| { @@ -430,6 +1057,7 @@ mod tests { let globals = dict(&globals); let bound = PyDict::new(py); let projected = project_litellm_params( + names.iter().copied(), |name| present(&kwargs, &bound, name), |name| present(&globals, &bound, name), ) diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index a9ae01bd7ee..9a1de064051 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,7 +1,15 @@ +mod agentic; +mod answer; mod host; +mod modify_params; +mod pre_request; +mod safeguards; +mod transport; +pub(crate) mod workload_identity; use host::MessagesPythonHost; use litellm_callbacks_legacy_python::LoggingOperation; +use modify_params::ModifyParams; use pyo3::prelude::*; use super::NativeCall; @@ -9,23 +17,61 @@ use super::NativeCall; fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { let (arguments, hooks) = crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?; - crate::routes::run_public_call( + let hooked_request = agentic::HookedRequest::default(); + let hooks = litellm_host_python::HookChain::new() + .with(agentic::AgenticLoop::new( + asynchronous, + std::sync::Arc::clone(&hooked_request), + )) + .with(hooks); + let modify_params = ModifyParams::new(&call.args, &call.kwargs)?; + let classifies = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let classified = std::sync::Arc::clone(&classifies); + crate::routes::run_public_call_after_policy( py, arguments, move |py, arguments, request| { + let hooked = pre_request::runs_pre_request_hooks(py)?; let route = litellm_inference_messages::MessagesRoute::new( - crate::http::provider_client(py, arguments, asynchronous)? + crate::http::messages_client(py, arguments, asynchronous)? .map_err(crate::http::client_error)?, crate::http::resources().auth.clone(), crate::secrets::source(py)?, - ); + ) + .with_workload_identity(std::sync::Arc::new( + workload_identity::PythonWorkloadIdentity::new( + litellm_host_python::PythonContext::capture(py)?, + ), + )); Ok(litellm_host::call::hosted_call( request, None, - move |(call, selection): (_, crate::cache::Selection), + move |(call, selection, hosted): (_, crate::cache::Selection, bool), services, interceptors, observers| async move { + let route = if hosted { + route.with_transport(std::sync::Arc::new(transport::PythonTransport::new( + services.clone(), + ))) + } else { + route + }; + let route = match hooked { + true => route.with_pre_request_hooks(std::sync::Arc::new( + pre_request::HostPreRequestHooks::new(services.clone()), + )), + false => route, + }; + let route = match classified.load(std::sync::atomic::Ordering::Acquire) { + true => route.with_safeguards(std::sync::Arc::new( + safeguards::PythonSafeguards::new(services.clone()), + )), + false => route, + }; + let route = route.with_host_answer(std::sync::Arc::new( + answer::PythonAnswer::new(services.clone()), + )); let (cache, options) = selection.attach(services); let route = match cache { Some(cache) => route.with_cache(cache), @@ -44,24 +90,32 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR }, )) }, - MessagesPythonHost::new(call.view()?, asynchronous), + MessagesPythonHost::new(call.view()?, asynchronous, hooked_request, classifies), hooks, + modify_params, asynchronous, ) } -/// Python's sync Messages handler raises for every native provider before it sends anything -/// or opens a logging object (`BaseLLMHTTPHandler.anthropic_messages_handler`), and Rust copies -/// that until both sides serve sync calls. +/// Python's sync Messages handler answers a mock without logging, hands proxy MCP tools to the +/// gateway, and raises for every other native provider call before it sends anything or opens a +/// logging object (`BaseLLMHTTPHandler.anthropic_messages_handler`), and Rust copies that until +/// both sides serve sync calls. const SYNC_NOT_IMPLEMENTED: &str = "anthropic_messages_handler is not implemented for sync calls"; #[pyfunction] #[pyo3(signature = (call))] -pub(crate) fn messages(call: NativeCall<'_>) -> PyResult> { - drop(call); - Err(pyo3::exceptions::PyValueError::new_err( - SYNC_NOT_IMPLEMENTED, - )) +pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + let answer = py + .import(host::ROUTE_HOST_MODULE)? + .getattr("sync_answer")? + .call1((call.resolved()?,))?; + if answer.is_none() { + return Err(pyo3::exceptions::PyValueError::new_err( + SYNC_NOT_IMPLEMENTED, + )); + } + Ok(answer.unbind()) } #[pyfunction] diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs b/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs new file mode 100644 index 00000000000..c6f846478ad --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/modify_params.rs @@ -0,0 +1,152 @@ +use litellm_host::hooks::CallHooks; +use litellm_host_python::{PythonOwned, PythonRuntime}; +use pyo3::{ + exceptions::PyException, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyDict, PyTuple}, +}; + +use super::host::ROUTE_HOST_MODULE; + +pub(super) struct ModifyParams { + call: Option<(Py, Py)>, +} + +impl ModifyParams { + pub(super) fn new(args: &Bound<'_, PyTuple>, kwargs: &Bound<'_, PyDict>) -> PyResult { + let model = match args.get_item(0) { + Ok(model) => model, + Err(_) => kwargs + .get_item("model")? + .unwrap_or_else(|| args.py().None().into_bound(args.py())), + }; + Ok(Self { + call: Some((model.unbind(), args.clone().unbind())), + }) + } +} + +impl CallHooks for ModifyParams { + fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py) -> PyResult<()> { + let Some((model, args)) = &self.call else { + return Ok(()); + }; + let clamped = py + .import(ROUTE_HOST_MODULE) + .and_then(|module| module.getattr("modify_max_tokens")) + .and_then(|clamp| clamp.call1((model, args, arguments))); + match clamped { + Err(error) if !error.is_instance_of::(py) => Err(error), + _ => Ok(()), + } + } +} + +impl PythonOwned for ModifyParams { + fn close(&mut self, _: Python<'_>) { + self.call = None; + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some((model, args)) = &self.call { + visit.call(model)?; + visit.call(args)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, + }; + use rstest::rstest; + + use super::{ModifyParams, ROUTE_HOST_MODULE}; + use litellm_host::hooks::CallHooks; + + #[rstest] + #[case::keyword_model( + "()", + "{'model': 'anthropic/claude-sonnet-4-5'}", + "'anthropic/claude-sonnet-4-5'" + )] + #[case::first_positional_wins("(64, [])", "{'model': 'anthropic/claude-sonnet-4-5'}", "64")] + #[case::no_model("()", "{}", "None")] + fn reads_the_model_the_client_wrapper_reads( + #[case] args: &str, + #[case] kwargs: &str, + #[case] expected: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let eval = |source: &str| { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + }; + let args = eval(args).cast_into::().unwrap(); + let kwargs = eval(kwargs).cast_into::().unwrap(); + let hook = ModifyParams::new(&args, &kwargs).unwrap(); + let (model, captured) = hook.call.as_ref().unwrap(); + assert!(model.bind(py).eq(eval(expected)).unwrap()); + assert!(captured.bind(py).is(&args)); + }); + } + + #[rstest] + #[case::exception_leaves_max_tokens("ValueError", true)] + #[case::base_exception_propagates("KeyboardInterrupt", false)] + fn a_failing_clamp_behaves_like_the_client_wrappers_swallowed_try( + #[case] raised: &str, + #[case] swallowed: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("name", ROUTE_HOST_MODULE).unwrap(); + locals.set_item("raised", raised).unwrap(); + py.run( + c" +import builtins, sys, types +for parent in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.routes'): + sys.modules.setdefault(parent, types.ModuleType(parent)) +saved = sys.modules.get(name) +stub = types.ModuleType(name) +def modify_max_tokens(model, args, kwargs): + kwargs['max_tokens'] = 1 + raise getattr(builtins, raised)() +stub.modify_max_tokens = modify_max_tokens +sys.modules[name] = stub +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let args = PyTuple::empty(py); + let kwargs = PyDict::new(py); + kwargs + .set_item("model", "anthropic/claude-sonnet-4-5") + .unwrap(); + let mut hook = ModifyParams::new(&args, &kwargs).unwrap(); + let result = hook.arguments_prepared(py, &kwargs.clone().unbind()); + py.run( + c" +if saved is None: + sys.modules.pop(name, None) +else: + sys.modules[name] = saved +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + assert_eq!(result.is_ok(), swallowed); + if let Err(error) = result { + assert!(error.is_instance_of::(py)); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/pre_request.rs b/litellm-rust/crates/python-bridge/src/routes/messages/pre_request.rs new file mode 100644 index 00000000000..0553fc0fa7f --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/pre_request.rs @@ -0,0 +1,35 @@ +use futures_util::future::BoxFuture; +use litellm_host::machine::HostServices; +use litellm_inference_messages::{Error, MessagesCall, PreRequestHooks}; +use pyo3::prelude::*; + +use super::host::{MessagesHostCall, MessagesHosted, ROUTE_HOST_MODULE}; + +pub(super) fn runs_pre_request_hooks(py: Python<'_>) -> PyResult { + py.import(ROUTE_HOST_MODULE)? + .getattr("runs_pre_request_hooks")? + .call0()? + .extract() +} + +pub(super) struct HostPreRequestHooks { + services: HostServices, +} + +impl HostPreRequestHooks { + pub(super) fn new(services: HostServices) -> Self { + Self { services } + } +} + +impl PreRequestHooks for HostPreRequestHooks { + fn run(&self, call: MessagesCall) -> BoxFuture<'_, Result> { + drop(call); + Box::pin(async move { + self.services + .call(MessagesHostCall::PreRequest) + .await + .and_then(|hooked| hooked) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/safeguards.rs b/litellm-rust/crates/python-bridge/src/routes/messages/safeguards.rs new file mode 100644 index 00000000000..af478521cb3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/safeguards.rs @@ -0,0 +1,183 @@ +use bytes::Bytes; +use futures_util::future::BoxFuture; +use litellm_host::{machine::HostServices, protocol::Reply}; +use litellm_host_python::{from_py, to_py}; +use litellm_inference_messages::{Error, MessagesBody, Safeguards}; +use pyo3::{ + exceptions::PyRuntimeError, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyBytes, PyDict}, +}; +use serde_json::Value; + +use super::host::{MessagesHostCall, MessagesHosted, ROUTE_HOST_MODULE}; + +pub(super) enum SafeguardsCall { + Results { + response: MessagesBody, + reply: Reply>, + }, + OpenStream { + reply: Reply, + }, + Frames { + chunk: Bytes, + reply: Reply>, + }, + Flush { + reply: Reply>, + }, +} + +pub(super) struct PythonSafeguards { + services: HostServices, +} + +impl PythonSafeguards { + pub(super) fn new(services: HostServices) -> Self { + Self { services } + } + + async fn ask( + &self, + call: impl FnOnce(Reply) -> SafeguardsCall + Send, + ) -> Result { + self.services + .call(|reply| MessagesHostCall::Safeguards(call(reply))) + .await + } +} + +impl Safeguards for PythonSafeguards { + fn results<'a>( + &'a self, + response: &'a MessagesBody, + ) -> BoxFuture<'a, Result, Error>> { + let response = response.clone(); + Box::pin(self.ask(|reply| SafeguardsCall::Results { response, reply })) + } + + fn open_stream(&self) -> BoxFuture<'_, Result> { + Box::pin(self.ask(|reply| SafeguardsCall::OpenStream { reply })) + } + + fn frames(&self, chunk: Bytes) -> BoxFuture<'_, Result, Error>> { + Box::pin(self.ask(|reply| SafeguardsCall::Frames { chunk, reply })) + } + + fn flush(&self) -> BoxFuture<'_, Result, Error>> { + Box::pin(self.ask(|reply| SafeguardsCall::Flush { reply })) + } +} + +enum Pending { + Results(Reply>), + OpenStream(Reply), + Frames(Reply>), +} + +#[derive(Default)] +pub(super) struct SafeguardsHost { + classifier: Option>, + pending: Option, +} + +impl SafeguardsHost { + pub(super) fn configure( + &mut self, + py: Python<'_>, + request: &Bound<'_, PyDict>, + ) -> PyResult { + let classifier = py + .import(ROUTE_HOST_MODULE)? + .getattr("safeguards")? + .call1((request,))?; + self.classifier = (!classifier.is_none()).then(|| classifier.unbind()); + Ok(self.classifier.is_some()) + } + + pub(super) fn rebind(&mut self, py: Python<'_>, request: &Bound<'_, PyDict>) -> PyResult<()> { + if self.classifier.is_none() { + return Ok(()); + } + let classifier = py + .import(ROUTE_HOST_MODULE)? + .getattr("safeguards")? + .call1((request,))?; + if !classifier.is_none() { + self.classifier = Some(classifier.unbind()); + } + Ok(()) + } + + pub(super) fn begin(&mut self, py: Python<'_>, call: SafeguardsCall) -> PyResult> { + let classifier = self + .classifier + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("safeguards call without a classifier"))? + .bind(py); + let (awaitable, pending) = match call { + SafeguardsCall::Results { response, reply } => ( + classifier.call_method1("results", (to_py(py, &response)?,))?, + Pending::Results(reply), + ), + SafeguardsCall::OpenStream { reply } => ( + classifier.call_method0("open_stream")?, + Pending::OpenStream(reply), + ), + SafeguardsCall::Frames { chunk, reply } => ( + classifier.call_method1("frames", (PyBytes::new(py, &chunk),))?, + Pending::Frames(reply), + ), + SafeguardsCall::Flush { reply } => { + (classifier.call_method0("flush")?, Pending::Frames(reply)) + } + }; + self.pending = Some(pending); + Ok(awaitable.unbind()) + } + + pub(super) fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult<()> { + let pending = self + .pending + .take() + .ok_or_else(|| PyRuntimeError::new_err("safeguards reply without a pending call"))?; + let value = result?.into_bound(py); + match pending { + Pending::Results(reply) => reply.send(if value.is_none() { + None + } else { + Some(from_py(&value)?) + }), + Pending::OpenStream(reply) => reply.send(value.extract()?), + Pending::Frames(reply) => reply.send( + value + .extract::>>()? + .iter() + .map(|frame| Bytes::copy_from_slice(frame.as_bytes())) + .collect(), + ), + } + Ok(()) + } + + pub(super) fn published(&self, py: Python<'_>, public: Py) -> PyResult> { + match &self.classifier { + Some(classifier) => classifier + .bind(py) + .call_method1("published", (public,)) + .map(Bound::unbind), + None => Ok(public), + } + } + + pub(super) fn close(&mut self) { + self.classifier = None; + self.pending = None; + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.classifier) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/transport.rs b/litellm-rust/crates/python-bridge/src/routes/messages/transport.rs new file mode 100644 index 00000000000..12e782f61aa --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/transport.rs @@ -0,0 +1,275 @@ +use std::fmt; + +use bytes::Bytes; +use futures_util::{StreamExt, future::BoxFuture}; +use litellm_host::{ + machine::HostServices, + protocol::{Protocol, Reply}, +}; +use litellm_host_python::{json_to_py, to_py}; +use litellm_inference_messages::{ + Delivered, Error, Transport, TransportBody, TransportRequest, TransportResponse, +}; +use pyo3::{ + exceptions::{PyBaseException, PyException, PyStopAsyncIteration}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyBytes, PyDict}, +}; + +use super::host::ROUTE_HOST_MODULE; + +pub(crate) enum TransportCall { + Send { + request: TransportRequest, + reply: Reply>, + }, + Read { + reply: Reply, Error>>, + }, +} + +pub(crate) struct Sent { + status: u16, + headers: Vec<(String, String)>, + body: SentBody, +} + +enum SentBody { + Text(String), + Stream, + Rejected { text: String, error: Error }, +} + +pub(super) struct PythonTransport { + services: HostServices

, +} + +impl PythonTransport

{ + pub(super) fn new(services: HostServices

) -> Self { + Self { services } + } +} + +impl

Transport for PythonTransport

+where + P: Protocol, + P::HostCall: From, +{ + fn send( + &self, + request: TransportRequest, + ) -> BoxFuture<'static, Result> { + let services = self.services.clone(); + Box::pin(async move { + let sent = services + .call(|reply| TransportCall::Send { request, reply }.into()) + .await??; + let delivered = match sent.body { + SentBody::Text(text) => Delivered::Complete(text), + SentBody::Rejected { text, error } => Delivered::Rejected { text, error }, + SentBody::Stream => Delivered::Stream( + futures_util::stream::try_unfold(services, |services| async move { + let chunk = services + .call(|reply| TransportCall::Read { reply }.into()) + .await??; + Ok(chunk.map(|chunk| (chunk, services))) + }) + .boxed(), + ), + }; + Ok(TransportResponse { + status: sent.status, + headers: sent.headers, + delivered, + }) + }) + } +} + +#[derive(thiserror::Error)] +#[error("the Python HTTP client failed")] +struct HostTransportFailure { + error: Py, + raw: bool, +} + +impl fmt::Debug for HostTransportFailure { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("HostTransportFailure") + } +} + +fn failed(py: Python<'_>, error: PyErr, raw: bool) -> Error { + Error::transport(HostTransportFailure { + error: error.into_value(py), + raw, + }) +} + +pub(super) fn python_failure( + py: Python<'_>, + error: &(dyn std::error::Error + 'static), +) -> Option<(PyErr, bool)> { + litellm_error::find::(error).map(|failure| { + ( + PyErr::from_value(failure.error.clone_ref(py).into_bound(py).into_any()), + failure.raw, + ) + }) +} + +enum Pending { + Send(Reply>), + Read(Reply, Error>>), +} + +#[derive(Default)] +pub(super) struct HostedTransport { + handler: Option>, + chunks: Option>, + pending: Option, +} + +impl HostedTransport { + pub(super) fn select(&mut self, py: Python<'_>, request: &Bound<'_, PyDict>) -> PyResult { + let handler = py + .import(ROUTE_HOST_MODULE)? + .getattr("transport_handler")? + .call1((request,))?; + self.handler = (!handler.is_none()).then(|| handler.unbind()); + Ok(self.handler.is_some()) + } + + pub(super) fn pending(&self) -> bool { + self.pending.is_some() + } + + pub(super) fn begin( + &mut self, + py: Python<'_>, + call: TransportCall, + request: &Bound<'_, PyDict>, + ) -> PyResult>> { + match call { + TransportCall::Send { + request: sent, + reply, + } => { + let handler = self + .handler + .as_ref() + .ok_or(litellm_host_python::Error::MissingState)?; + let body = match &sent.body { + TransportBody::Json(body) => json_to_py(py, body)?.unbind(), + TransportBody::Signed(bytes) => PyBytes::new(py, bytes).into_any().unbind(), + }; + let awaitable = py + .import(ROUTE_HOST_MODULE)? + .getattr("transport_send")? + .call1(( + handler, + request, + &sent.url, + to_py(py, &sent.headers)?, + body, + sent.stream, + ))?; + self.pending = Some(Pending::Send(reply)); + Ok(Some(awaitable.unbind())) + } + TransportCall::Read { reply } => { + let Some(chunks) = self.chunks.as_ref() else { + reply.send(Ok(None)); + return Ok(None); + }; + let awaitable = chunks.bind(py).call_method0("__anext__")?; + self.pending = Some(Pending::Read(reply)); + Ok(Some(awaitable.unbind())) + } + } + } + + pub(super) fn resume( + &mut self, + py: Python<'_>, + result: PyResult>, + ) -> PyResult>> { + let result = match result { + Err(error) if !error.is_instance_of::(py) => { + self.pending = None; + return Err(error); + } + result => result, + }; + match self.pending.take() { + Some(Pending::Send(reply)) => { + reply.send(match result { + Ok(sent) => self.sent(py, sent.bind(py))?, + Err(error) => Err(self.send_failure(py, error)?), + }); + } + Some(Pending::Read(reply)) => reply.send(match result { + Ok(chunk) => Ok(Some(Bytes::copy_from_slice( + chunk.bind(py).cast::()?.as_bytes(), + ))), + Err(error) if error.is_instance_of::(py) => { + self.chunks = None; + Ok(None) + } + Err(error) => { + self.chunks = None; + Err(failed(py, error, true)) + } + }), + None => return Err(litellm_host_python::Error::MissingState.into()), + } + Ok(None) + } + + fn send_failure(&self, py: Python<'_>, error: PyErr) -> PyResult { + let raw = py + .import(ROUTE_HOST_MODULE)? + .getattr("RawTransportFailure")?; + if error.value(py).is_instance(&raw)? { + let raised = error.value(py).getattr("error")?; + return Ok(failed(py, PyErr::from_value(raised), true)); + } + Ok(failed(py, error, false)) + } + + fn sent(&mut self, py: Python<'_>, sent: &Bound<'_, PyAny>) -> PyResult> { + let status: u16 = sent.getattr("status")?.extract()?; + let headers: Vec<(String, String)> = sent.getattr("headers")?.extract()?; + let text: Option = sent.getattr("text")?.extract()?; + let rejection = sent.getattr("rejection")?; + let chunks = sent.getattr("chunks")?; + let body = if !rejection.is_none() { + SentBody::Rejected { + text: text.unwrap_or_default(), + error: failed(py, PyErr::from_value(rejection), false), + } + } else if !chunks.is_none() { + self.chunks = Some(chunks.unbind()); + SentBody::Stream + } else { + SentBody::Text(text.unwrap_or_default()) + }; + Ok(Ok(Sent { + status, + headers, + body, + })) + } + + pub(super) fn close(&mut self) { + self.pending = None; + self.chunks = None; + self.handler = None; + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.handler)?; + visit.call(&self.chunks) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs b/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs new file mode 100644 index 00000000000..b8824d96584 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/workload_identity.rs @@ -0,0 +1,74 @@ +use std::fmt; + +use futures_util::future::BoxFuture; +use litellm_host_python::{PythonContext, attach_blocking, to_py}; +use litellm_inference_messages::WorkloadIdentity; +use litellm_router_types::{ExtraParams, LitellmParams}; +use litellm_secret_value::SecretValue; +use pyo3::{exceptions::PyBaseException, prelude::*}; + +const MODULE: &str = "litellm.rust_bridge.host.routes.messages"; + +pub(super) struct PythonWorkloadIdentity { + context: PythonContext, +} + +impl PythonWorkloadIdentity { + pub(super) fn new(context: PythonContext) -> Self { + Self { context } + } +} + +impl WorkloadIdentity for PythonWorkloadIdentity { + fn token<'a>( + &'a self, + litellm_params: &'a LitellmParams, + api_base: Option<&'a str>, + ) -> BoxFuture<'a, Result, litellm_auth::Error>> { + let params = litellm_params.extra.clone(); + let api_base = api_base.map(str::to_owned); + let context = self.context.clone(); + Box::pin(async move { + attach_blocking(context, move |py| { + mint(py, ¶ms, api_base).map_err(|error| raised(py, error)) + }) + .await + .map_err(|error| Python::attach(|py| raised(py, error))) + .and_then(|minted| minted) + .map(|token| token.map(SecretValue::new)) + }) + } +} + +fn mint( + py: Python<'_>, + params: &ExtraParams, + api_base: Option, +) -> PyResult> { + py.import(MODULE)? + .getattr("workload_identity_token")? + .call1((to_py(py, &**params)?, api_base))? + .extract() +} + +#[derive(thiserror::Error)] +#[error("Python workload identity federation failed")] +struct PythonWorkloadIdentityError(Py); + +impl fmt::Debug for PythonWorkloadIdentityError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("PythonWorkloadIdentityError") + } +} + +fn raised(py: Python<'_>, error: PyErr) -> litellm_auth::Error { + litellm_auth::Error::provider(PythonWorkloadIdentityError(error.into_value(py))) +} + +pub(crate) fn python_error( + py: Python<'_>, + error: &(dyn std::error::Error + 'static), +) -> Option { + litellm_error::find::(error) + .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f51c755c33e..9fcd792bdaa 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -166,6 +166,38 @@ fn run_public_call( hooks: impl PythonCallHooks + 'static, asynchronous: bool, ) -> PyResult> +where + H: PythonBinding + PythonHostCalls + 'static, + M: Machine + 'static, + M::Complete: Into::Response>>, +{ + run_public_call_after_policy( + py, + arguments, + start, + host, + hooks, + HookChain::new(), + asynchronous, + ) +} + +fn run_public_call_after_policy( + py: Python<'_>, + arguments: Py, + start: impl FnOnce( + Python<'_>, + &Bound<'_, PyDict>, + ::Request, + ) -> PyResult + + Send + + Sync + + 'static, + host: H, + hooks: impl PythonCallHooks + 'static, + after_policy: impl PythonCallHooks + 'static, + asynchronous: bool, +) -> PyResult> where H: PythonBinding + PythonHostCalls + 'static, M: Machine + 'static, @@ -179,7 +211,8 @@ where host, HookChain::new() .with(hooks) - .with(crate::preflight::SdkPolicy), + .with(crate::preflight::SdkPolicy) + .with(after_policy), arguments, CallOptions::new(asynchronous, execution_steps), ) diff --git a/litellm-rust/crates/python-compat/src/lib.rs b/litellm-rust/crates/python-compat/src/lib.rs index 2cc0eeac397..e23dbdb285e 100644 --- a/litellm-rust/crates/python-compat/src/lib.rs +++ b/litellm-rust/crates/python-compat/src/lib.rs @@ -7,6 +7,7 @@ //! |---|---| //! | [`literal`] | `ast.literal_eval(text)` | //! | [`repr`] | `repr(value)` and `str(value)` | +//! | [`reprlib`] | `reprlib.repr(value)` | //! | [`json`] | `json.dumps(value)` and `json.loads(text)` | //! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data | //! | [`truthy`] | `bool(value)` | @@ -29,6 +30,7 @@ pub mod literal; pub mod number; pub mod pickle; pub mod repr; +pub mod reprlib; pub mod truthy; mod value; diff --git a/litellm-rust/crates/python-compat/src/reprlib.rs b/litellm-rust/crates/python-compat/src/reprlib.rs new file mode 100644 index 00000000000..5d04b4e1ca7 --- /dev/null +++ b/litellm-rust/crates/python-compat/src/reprlib.rs @@ -0,0 +1,179 @@ +use crate::{Value, repr::repr}; + +const MAX_LEVEL: usize = 6; +const MAX_TUPLE: usize = 6; +const MAX_LIST: usize = 6; +const MAX_SET: usize = 6; +const MAX_DICT: usize = 4; +const MAX_STRING: usize = 30; +const MAX_LONG: usize = 40; +const MAX_OTHER: usize = 30; +const FILL: &str = "..."; + +pub fn reprlib(value: &Value) -> String { + repr1(value, MAX_LEVEL) +} + +fn repr1(value: &Value, level: usize) -> String { + match value { + Value::Str(text) => repr_str(text), + Value::Int(_) => shortened(repr(value), MAX_LONG), + Value::Tuple(values) => repr_iterable(values, level, ("(", ")"), MAX_TUPLE, ","), + Value::List(values) => repr_iterable(values, level, ("[", "]"), MAX_LIST, ""), + Value::Set(values) if values.is_empty() => "set()".to_owned(), + Value::Set(values) => repr_iterable(&sorted(values), level, ("{", "}"), MAX_SET, ""), + Value::Dict(entries) => repr_dict(entries, level), + Value::None + | Value::Bool(_) + | Value::Float(_) + | Value::Complex { .. } + | Value::Bytes(_) => shortened(repr(value), MAX_OTHER), + } +} + +fn shortened(text: String, limit: usize) -> String { + let chars: Vec = text.chars().collect(); + if chars.len() <= limit { + return text; + } + let head = (limit - FILL.len()) / 2; + let tail = limit - FILL.len() - head; + let mut out: String = chars[..head].iter().collect(); + out.push_str(FILL); + out.extend(&chars[chars.len() - tail..]); + out +} + +fn repr_str(text: &str) -> String { + let chars: Vec = text.chars().collect(); + let quoted = repr(&Value::Str(chars.iter().take(MAX_STRING).collect())); + if quoted.chars().count() <= MAX_STRING { + return quoted; + } + let head = (MAX_STRING - FILL.len()) / 2; + let tail = MAX_STRING - FILL.len() - head; + let kept: String = chars[..head.min(chars.len())] + .iter() + .chain(&chars[chars.len().saturating_sub(tail)..]) + .collect(); + let quoted: Vec = repr(&Value::Str(kept)).chars().collect(); + let mut out: String = quoted[..head].iter().collect(); + out.push_str(FILL); + out.extend("ed[quoted.len() - tail..]); + out +} + +fn repr_iterable( + values: &[Value], + level: usize, + (left, right): (&str, &str), + max_items: usize, + trail: &str, +) -> String { + if level == 0 && !values.is_empty() { + return format!("{left}{FILL}{right}"); + } + let mut pieces: Vec = values + .iter() + .take(max_items) + .map(|value| repr1(value, level - 1)) + .collect(); + if values.len() > max_items { + pieces.push(FILL.to_owned()); + } + let trail = if values.len() == 1 { trail } else { "" }; + format!("{left}{}{trail}{right}", pieces.join(", ")) +} + +fn repr_dict(entries: &[(Value, Value)], level: usize) -> String { + if entries.is_empty() { + return "{}".to_owned(); + } + if level == 0 { + return format!("{{{FILL}}}"); + } + let mut pieces: Vec = sorted_entries(entries) + .into_iter() + .take(MAX_DICT) + .map(|(key, value)| format!("{}: {}", repr1(key, level - 1), repr1(value, level - 1))) + .collect(); + if entries.len() > MAX_DICT { + pieces.push(FILL.to_owned()); + } + format!("{{{}}}", pieces.join(", ")) +} + +fn sorted_entries(entries: &[(Value, Value)]) -> Vec<&(Value, Value)> { + let mut ordered: Vec<&(Value, Value)> = entries.iter().collect(); + if entries.iter().all(|(key, _)| matches!(key, Value::Str(_))) { + ordered.sort_by(|(left, _), (right, _)| match (left, right) { + (Value::Str(left), Value::Str(right)) => left.cmp(right), + _ => std::cmp::Ordering::Equal, + }); + } + ordered +} + +fn sorted(values: &[Value]) -> Vec { + let mut ordered = values.to_vec(); + if values.iter().all(|value| matches!(value, Value::Str(_))) { + ordered.sort_by(|left, right| match (left, right) { + (Value::Str(left), Value::Str(right)) => left.cmp(right), + _ => std::cmp::Ordering::Equal, + }); + } + ordered +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn list(values: impl IntoIterator) -> Value { + Value::List(values.into_iter().collect()) + } + + fn dict(entries: &[(&str, Value)]) -> Value { + Value::Dict( + entries + .iter() + .map(|(key, value)| (Value::from(*key), value.clone())) + .collect(), + ) + } + + fn nested(depth: usize) -> Value { + (0..depth).fold(list([Value::from(1)]), |inner, _| list([inner])) + } + + #[rstest] + #[case::none(Value::None, "None")] + #[case::bool(Value::Bool(true), "True")] + #[case::int(Value::from(3), "3")] + #[case::wide_int(Value::Int("1".repeat(45).parse().unwrap()), "111111111111111111...1111111111111111111")] + #[case::float(Value::Float(0.1), "0.1")] + #[case::long_float(Value::Float(1.0 / 3.0), "0.3333333333333333")] + #[case::nan(Value::Float(f64::NAN), "nan")] + #[case::short_str(Value::from("high"), "'high'")] + #[case::long_str( + Value::from("abcdefghijklmnopqrstuvwxyz0123456789"), + "'abcdefghijkl...xyz0123456789'" + )] + #[case::quoted_str( + Value::from("it's a fairly long string with a quote"), + "\"it's a fairl... with a quote\"" + )] + #[case::list(list([Value::from(1)]), "[1]")] + #[case::long_list(list((0..8).map(Value::from)), "[0, 1, 2, 3, 4, 5, ...]")] + #[case::one_tuple(Value::Tuple(vec![Value::from(1)]), "(1,)")] + #[case::sorted_dict(dict(&[("b", Value::from(2)), ("a", Value::from(1))]), "{'a': 1, 'b': 2}")] + #[case::long_dict(dict(&[("e", Value::None), ("d", Value::None), ("c", Value::None), ("b", Value::None), ("a", Value::None)]), "{'a': None, 'b': None, 'c': None, 'd': None, ...}")] + #[case::empty_dict(dict(&[]), "{}")] + #[case::deep(nested(6), "[[[[[[[...]]]]]]]")] + #[case::deep_dict(dict(&[("a", dict(&[("b", dict(&[("c", dict(&[("d", dict(&[("e", dict(&[("f", dict(&[("g", Value::None)]))]))]))]))]))]))]), "{'a': {'b': {'c': {'d': {'e': {'f': {...}}}}}}}")] + fn values_shorten_as_reprlib_shortens_them(#[case] value: Value, #[case] expected: &str) { + assert_eq!(reprlib(&value), expected); + } +} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index b75c1a14446..6955b8a927a 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -31,6 +31,8 @@ impl Router { .model_list .iter() .map(|model| { + let drop_params = + drops_params_globally || is_on(model.litellm_params.drop_params.as_ref(), env); ( model.model_name.clone(), Deployment { @@ -45,8 +47,8 @@ impl Router { litellm_params: model.litellm_params.clone(), shaping: MessagesShaping { settings: MessagesSettings { - drop_params: drops_params_globally - || is_on(model.litellm_params.drop_params.as_ref(), env), + drop_params, + drops_unsupported_params: drop_params, ..MessagesSettings::default() }, ..MessagesShaping::default() @@ -113,15 +115,18 @@ mod tests { )) .unwrap(); + let settings = Router::from_config(&config, &|name: &str| { + (name == "DROP_PARAMS_FLAG").then(|| " Yes ".to_owned()) + }) + .get("public-model") + .unwrap() + .shaping + .settings + .clone(); + assert_eq!( - Router::from_config(&config, &|name: &str| (name == "DROP_PARAMS_FLAG") - .then(|| " Yes ".to_owned())) - .get("public-model") - .unwrap() - .shaping - .settings - .drop_params, - expected + (settings.drop_params, settings.drops_unsupported_params), + (expected, expected) ); } } diff --git a/litellm/__init__.py b/litellm/__init__.py index bfe0e2b5f3e..8a48b0812d9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -27,7 +27,7 @@ from typing import ( from collections.abc import Mapping from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams -from litellm.litellm_core_utils.core_helpers import drop_params_env_flag +from litellm.litellm_core_utils.core_helpers import drop_params_env_flag, env_flag from litellm.types.integrations.pointfive import PointFiveInitParams from litellm.types.integrations.zerobus import ZerobusInitParams from litellm._logging import ( @@ -227,7 +227,7 @@ token: Optional[str] = ( ) max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = drop_params_env_flag(os.environ, verbose_logger) -modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) +modify_params = env_flag(os.environ, "LITELLM_MODIFY_PARAMS", verbose_logger) bedrock_neutralize_orphaned_tool_blocks: bool = True use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index fd72bf7d2e7..b93c88cb935 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -49,7 +49,6 @@ _BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) _DECODED_JSON: Final = TypeAdapter(object) BetaPlacement: TypeAlias = Literal["header", "body"] -UnlistedBetas: TypeAlias = Literal["drop", "forward"] _HEADER_ONLY_COLUMN_SUFFIX: Final = "_header_only" @@ -281,13 +280,12 @@ def filter_and_transform_beta_headers( beta_headers: list[str], provider: str, placement: BetaPlacement = "header", - unlisted: UnlistedBetas = "drop", ) -> list[str]: """ Filter and transform beta headers based on provider's mapping configuration. This function: - 1. Only allows headers that are present in the provider's mapping keys, unless ``unlisted`` is "forward" + 1. Only allows headers that are present in the provider's mapping keys 2. Filters out headers with null values (unsupported) 3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool) @@ -295,7 +293,6 @@ def filter_and_transform_beta_headers( beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") placement: Whether the betas travel as the anthropic-beta header or the anthropic_beta body field - unlisted: Whether a beta missing from the mapping is dropped or forwarded as sent Returns: List of filtered and transformed beta headers for the provider @@ -310,10 +307,6 @@ def filter_and_transform_beta_headers( for header in beta_headers: header = header.strip() - if header not in provider_mapping and unlisted == "forward": - filtered_headers.add(header) - continue - # Check if header is in the mapping if header not in provider_mapping: verbose_logger.debug( @@ -426,7 +419,6 @@ def update_headers_with_filtered_beta( headers: dict, provider: str, api_base: str | None = None, - unlisted: UnlistedBetas = "drop", keep: frozenset[str] = frozenset(), ) -> dict: """ @@ -437,7 +429,6 @@ def update_headers_with_filtered_beta( headers: Request headers dict (will be modified in place) provider: Provider name api_base: Deployment api_base, used to pick the allowlist row (Foundry hosts use ``azure_ai``) - unlisted: Whether a beta missing from the row is dropped or forwarded as sent keep: Betas the routed model accepts, forwarded as sent whatever the row says Returns: @@ -457,7 +448,6 @@ def update_headers_with_filtered_beta( *filter_and_transform_beta_headers( beta_headers=[b for b in beta_values if b not in keep], provider=resolve_beta_provider(provider, api_base), - unlisted=unlisted, ), *(b for b in beta_values if b in keep), ] diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index a3cf54916f1..1e89b700137 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -19,6 +19,7 @@ import datetime import inspect import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeVar from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError @@ -80,6 +81,19 @@ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper EmbeddingCacheInputElement = str | list[int] | ChatCompletionFileObject +@dataclass(frozen=True, slots=True) +class CacheRead: + kwargs: dict[str, object] # mutable-ok: the @client call kwargs the embedding path keeps reading and rewriting + checked: bool = False + cached_result: object | None = None + model: str = "" + end_time: datetime.datetime | None = None + + @property + def hit(self) -> bool: + return self.end_time is not None + + class CachingHandlerResponse(LiteLLMBaseModel): """ This is the response object for the caching handler. We need to separate embedding cached responses and (completion / text_completion / transcription) cached responses @@ -313,6 +327,10 @@ def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None: return request_kwargs.get("cache_key", None) +def _is_cache_miss(cached_result: object) -> bool: + return cached_result is None or isinstance(cached_result, list) + + def _set_cached_hidden_param(response: object, key: str, value: object) -> None: if isinstance(response, TextCompletionResponse): setattr(response.hidden_params, key, value) @@ -406,121 +424,146 @@ class LLMCachingHandler: Raises: None """ + if not self.reads_cache(model=model, kwargs=kwargs): + return None + args = args or () + final_embedding_cached_response: EmbeddingResponse | None = None + embedding_all_elements_cache_hit: bool = False + read: Final = await self.async_read_cache( + model=model, + original_function=original_function, + logging_obj=logging_obj, + call_type=call_type, + kwargs=kwargs, + args=args, + ) + kwargs = read.kwargs + cached_result: object | None = read.cached_result + if litellm.cache is not None and read.checked: + if read.end_time is not None: + cached_result = self._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=original_function.__name__, + kwargs=kwargs, + logging_obj=logging_obj, + model=read.model, + custom_llm_provider=kwargs.get("custom_llm_provider", None), + args=args, + ) + if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result): + # LOG SUCCESS + self._async_log_cache_hit_on_callbacks( + logging_obj=logging_obj, + cached_result=cached_result, + start_time=start_time, + end_time=read.end_time, + cache_hit=True, + ) + _set_cached_hidden_param(cached_result, "cache_key", self.hit_cache_key()) + return CachingHandlerResponse(cached_result=cached_result) + elif ( + call_type == CallTypes.aembedding.value + and cached_result is not None + and isinstance(cached_result, list) + and litellm.cache is not None + and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude. + ): + ( + final_embedding_cached_response, + embedding_all_elements_cache_hit, + ) = self._process_async_embedding_cached_response( + final_embedding_cached_response=final_embedding_cached_response, + cached_result=cached_result, + kwargs=kwargs, + logging_obj=logging_obj, + start_time=start_time, + model=model, + ) + return CachingHandlerResponse( + final_embedding_cached_response=final_embedding_cached_response, + embedding_all_elements_cache_hit=embedding_all_elements_cache_hit, + embedding_uncached_input=self.handle_kwargs_input_list_or_str(kwargs), + ) + + verbose_logger.debug("CACHE RESULT: %s", cached_result) + return CachingHandlerResponse( + cached_result=cached_result, + final_embedding_cached_response=final_embedding_cached_response, + ) + + _async_get_cache = async_get_cache + + def reads_cache( + self, + model: str, + kwargs: dict[str, Any], # pyright: ignore[reportExplicitAny] # mutable-ok: the @client call kwargs, read as the wrapper reads them + ) -> bool: if _is_response_cache_excluded(model=model, kwargs=kwargs): self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs) - return None - - # Check if caching should be performed BEFORE doing expensive operations - if ( + return False + return ( (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True ) and ( not _cache_reads_disabled(kwargs) - ): # allow users to control returning cached responses from the completion function - args = args or () - final_embedding_cached_response: EmbeddingResponse | None = None - embedding_all_elements_cache_hit: bool = False - cached_result: object | None = None - kwargs = kwargs.copy() - ######################################################### - # Init cache timing metrics - ######################################################### - cache_check_start_time: Final = time.perf_counter() - cache_check_end_time: float | None = None - ######################################################### - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - kwargs["parent_otel_span"] = parent_otel_span + ) # allow users to control returning cached responses from the completion function - if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): - verbose_logger.debug("Checking Async Cache") - logging_obj.reset_caching_details() - cached_result = await self._retrieve_from_cache( - call_type=call_type, - kwargs=kwargs, - args=args, - ) - cache_check_end_time = time.perf_counter() + async def async_read_cache( + self, + model: str, + original_function: Callable[..., object], + logging_obj: LiteLLMLoggingObj, + call_type: str, + kwargs: Mapping[str, Any], # pyright: ignore[reportExplicitAny] # the @client call kwargs, read as the wrapper reads them + args: tuple[object, ...], + ) -> CacheRead: + cache_check_start_time: Final = time.perf_counter() + read_kwargs: Final = dict(kwargs) + read_kwargs["parent_otel_span"] = get_parent_otel_span_from_kwargs(read_kwargs) + if litellm.cache is None or not self._is_call_type_supported_by_cache(original_function=original_function): + return CacheRead(kwargs=read_kwargs) + verbose_logger.debug("Checking Async Cache") + logging_obj.reset_caching_details() + cached_result: Final = await self._retrieve_from_cache( + call_type=call_type, + kwargs=read_kwargs, + args=args, + ) + cache_check_end_time: Final = time.perf_counter() + if _is_cache_miss(cached_result): + return CacheRead(kwargs=read_kwargs, checked=True, cached_result=cached_result) + verbose_logger.debug("Cache Hit!") + end_time: Final = datetime.datetime.now() + provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, + custom_llm_provider=read_kwargs.get("custom_llm_provider", None), + api_base=read_kwargs.get("api_base", None), + api_key=read_kwargs.get("api_key", None), + ) + self._update_litellm_logging_obj_environment( + logging_obj=logging_obj, + model=provider_model, + kwargs=read_kwargs, + cached_result=cached_result, + is_async=True, + custom_llm_provider=custom_llm_provider, + cache_duration_ms=(cache_check_end_time - cache_check_start_time) * 1000, + ) + return CacheRead( + kwargs=read_kwargs, + checked=True, + cached_result=cached_result, + model=provider_model, + end_time=end_time, + ) - if cached_result is not None and not isinstance(cached_result, list): - verbose_logger.debug("Cache Hit!") - cache_hit: Final = True - end_time: Final = datetime.datetime.now() - provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, - custom_llm_provider=kwargs.get("custom_llm_provider", None), - api_base=kwargs.get("api_base", None), - api_key=kwargs.get("api_key", None), - ) - cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000 - self._update_litellm_logging_obj_environment( - logging_obj=logging_obj, - model=provider_model, - kwargs=kwargs, - cached_result=cached_result, - is_async=True, - custom_llm_provider=custom_llm_provider, - cache_duration_ms=cache_duration_ms, - ) - - call_type = original_function.__name__ - - cached_result = self._convert_cached_result_to_model_response( - cached_result=cached_result, - call_type=call_type, - kwargs=kwargs, - logging_obj=logging_obj, - model=provider_model, - custom_llm_provider=kwargs.get("custom_llm_provider", None), - args=args, - ) - if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result): - # LOG SUCCESS - self._async_log_cache_hit_on_callbacks( - logging_obj=logging_obj, - cached_result=cached_result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - ) - cache_key: Final = ( - self.preset_cache_key - or self.request_kwargs.get("cache_key") - or litellm.cache.get_cache_key(**self.request_kwargs) - ) - _set_cached_hidden_param(cached_result, "cache_key", cache_key) - return CachingHandlerResponse(cached_result=cached_result) - elif ( - call_type == CallTypes.aembedding.value - and cached_result is not None - and isinstance(cached_result, list) - and litellm.cache is not None - and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude. - ): - ( - final_embedding_cached_response, - embedding_all_elements_cache_hit, - ) = self._process_async_embedding_cached_response( - final_embedding_cached_response=final_embedding_cached_response, - cached_result=cached_result, - kwargs=kwargs, - logging_obj=logging_obj, - start_time=start_time, - model=model, - ) - return CachingHandlerResponse( - final_embedding_cached_response=final_embedding_cached_response, - embedding_all_elements_cache_hit=embedding_all_elements_cache_hit, - embedding_uncached_input=self.handle_kwargs_input_list_or_str(kwargs), - ) - - verbose_logger.debug("CACHE RESULT: %s", cached_result) - return CachingHandlerResponse( - cached_result=cached_result, - final_embedding_cached_response=final_embedding_cached_response, - ) - # Caching disabled - return None to indicate no caching attempted - return None - - _async_get_cache = async_get_cache + def hit_cache_key(self) -> object: + if litellm.cache is None: + return None + return ( + self.preset_cache_key + or self.request_kwargs.get("cache_key") + or litellm.cache.get_cache_key(**self.request_kwargs) + ) def sync_get_cache( self, @@ -893,14 +936,14 @@ class LLMCachingHandler: _combine_cached_embedding_response_with_api_result = combine_cached_embedding_response_with_api_result - def _async_log_cache_hit_on_callbacks( + def async_log_cache_hit_on_callbacks( self, logging_obj: LiteLLMLoggingObj, cached_result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool, - ): + ) -> None: """ Helper function to log the success of a cached result on callbacks @@ -941,6 +984,8 @@ class LLMCachingHandler: cache_hit=cache_hit, ) + _async_log_cache_hit_on_callbacks = async_log_cache_hit_on_callbacks + async def _retrieve_from_cache( self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...] ) -> object | None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 71b20f92cf7..70cb517168c 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -259,19 +259,21 @@ def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool DROP_PARAMS_ENV_VAR: Final = "LITELLM_DROP_PARAMS" -def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: - configured: Final = environ.get(DROP_PARAMS_ENV_VAR, "").strip() +def env_flag(environ: Mapping[str, str], name: str, logger: logging.Logger) -> bool: + configured: Final = environ.get(name, "").strip() if configured == "": return False - normalized: Final = normalize_drop_params(configured) + normalized: Final = lax_bool(configured) if normalized is None: - logger.warning( - "%s=%r is not a flag value, treating it as on. Set it to true or false", DROP_PARAMS_ENV_VAR, configured - ) + logger.warning("%s=%r is not a flag value, treating it as on. Set it to true or false", name, configured) return True return normalized +def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: + return env_flag(environ, DROP_PARAMS_ENV_VAR, logger) + + def safe_divide( numerator: float, denominator: float, diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 08c3e0b48a8..16b03a711ac 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -467,6 +467,39 @@ def _short_circuit_safeguards_evaluator( return request_safeguards_evaluator(kwargs, messages) +def provider_after_pre_request_hooks( + model: str, custom_llm_provider: str | None, request_kwargs: Mapping[str, object] +) -> str | None: + if custom_llm_provider: + return custom_llm_provider + hook_params: Final = cast( # cast-ok: read as the hooks left it, so a non-mapping raises as Python's handler raises + Mapping[str, object], request_kwargs.get("litellm_params", {}) + ) + derived: Final = cast( # cast-ok: handed on unchecked, as Python's handler does + str | None, hook_params.get("custom_llm_provider") + ) + if derived: + return derived + try: + _, resolved, _, _ = litellm.get_llm_provider(model=model) + except Exception: # noqa: BLE001 # an unknown model keeps the provider the hooks left + return derived + return resolved + + +def tools_marked_for_cache_points( + tools: object, + request_kwargs: dict[str, object], # mutable-ok: pops the points and records the mark in the handler's kwargs +) -> Sequence[object] | None: + marked: Final = tools_marked_for_a_tool_config_point( + tools, request_kwargs.pop("cache_control_injection_points", None) + ) + if marked is None: + return None + AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1) + return list(marked) + + @client async def anthropic_messages( max_tokens: int, @@ -527,16 +560,7 @@ async def anthropic_messages( tool_choice = request_kwargs.pop("tool_choice", tool_choice) top_k = request_kwargs.pop("top_k", top_k) top_p = request_kwargs.pop("top_p", top_p) - # Propagate the provider derived inside pre-request hooks, if not already set. - # The litellm_params dict may have been overwritten by **kwargs in - # _execute_pre_request_hooks, so fall back to get_llm_provider() if needed. - if not custom_llm_provider: - custom_llm_provider = request_kwargs.get("litellm_params", {}).get("custom_llm_provider") - if not custom_llm_provider: - try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) - except Exception: - pass + custom_llm_provider = provider_after_pre_request_hooks(model, custom_llm_provider, request_kwargs) # Remove litellm_params from kwargs (only needed for hooks) request_kwargs.pop("litellm_params", None) # Merge back any other modifications @@ -869,12 +893,8 @@ def anthropic_messages_handler( ) request_kwargs: Final = cast(dict[str, object], kwargs) # cast-ok: **kwargs of the handler is untyped - marked_tools: Final = tools_marked_for_a_tool_config_point( - local_vars.get("tools"), request_kwargs.pop("cache_control_injection_points", None) - ) - if marked_tools is not None: - AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1) - local_vars.update({**request_kwargs, **({"tools": list(marked_tools)} if marked_tools is not None else {})}) + marked_tools: Final = tools_marked_for_cache_points(local_vars.get("tools"), request_kwargs) + local_vars.update({**request_kwargs, **({"tools": marked_tools} if marked_tools is not None else {})}) anthropic_messages_optional_request_params: Final = ( AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( params=local_vars, diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index b0558292271..c2c7244259b 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -105,7 +105,6 @@ _MESSAGES_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): - _workload_identity_eligible: ClassVar[bool] = True _report_reasoning_drops: ClassVar[bool] = True @property @@ -372,12 +371,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): self._apply_env_auth_header( headers, self._require_auth_header( - AnthropicModelInfo.get_auth_header( - api_key, - api_base=api_base, - litellm_params=litellm_params, - allow_workload_identity=self._allows_workload_identity, - ), + AnthropicModelInfo.get_auth_header(api_key, api_base=api_base), model=model, ), ) @@ -412,12 +406,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): self._apply_env_auth_header( oauth_headers, self._require_auth_header( - await AnthropicModelInfo.aget_auth_header( - oauth_api_key, - api_base=api_base, - litellm_params=litellm_params, - allow_workload_identity=self._allows_workload_identity, - ), + await AnthropicModelInfo.aget_auth_header(oauth_api_key, api_base=api_base), model=model, ), ) @@ -447,14 +436,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if merged_beta: headers["anthropic-beta"] = merged_beta - @property - def _allows_workload_identity(self) -> bool: - """Subclasses reuse this validate step for their own /v1/messages-compatible providers, so - eligibility is declared per class and never inherited.""" - from litellm.llms.anthropic.common_utils import config_allows_workload_identity - - return config_allows_workload_identity(self) - def finalize_messages_headers( self, headers: dict, # mutable-ok: out-param diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index a5b15a24a7f..b290a7abf3d 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -345,6 +345,10 @@ def _set_ciphers( ) +def rsa_key_exchange_and_chacha20_excluded() -> bool: + return _rsa_key_exchange_and_chacha20_excluded.is_set() + + def exclude_rsa_key_exchange_and_chacha20(excluded: bool) -> None: """ Turn RSA key exchange and TLS 1.2 ChaCha20 off or back on for new SSL contexts and every cached one, which clients diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f65c4205944..68b13c6172f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -41,7 +41,6 @@ from litellm._internal_context import internal_sub_call from litellm._logging import redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import ( aload_beta_headers_config, - resolve_beta_provider, update_headers_with_filtered_beta, ) from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -82,7 +81,7 @@ from litellm.litellm_core_utils.realtime_errors import ( ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers +from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -2633,17 +2632,10 @@ class BaseLLMHTTPHandler: litellm_params.get("additional_drop_params") or (), anthropic_messages_optional_request_params, ) - affinity_api_base: Final = ( - AnthropicModelInfo.get_api_base(api_base) if custom_llm_provider == "anthropic" else api_base - ) - takes_another_beta_row: Final = ( - resolve_beta_provider(custom_llm_provider, affinity_api_base) != custom_llm_provider - ) - provider_filters_betas: Final = anthropic_messages_provider_config.should_filter_anthropic_beta_headers() - filters_beta_headers: Final = provider_filters_betas or takes_another_beta_row + filters_beta_headers: Final = anthropic_messages_provider_config.should_filter_anthropic_beta_headers() if request_params.get("safeguards") is not None and filters_beta_headers: await aload_beta_headers_config() - default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) + default_affinity_header: Final = get_provider_affinity_header(api_base) ( headers, api_base, @@ -2666,8 +2658,6 @@ class BaseLLMHTTPHandler: headers = update_headers_with_filtered_beta( headers=headers, provider=custom_llm_provider, - api_base=affinity_api_base, - unlisted="drop" if provider_filters_betas else "forward", keep=anthropic_messages_provider_config.model_specific_anthropic_betas(model, litellm_params), ) diff --git a/litellm/llms/dashscope/messages/transformation.py b/litellm/llms/dashscope/messages/transformation.py index 9d0d1db91c8..81f93a32166 100644 --- a/litellm/llms/dashscope/messages/transformation.py +++ b/litellm/llms/dashscope/messages/transformation.py @@ -2,7 +2,7 @@ DashScope Anthropic-compatible Messages transformation config. """ -from typing import ClassVar, Final +from typing import Final import litellm from litellm.exceptions import AuthenticationError @@ -25,8 +25,6 @@ class DashScopeAnthropicMessagesConfig(AnthropicMessagesConfig): ``dashscope.aliyuncs.com`` and ``{workspace}.{region}.maas.aliyuncs.com``. """ - _workload_identity_eligible: ClassVar[bool] = False - @property def custom_llm_provider(self) -> str | None: return "dashscope" diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index faa604a25ce..b974e677386 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -3,14 +3,14 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Itera from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable -from litellm.exceptions import BadRequestError +from litellm.exceptions import BadRequestError, UnsupportedParamsError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.messages.native import handler as main from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import detectors +from litellm.rust_bridge import catalog, detectors from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target +from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch -from litellm.rust_bridge.messages import detectors as route_detectors from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, NATIVE_MESSAGES, @@ -22,7 +22,6 @@ from litellm.rust_bridge.public_call import ( native_acall_hook, native_call, native_call_hook, - optional_sequence, optional_str, signature, ) @@ -66,7 +65,7 @@ def _native_kwargs(fields: Mapping[str, object], kwargs: Mapping[str, object]) - request_timeout=kwargs.get("request_timeout"), stream_timeout=kwargs.get("stream_timeout"), ), - stream=fields.get("stream") is True, + stream=bool(fields.get("stream")), custom_llm_provider=provider or "", ) if timeout is None: @@ -81,9 +80,7 @@ def _public_request( if fields is None: return None model: Final = fields.get("model") - messages: Final = optional_sequence(fields.get("messages")) - max_tokens: Final = fields.get("max_tokens") - if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int): + if not isinstance(model, str): return None return native_call(legacy, args, _native_kwargs(fields, kwargs)) @@ -133,30 +130,20 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False) -_FEATURES: Final[detectors.Detectors] = MappingProxyType( +_NATIVE_COMMON_FEATURES: Final = frozenset( { - **{feature: detect for feature, detect in detectors.COMMON.items() if feature is not Feature.TIMEOUT_VARIANTS}, - Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, - Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, - Feature.RESPONSE_CACHE: detectors.response_cache, - Feature.MODIFY_PARAMS: detectors.modify_params, - Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, - Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, - Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, - Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry, - Feature.TOOL_CUSTOM_FIELD: route_detectors.tool_custom_field, + Feature.TIMEOUT_VARIANTS, + Feature.LATEST_ALIAS, + Feature.MOCK_RESPONSE, + Feature.CALLER_CLIENT, + Feature.TLS_POLICY, + Feature.AGENTIC_LOOP_HOOK, + Feature.PROXY_MCP_TOOLS, } ) -_ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( - { - **_FEATURES, - Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook, - Feature.FOUNDRY_BETA_ROW: route_detectors.foundry_beta_row, - Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception, - Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted, - Feature.STRICT_TOOLS_RETRY: route_detectors.strict_tools_retry, - } +_FEATURES: Final[detectors.Detectors] = MappingProxyType( + {feature: detect for feature, detect in detectors.COMMON.items() if feature not in _NATIVE_COMMON_FEATURES} ) @@ -181,32 +168,87 @@ _DISPATCH: Final = PublicDispatch( _ADISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs), - context=lambda request: _context(request, _ASYNC_FEATURES), + context=lambda request: _context(request, _FEATURES), ) +def _unsupported_features(features: detectors.Detectors, request: NativeCall) -> tuple[str, ...]: + try: + return tuple(sorted(feature.value for feature in detectors.detected(features, request))) + except Exception: # noqa: BLE001 # a failing detector names nothing, and the call is refused all the same + return () + + +def _checked_target(request: NativeCall, model: str, provider: str | None) -> Target | None: + try: + return _target(request, model, provider) + except Exception: # noqa: BLE001 # the provider alone still decides whether the call has a Python path + return None + + +def _rust_only_refusal( + legacy: inspect.Signature, features: detectors.Detectors, args: tuple[object, ...], kwargs: Mapping[str, object] +) -> UnsupportedParamsError | None: + if bind(legacy, args, kwargs) is None: + return None + request: Final = native_call(legacy, args, kwargs) + requested_model: Final = request.resolved.get("model") + if not isinstance(requested_model, str): + return None + model, provider = _resolved(request) + context: Final = RouteContext( + Route.MESSAGES, provider=provider, model=requested_model, target=_checked_target(request, model, provider) + ) + if catalog.rollout(context) is not Rollout.RUST_REQUIRED: + return None + unsupported: Final = _unsupported_features(features, request) + return UnsupportedParamsError( + message=( + f"provider {provider} on /v1/messages runs only on the Rust path, which does not support " + f"{', '.join(unsupported) if unsupported else 'this call'}" + ), + llm_provider=provider, + model=requested_model, + ) + + +def _python_messages_handler( + *args: object, + **kwargs: object, # kwargs-ok: forwards the public call shape +) -> MessagesResult | Coroutine[object, object, MessagesResult]: + refusal: Final = _rust_only_refusal(_MESSAGES, _FEATURES, args, kwargs) + if refusal is not None: + raise refusal + return _PYTHON_MESSAGES(*args, **kwargs) + + +async def _python_anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: # kwargs-ok: public call shape + refusal: Final = _rust_only_refusal(_AMESSAGES, _FEATURES, args, kwargs) + if refusal is not None: + raise refusal + return await _PYTHON_AMESSAGES(*args, **kwargs) + + def anthropic_messages_handler( *args: object, **kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape ) -> MessagesResult | Coroutine[object, object, MessagesResult]: - python: Final = _PYTHON_MESSAGES if _FLAG.requested(kwargs): return anthropic_messages(**_FLAG.keywords(args, kwargs)) return _DISPATCH.run( args, _FLAG.dropped(kwargs), - python=python, + python=_python_messages_handler, binding=NATIVE_MESSAGES, native=native_call_hook, ) async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: # kwargs-ok: public call shape - python: Final = _PYTHON_AMESSAGES return await _ADISPATCH.arun( args, _FLAG.dropped(kwargs), - python=python, + python=_python_anthropic_messages, binding=NATIVE_AMESSAGES, native=native_acall_hook, ) diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index a13f57fae4b..ca71fc59c28 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -21,8 +21,11 @@ - `tests/test_litellm_rust/test_public_entries.py` marks such a route `python=None`, so dropping a route's Python changes its entries, its rules and that field together - A detector reads the same inputs the legacy Python path reads, so the catalog and the fallback agree on which requests need Python - A route's context, its target as well as its features, reads `NativeCall.resolved`, so an argument counts the same whether it came by position, by keyword or as a signature default - - A route's sync and async dispatches read different detector maps when the legacy handlers differ: Messages runs pre-request hooks, web search interception and interceptors only on async calls, and the sync Responses handler runs the deployment hook itself + - A route's sync and async dispatches read different detector maps when the legacy handlers differ, as the sync Responses handler runs the deployment hook itself. Messages reads one map for both, since Rust serves its async-only pre-request hooks, web search short-circuit and interceptors through host calls - A detector that raises keeps the call on Python. `PublicDispatch` logs it and never fails the call over routing +- Messages for provider `anthropic` is `RUST_REQUIRED` and has no Python path + - a call the rules would send to Python (a `MESSAGES_PYTHON_FEATURES` feature or a raising detector) raises `UnsupportedParamsError` from the Python fallback in `litellm/messages/dispatch.py`, which names the detected features + - other Messages providers keep the Python fallback - `catalog.RULES` is rollout policy with no tests of its own. Tests pass the rules they need, so check these invariants by review whenever `RULES` changes # Public entries diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index ebbd45c7382..bde79aeac48 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -21,6 +21,7 @@ class NativeRedisStoreError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... +def unsupported_tls_settings(ssl_security_level: str | None = None, ssl_ecdh_curve: str | None = None) -> list[str]: ... @final class NativeCostMapSource: def __new__(cls, url: str, *, use_http_settings: bool) -> NativeCostMapSource: ... @@ -732,6 +733,7 @@ __all__ = [ "responses", "trace_encode_error", "transcription", + "unsupported_tls_settings", ] @final diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c493617003d..87f4c855561 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -32,8 +32,6 @@ class Target(str, Enum): class Feature(str, Enum): - SAFEGUARDS_CLASSIFIER = "safeguards_classifier" - """A non-Claude model with dangerous_tool_use safeguards, which the proxy classifier answers.""" MOCK_RESPONSE = "mock_response" """mock_response short-circuits before any provider call.""" PROXY_MCP_TOOLS = "proxy_mcp_tools" @@ -66,14 +64,6 @@ class Feature(str, Enum): """Extra body fields that Python transforms and merges into the provider request.""" PROVIDER_AFFINITY = "provider_affinity" """A provider affinity header that Python validates and populates from the session id.""" - FOUNDRY_BETA_ROW = "foundry_beta_row" - """A Foundry api_base, where Python drops the anthropic-beta values Foundry rejects.""" - PRE_REQUEST_HOOK = "pre_request_hook" - """A logger overriding async_pre_request_hook, which may rewrite the request.""" - WEBSEARCH_INTERCEPTION = "websearch_interception" - """A WebSearchInterceptionLogger that may answer without calling the provider.""" - MESSAGES_INTERCEPTOR = "messages_interceptor" - """A registered MessagesInterceptor that takes over the call.""" AGENTIC_LOOP_HOOK = "agentic_loop_hook" """A logger overriding async_should_run_agentic_loop, which may run follow-up calls.""" RESPONSE_CACHE = "response_cache" @@ -94,14 +84,6 @@ class Feature(str, Enum): """No static key, so Python mints a workload identity federation token.""" CACHE_CONTROL_INJECTION = "cache_control_injection" """Prompt caching injection that Python applies to messages and system.""" - STREAM_COST_INJECTION = "stream_cost_injection" - """litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage.""" - THINKING_SIGNATURE_RETRY = "thinking_signature_retry" - """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" - TOOL_CUSTOM_FIELD = "tool_custom_field" - """Tools carrying a non-spec custom object, which Python drops or hoists for the provider.""" - STRICT_TOOLS_RETRY = "strict_tools_retry" - """Strict tools under drop_params, which Python sends again without strict if Claude rejects their limits.""" CLIENT_HEADERS = "client_headers" """Proxy-forwarded headers that Python merges into extra_headers.""" DEPLOYMENT_HOOK = "deployment_hook" @@ -172,21 +154,14 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS}) | { - Feature.PROVIDER_AFFINITY, - Feature.FOUNDRY_BETA_ROW, - Feature.SAFEGUARDS_CLASSIFIER, - Feature.PRE_REQUEST_HOOK, - Feature.WEBSEARCH_INTERCEPTION, - Feature.MESSAGES_INTERCEPTOR, - Feature.RESPONSE_CACHE, - Feature.MODIFY_PARAMS, - Feature.WORKLOAD_IDENTITY, - Feature.CACHE_CONTROL_INJECTION, - Feature.STREAM_COST_INJECTION, - Feature.THINKING_SIGNATURE_RETRY, - Feature.TOOL_CUSTOM_FIELD, - Feature.STRICT_TOOLS_RETRY, +MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES - { + Feature.TIMEOUT_VARIANTS, + Feature.LATEST_ALIAS, + Feature.MOCK_RESPONSE, + Feature.CALLER_CLIENT, + Feature.TLS_POLICY, + Feature.AGENTIC_LOOP_HOOK, + Feature.PROXY_MCP_TOOLS, } RESPONSES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { @@ -220,7 +195,7 @@ RULES: Final[Rules] = ( RouteRule(Route.OCR, Rollout.RUST_REQUIRED), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES), - RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, providers=frozenset({"anthropic"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY, features=RESPONSES_PYTHON_FEATURES), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index c38c3ea7135..1dd3dc6b35c 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -2,10 +2,11 @@ from __future__ import annotations +import functools import os from collections.abc import Callable, Iterable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias, cast # noqa: TID251 # tool lists stay caller-owned +from typing import Final, TypeAlias import httpx @@ -16,9 +17,6 @@ from litellm.litellm_core_utils.model_latest_aliases import is_model_latest_alia from litellm.rust_bridge.catalog import Feature from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str -if TYPE_CHECKING: - from litellm.responses.mcp.litellm_proxy_mcp_handler import ToolParam - Detector: TypeAlias = Callable[[NativeCall], bool] Detectors: TypeAlias = Mapping[Feature, Detector] @@ -69,12 +67,10 @@ def timeout_variants(request: NativeCall) -> bool: def proxy_mcp_tools(request: NativeCall) -> bool: """Python expands proxy MCP tools and may execute their tool calls.""" - from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway tools: Final = optional_sequence(request.resolved.get("tools")) - return tools is not None and LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway( - cast("Iterable[ToolParam]", tools) # cast-ok: the predicate skips entries that are not MCP tool dicts - ) + return tools is not None and may_name_litellm_mcp_gateway(tools) def prompt_management(request: NativeCall) -> bool: @@ -168,15 +164,27 @@ def response_cache(request: NativeCall) -> bool: return litellm.cache is not None and cache.get("no-cache") is not True -_NATIVE_ECDH_CURVES: Final = frozenset({"x25519", "prime256v1", "secp256r1", "p-256", "secp384r1", "p-384"}) +@functools.lru_cache(maxsize=32) +def _tls_unsupported(security_level: str | None, ecdh_curve: str | None) -> bool: + from litellm.rust_bridge import _native + + return bool(_native.unsupported_tls_settings(security_level, ecdh_curve)) def tls_policy(_request: NativeCall) -> bool: """Python applies TLS settings or ECDH curves the native client cannot configure.""" - curve: Final = os.getenv("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve) - return bool(os.getenv("SSL_SECURITY_LEVEL", litellm.ssl_security_level)) or ( - bool(curve) and str(curve).strip().lower() not in _NATIVE_ECDH_CURVES + return unsupported_tls() + + +def unsupported_tls() -> bool: + configured: Final[tuple[object, object]] = ( + os.environ.get("SSL_SECURITY_LEVEL", litellm.ssl_security_level), + os.environ.get("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve), ) + if any(value and not isinstance(value, str) for value in configured): # pyright: ignore[reportUnnecessaryIsInstance] # callers assign the ssl settings untyped + return True + security_level, ecdh_curve = (value if isinstance(value, str) else None for value in configured) + return _tls_unsupported(security_level, ecdh_curve) COMMON: Final[Detectors] = MappingProxyType( diff --git a/litellm/rust_bridge/execution.py b/litellm/rust_bridge/execution.py index 430de8d1ad6..a77fd162bf4 100644 --- a/litellm/rust_bridge/execution.py +++ b/litellm/rust_bridge/execution.py @@ -101,8 +101,10 @@ def synchronous(native: Callable[[_Call], object]) -> Callable[[_Call], object]: return lambda call: run(_execution(native(call)), SyncStream) -def asynchronous(native: Callable[[_Call], object]) -> Callable[[_Call], Awaitable[object]]: - return lambda call: drive(_execution(native(call)), AsyncStream) +def asynchronous( + native: Callable[[_Call], object], open_stream: Callable[[Execution, object], object] | None = None +) -> Callable[[_Call], Awaitable[object]]: + return lambda call: drive(_execution(native(call)), open_stream or AsyncStream) def _execution(value: object) -> Execution: diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index c249eb82ee0..da933f3b467 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -25,7 +25,7 @@ Python behavior Rust needs but does not own, behind one module per question. Rus - `logger`: where native diagnostics are emitted - `preflight`: the SDK globals the native preflight borrows - `secret_manager`: the configured secret manager, its native config and the error text Python's managers produce - - `http_transport`: which async transport is in use + - `http_transport`: which async transport is in use, and whether a call must send through Python's `AsyncHTTPHandler` - `timeouts`: timeout to seconds - `tokenizers`: bundled tiktoken files and Hugging Face tokenizer downloads - `randomness`: the process `random` generator the Python router draws from diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py index bf2655f9f62..38a18b9be57 100644 --- a/litellm/rust_bridge/host/cache.py +++ b/litellm/rust_bridge/host/cache.py @@ -1,7 +1,8 @@ from __future__ import annotations -from collections.abc import Mapping -from typing import TYPE_CHECKING, Final +import weakref +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # bounded cast of the served response from pydantic import TypeAdapter @@ -9,9 +10,13 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_par if TYPE_CHECKING: from litellm.caching.caching import Cache + from litellm.litellm_core_utils.litellm_logging import Logging _RUST_RESPONSE_HEADER: Final = "x-litellm-rust" _HEADERS: Final = TypeAdapter(dict[str, object]) +_EVENTS: Final = TypeAdapter(list[str]) +_HITS: Final[weakref.WeakKeyDictionary[Logging, Mapping[str, object]]] = weakref.WeakKeyDictionary() +_UNCLAMPED: Final[weakref.WeakKeyDictionary[Logging, object]] = weakref.WeakKeyDictionary() def configured_cache() -> Cache | None: @@ -25,17 +30,14 @@ def _python_entry(response: Mapping[str, object]) -> Mapping[str, object]: headers: Final = hidden.get("additional_headers") if not isinstance(headers, Mapping) or _RUST_RESPONSE_HEADER not in headers: return response - return { - **response, - "_hidden_params": { - **hidden, - "additional_headers": { - name: value - for name, value in _HEADERS.validate_python(headers).items() - if name != _RUST_RESPONSE_HEADER - }, - }, + kept: Final = { + name: value for name, value in _HEADERS.validate_python(headers).items() if name != _RUST_RESPONSE_HEADER } + stored: Final = {name: value for name, value in hidden.items() if name != "additional_headers" or kept} | ( + {"additional_headers": kept} if kept else {} + ) + entry: Final = {name: value for name, value in response.items() if name != "_hidden_params"} + return {**entry, "_hidden_params": stored} if stored else entry def store_response( @@ -44,11 +46,183 @@ def store_response( /, **kwargs: object, # kwargs-ok: the cache call's own arguments ) -> None: - from litellm.caching.caching_handler import create_cache_write_task + from litellm.caching.caching_handler import create_cache_write_task, is_response_without_output + if is_response_without_output(response): + return create_cache_write_task( lambda: cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped _python_entry(response), **kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them ) ) + + +def store_verbatim( + cache: Cache, + response: Mapping[str, object], + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> None: + from litellm.caching.caching_handler import create_cache_write_task, is_response_without_output + + if is_response_without_output(response): + return + create_cache_write_task( + lambda: cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped + response, + **kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them + ) + ) + + +async def store_stream( + cache: Cache, + entry: Mapping[str, object], + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> None: + from litellm._logging import verbose_logger + from litellm.caching.caching_handler import LLMCachingHandler, create_cache_write_task + from litellm.types.caching import CACHED_STREAM_EVENTS_KEY + + events: Final = entry.get(CACHED_STREAM_EVENTS_KEY) + handler: Final = getattr(kwargs.get("litellm_logging_obj"), "llm_caching_handler", None) + preset_cache_key: Final = handler.preset_cache_key if isinstance(handler, LLMCachingHandler) else None + cache_key_override: Final[Mapping[str, object]] = ( + {"cache_key": preset_cache_key} if preset_cache_key is not None else {} + ) + + async def write() -> None: + try: + await cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped + {**entry, CACHED_STREAM_EVENTS_KEY: tuple(_EVENTS.validate_python(events))} + if isinstance(events, list) + else entry, + dynamic_cache_object=handler.dual_cache if isinstance(handler, LLMCachingHandler) else None, + **{**kwargs, **cache_key_override}, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them + ) + except Exception as e: # noqa: BLE001 # a cache write must never surface as a client-visible stream error + verbose_logger.exception("Anthropic Messages stream cache write failed: %s", e) + + create_cache_write_task(write) + + +def caching_handler( + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> object | None: + from litellm.caching.caching_handler import LLMCachingHandler + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging) or not isinstance(logging_obj.llm_caching_handler, LLMCachingHandler): + return None + return logging_obj.llm_caching_handler + + +def wrap_stream( + _cache: Cache, + original_function: Callable[..., object], + stream: object, + read_handler: object | None, + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> object: + from litellm.caching.caching_handler import LLMCachingHandler + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return stream + handler: Final = ( + read_handler + if isinstance(read_handler, LLMCachingHandler) + else LLMCachingHandler( + original_function=original_function, request_kwargs=dict(kwargs), start_time=logging_obj.start_time + ) + ) + return handler.wrap_streaming_result_for_cache(result=stream, call_type=original_function.__name__) + + +def remember_unclamped(logging_obj: object, max_tokens: object) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + if isinstance(logging_obj, Logging): + _UNCLAMPED[logging_obj] = max_tokens + + +async def lookup( + _cache: Cache, + original_function: Callable[..., object], + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> object | None: + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + return await _read( + original_function, + logging_obj, + {**kwargs, "max_tokens": _UNCLAMPED.pop(logging_obj)} if logging_obj in _UNCLAMPED else kwargs, + ) + + +async def _read( + original_function: Callable[..., object], + logging_obj: Logging, + kwargs: dict[str, object], # mutable-ok: the read call's own keyword arguments, handed on as the handler takes them +) -> object | None: + from litellm.caching.caching_handler import LLMCachingHandler + + handler: Final = LLMCachingHandler( + original_function=original_function, request_kwargs=dict(kwargs), start_time=logging_obj.start_time + ) + logging_obj.llm_caching_handler = handler + model: Final = kwargs.get("model") + model_name: Final = model if isinstance(model, str) else "" + if not handler.reads_cache(model=model_name, kwargs=kwargs): + return None + read: Final = await handler.async_read_cache( + model=model_name, + original_function=original_function, + logging_obj=logging_obj, + call_type=original_function.__name__, + kwargs=kwargs, + args=(), + ) + if not read.hit: + return None + _remember(logging_obj, read.cached_result) + return read.cached_result + + +def _remember(logging_obj: Logging, cached: object) -> None: + if isinstance(cached, Mapping): + _HITS[logging_obj] = _HEADERS.validate_python(cached) + + +def _replace(response: object, cached: Mapping[str, object] | None) -> None: + if isinstance(response, dict) and cached is not None: + served: Final = cast( # cast-ok: a Messages response is the str-keyed dict the route serves + "dict[str, object]", response + ) + served.clear() + served.update(cached) + + +def restore_hit(response: object, logger: Logging) -> None: + from litellm.caching.caching_handler import ( + _set_cached_hidden_param, # pyright: ignore[reportPrivateUsage] # the @client cache hit's own hidden params + ) + + _replace(response, _HITS.pop(logger, None)) + _set_cached_hidden_param(response, "cache_hit", True) + + +def hit_cache_key(logger: Logging) -> object: + from litellm.caching.caching_handler import LLMCachingHandler + + handler: Final = logger.llm_caching_handler + return handler.hit_cache_key() if isinstance(handler, LLMCachingHandler) else None diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index ba557f769eb..19124fcffc0 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -15,7 +15,7 @@ import sys import traceback import uuid import weakref -from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -27,6 +27,7 @@ from typing import ( import httpx from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -41,6 +42,7 @@ class MetadataUpdater(Protocol): kwargs: dict[str, object], start_time: datetime.datetime, end_time: datetime.datetime, + cache_hit: bool | None = None, ) -> None: ... @@ -97,6 +99,7 @@ def finalize( kwargs: dict[str, object], start_time: datetime.datetime, end_time: datetime.datetime, + cache_hit: bool, ) -> None: from litellm.litellm_core_utils.llm_response_utils import response_metadata @@ -104,6 +107,16 @@ def finalize( update: Final = cast( # cast-ok: legacy metadata function accepts concrete kwargs MetadataUpdater, response_metadata.update_response_metadata ) + if cache_hit: + from litellm.caching.caching_handler import ( + _set_cached_hidden_param, # pyright: ignore[reportPrivateUsage] # the @client cache hit's own hidden params + ) + from litellm.rust_bridge.host.cache import hit_cache_key, restore_hit + + restore_hit(response, logger) + update(response, logger, logger.model, kwargs, start_time, end_time, cache_hit=True) + _set_cached_hidden_param(response, "cache_key", hit_cache_key(logger)) + return update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) cache_key: Final = logger.model_call_details.get("cache_key") if logger.model_call_details.get("cache_hit") is True and isinstance(cache_key, str): @@ -277,14 +290,76 @@ def post_call( logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) -def record_response(logger: Logging, status: int, headers: Sequence[tuple[str, str]], body: str) -> None: +class _UnreadBody(httpx.SyncByteStream, httpx.AsyncByteStream): + def __iter__(self) -> Iterator[bytes]: + return iter(()) + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in (): + yield chunk + + +def wire_headers(headers: Sequence[tuple[str, str]]) -> httpx.Headers: + return httpx.Headers([(name.encode(), value.encode()) for name, value in headers]) + + +def _sent_request(sent: tuple[str, Mapping[str, str], Mapping[str, object]]) -> httpx.Request: + url, headers, body = sent + return httpx.Request("POST", url, headers=list(headers.items()), content=json.dumps(body, default=str).encode()) + + +def record_response( + logger: Logging, + status: int, + headers: Sequence[tuple[str, str]], + body: str, + sent: tuple[str, Mapping[str, str], Mapping[str, object]] | None, +) -> None: """Stores the provider response the way Python's Messages handler does, for logging to replay.""" - response: Final = httpx.Response(status_code=status, headers=list(headers), content=body.encode()) + request: Final = None if sent is None else _sent_request(sent) + streamed: Final = sent is not None and sent[2].get("stream") is True + response: Final = ( + httpx.Response(status_code=status, stream=_UnreadBody(), request=request) + if streamed + else httpx.Response(status_code=status, content=body.encode(), request=request) + ) + response.headers = wire_headers(headers) details: Final = logger.model_call_details # the call's own logger: recording on it is the point details["httpx_response"] = response details["response_headers"] = dict(response.headers) +class _ThinkingSignatureRetry(TypedDict): + kind: ReadOnly[Literal["thinking_signature"]] + + +class _StrictToolsRetry(TypedDict): + kind: ReadOnly[Literal["strict_tools"]] + dropped: ReadOnly[int] + + +_RETRY_REASONS: Final = TypeAdapter(tuple[_ThinkingSignatureRetry | _StrictToolsRetry, ...]) + + +async def retried_request( + logger: Logging, + logged_body: dict[str, object], # mutable-ok: the retry rewrites the body pre_call logged, as Python's loop does + body: Mapping[str, object], + reasons: Sequence[Mapping[str, object]], +) -> None: + from litellm.llms.anthropic.common_utils import warn_claude_strict_tools_dropped + + if logger.baseline_cache_context is not None: + await logger.invalidate_baseline_cache_estimate("retried_request") + for reason in _RETRY_REASONS.validate_python(reasons): + if reason["kind"] == "strict_tools": + warn_claude_strict_tools_dropped(reason["dropped"]) + for name in logged_body.keys() - body.keys(): + del logged_body[name] # rebind-ok: the logged body is the one logging already holds + logged_body.update({name: value for name, value in body.items() if logged_body.get(name) != value}) + logger.model_call_details.update(logged_body) + + def record_api_window(logger: Logging, start: float, end: float) -> None: """Records the provider exchange's window the way Python's timed `AsyncHTTPHandler.post` does.""" from litellm.litellm_core_utils.logging_utils import ( @@ -296,6 +371,38 @@ def record_api_window(logger: Logging, start: float, end: float) -> None: start_time=datetime.datetime.fromtimestamp(start), # noqa: DTZ006 # the timed post records naive local times end_time=datetime.datetime.fromtimestamp(end), # noqa: DTZ006 # the timed post records naive local times ) + _log_api_window(logger, start, end) + + +def _log_api_window(logger: Logging, start: float, end: float) -> None: + from litellm._logging import verbose_logger + from litellm.litellm_core_utils.logging_utils import ( + _get_parent_otel_span_from_logging_obj, # pyright: ignore[reportPrivateUsage] # the timed post's own span lookup + _get_service_logger, # pyright: ignore[reportPrivateUsage] # the timed post's own service logger + ) + from litellm.types.services import ServiceTypes + + parent_otel_span: Final = _get_parent_otel_span_from_logging_obj(logger) + try: + service_logger: Final = _get_service_logger() + _ = asyncio.create_task( + service_logger.async_service_success_hook( # pyright: ignore[reportUnknownMemberType] # the hook's untyped event_metadata dict + service=ServiceTypes.LITELLM, + duration=end - start, + call_type="post <- track_llm_api_timing", + start_time=start, + end_time=end, + parent_otel_span=parent_otel_span, + ) + ) + except Exception as e: # noqa: BLE001 # the timed post logs and drops a service logging failure + verbose_logger.debug("Error in service logging: %s", e) + + +def times_provider_exchange(client: object, timeout: object) -> bool: + from litellm.rust_bridge.host import http_transport + + return not http_transport.sends_through_python(client, timeout) def defers_async_logging(logger: LoggingSurface) -> bool: @@ -369,6 +476,16 @@ def async_success_handler( return logger.async_success_handler(response, start, end, cache_hit=_cache_hit(logger)) +def cache_hit_success(logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime) -> None: + from litellm.caching.caching_handler import LLMCachingHandler + + handler: Final = logger.llm_caching_handler + if isinstance(handler, LLMCachingHandler): + handler.async_log_cache_hit_on_callbacks( + logging_obj=logger, cached_result=response, start_time=start, end_time=end, cache_hit=True + ) + + def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None: from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -414,7 +531,29 @@ def after_deployment_success(kwargs: dict[str, object], response: object, call_t hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook DeploymentSuccessHook, utils.async_post_call_success_deployment_hook ) - return hook(kwargs, response, CallTypes(call_type)) + typed: Final = CallTypes(call_type) + return hook(kwargs, _messages_stream(response) if typed == CallTypes.anthropic_messages else response, typed) + + +def _messages_stream(response: object) -> object: + from litellm.rust_bridge.execution import AsyncStream + + if not isinstance(response, AsyncStream): + return response + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, + ) + + return AnthropicMessagesStreamingResponse( + completion_stream=cast( # cast-ok: the native stream yields the provider's SSE bytes + AsyncIterator[bytes], response + ), + hidden_params=cast( # cast-ok: the stream's stamped hidden params, shared as Python's handler shares them + AnthropicMessagesStreamHiddenParams, + response._hidden_params, # pyright: ignore[reportPrivateUsage] # the native stream's own hidden params + ), + ) def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_type: str) -> Awaitable[None]: @@ -442,17 +581,25 @@ def stamp_stream_head( kwargs: dict[str, object], # mutable-ok: `finalize` takes the call kwargs as a dict start_time: datetime.datetime, end_time: datetime.datetime, + cache_hit: bool, ) -> None: + if kwargs.get("stream") is True: + logger.stream = True + logger.model_call_details["stream"] = True carrier: Final = _HiddenParamsCarrier(head) - finalize(carrier, logger, kwargs, start_time, end_time) + finalize(carrier, logger, kwargs, start_time, end_time, cache_hit) stamped: Final = carrier._hidden_params # pyright: ignore[reportPrivateUsage] # the carrier's own field if stamped is not head and isinstance(stamped, dict): head.update(cast("dict[str, object]", stamped)) # cast-ok: update_response_metadata writes a str-keyed dict -def stream_opened(logger: Logging, head: Mapping[str, object]) -> None: - logger.stream = True - logger.model_call_details["stream"] = True +def stream_opened(logger: Logging, head: Mapping[str, object], keeps_requested: bool = False) -> None: + requested: Final = logger.model_call_details.get("stream") + streamed: Final = requested if keeps_requested and requested else True + logger.stream = cast(bool, streamed) # cast-ok: Python's Messages handler logs the stream flag it sent + logger.model_call_details["stream"] = streamed + if "httpx_response" in logger.model_call_details: + return logger.model_call_details["response_headers"] = _STREAM_HEADERS.validate_python( head.get("additional_headers") or {} ) @@ -486,6 +633,12 @@ def stream_success( if first_chunk is not None: logger.completion_start_time = first_chunk logger.model_call_details["completion_start_time"] = first_chunk + superseded_error_event: Final = PassThroughStreamingHandler._error_frame_closed_at( # pyright: ignore[reportPrivateUsage] # the Messages stream's own close-at-frame check + EndpointType(endpoint_type), chunks + ) + if superseded_error_event is not None: + _log_superseded_stream_failure(logger, endpoint_type, request_body, chunks, superseded_error_event) + return build: Final = cast( # cast-ok: bounded adapter for the untyped pass-through logging builder StreamingLogBuilder, PassThroughStreamingHandler._route_streaming_logging_to_handler, # pyright: ignore[reportPrivateUsage] # the Messages stream iterator bills through the same builder @@ -518,6 +671,37 @@ def stream_success( enqueue_logging(coroutine) +def _log_superseded_stream_failure( + logger: Logging, + endpoint_type: str, + request_body: Mapping[str, object], + chunks: Sequence[bytes], + error_event: tuple[str, str, int], +) -> None: + from litellm._logging import verbose_proxy_logger + from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + error_type, message, status_code = error_event + PassThroughStreamingHandler._record_partial_usage_for_failure( # pyright: ignore[reportPrivateUsage] # the superseded stream's own partial usage record + litellm_logging_obj=logger, + endpoint_type=EndpointType(endpoint_type), + request_body=dict(request_body), + raw_bytes=tuple(chunks), + stream_context=None, + ) + try: + enqueue_logging( + logger.dispatch_attempt_failure_handlers( # pyright: ignore[reportArgumentType] # the dispatch snapshots now, its coroutine never receives a sent value + anthropic_error_frame_exception(error_type, message, status_code, str(logger.model)), + traceback.format_exc(), + ) + ) + except Exception as e: # noqa: BLE001 # the stream's failure scheduling logs and drops its own failure + verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e) + + def stream_failure( logger: Logging, endpoint_type: str, diff --git a/litellm/rust_bridge/host/http_transport.py b/litellm/rust_bridge/host/http_transport.py index 1bf0f890924..6155cfce2ea 100644 --- a/litellm/rust_bridge/host/http_transport.py +++ b/litellm/rust_bridge/host/http_transport.py @@ -1,7 +1,32 @@ from __future__ import annotations +from typing import Final + def uses_aiohttp_transport() -> bool: from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler return AsyncHTTPHandler.should_use_aiohttp_transport() + + +def _non_positive_timeout(timeout: object) -> bool: + import httpx + + seconds: Final = ( + (timeout.connect, timeout.read, timeout.write, timeout.pool) + if isinstance(timeout, httpx.Timeout) + else (timeout,) + ) + return any(isinstance(value, int | float) and value <= 0 for value in seconds) + + +def sends_through_python(client: object, timeout: object = None) -> bool: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, rsa_key_exchange_and_chacha20_excluded + from litellm.rust_bridge import detectors + + return ( + isinstance(client, AsyncHTTPHandler) + or detectors.unsupported_tls() + or rsa_key_exchange_and_chacha20_excluded() + or _non_positive_timeout(timeout) + ) diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c6568264edb..48de9b0ac03 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -1,39 +1,99 @@ from __future__ import annotations +import asyncio import datetime import errno +import functools +import inspect import ipaddress +import json import os import socket -from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict, dataclass +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence +from dataclasses import asdict, dataclass, replace from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payloads to public TypedDicts import httpx from aiohttp.http_exceptions import ContentLengthError, TransferEncodingError -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError import litellm +from litellm.caching.caching_handler import LLMCachingHandler +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import exception_mapping_utils -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + drops_unsupported_params, + normalize_drop_params, + proxy_request_headers, + stream_cost_injection_requested, +) +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.model_latest_aliases import is_model_latest_alias +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout -from litellm.llms.anthropic.common_utils import AnthropicError, anthropic_pass_through_api_base +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.anthropic.common_utils import ( + AnthropicError, + AnthropicModelInfo, + anthropic_pass_through_api_base, + strip_provider_specific_fields_from_anthropic_messages, +) from litellm.llms.anthropic.messages.native import handler +from litellm.llms.anthropic.messages.native import utils as native_utils +from litellm.llms.anthropic.messages.native.handler import ( + _bridge_target as _bridged, # pyright: ignore[reportPrivateUsage] # the handler's own bridge target +) +from litellm.llms.anthropic.messages.native.handler import ( + _BridgeTarget, # pyright: ignore[reportPrivateUsage] # the handler's own bridge target +) from litellm.llms.anthropic.messages.native.handler import ( _native_messages_config as native_messages_config, # pyright: ignore[reportPrivateUsage] # the handler's own config lookup ) +from litellm.llms.anthropic.messages.native.handler import ( + _resolve_requested_provider as _requested_provider, # pyright: ignore[reportPrivateUsage] # the handler's own provider resolution +) +from litellm.llms.anthropic.messages.native.handler import ( + _sets_own_entra_credential as _own_entra_credential, # pyright: ignore[reportPrivateUsage] # the handler's own credential check +) from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) +from litellm.llms.anthropic.messages.native.interceptors import MessagesInterceptor, get_messages_interceptors +from litellm.llms.anthropic.messages.native.utils import normalize_native_messages, prepare_native_messages +from litellm.llms.anthropic.messages.safeguards import ( + SafeguardsEvaluator, + StreamedSafeguardResults, + native_route_classifies, + with_safeguard_results, +) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled +from litellm.llms.anthropic.wif import get_anthropic_wif_token +from litellm.llms.base_llm.anthropic_messages.transformation import BaseAnthropicMessagesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.custom_httpx.http_handler import get_ssl_configuration -from litellm.rust_bridge import failures -from litellm.rust_bridge.host import callbacks_legacy_python, http_transport -from litellm.rust_bridge.public_call import known_params, optional_str +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params dict is untyped + get_ssl_configuration, +) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import ( + _merge_messages_header_sources as _merged_headers, # pyright: ignore[reportPrivateUsage] # the handler's own header merge +) +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway +from litellm.rust_bridge import detectors, execution, failures +from litellm.rust_bridge.execution import AsyncStream +from litellm.rust_bridge.host import cache, callbacks_legacy_python, http_transport +from litellm.rust_bridge.messages.entrypoints import MessagesStreamHead +from litellm.rust_bridge.public_call import NativeCall, known_params, optional_sequence, optional_str +from litellm.types.integrations.custom_logger import WEBSEARCH_CONVERTED_STREAM_KEY from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CallTypes, ProviderSpecificHeader, anthropic_wif_litellm_params +from litellm.utils import clamp_max_tokens _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( @@ -61,9 +121,12 @@ _BODY_FIELDS: Final = frozenset( "reasoning_effort", "safeguards", "service_tier", + "thread", } ) +WORKLOAD_IDENTITY_FIELDS: Final = anthropic_wif_litellm_params + @dataclass(frozen=True, slots=True) class MessagesSettings: @@ -71,6 +134,15 @@ class MessagesSettings: reasoning_auto_summary: bool additional_drop_params: Sequence[str] request_timeout: float | None + drops_unsupported_params: bool + model: str | None = None + mock_response: Mapping[str, object] | None = None + stream_cost_model: str | None = None + extra_headers: Mapping[str, object] | None = None + mcp_gateway: bool = False + ascii_response_headers: bool = False + joined_reads: bool = False + ascii_request_headers: bool = True def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: @@ -90,13 +162,37 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: return provider -def body_params(request: Mapping[str, object]) -> Mapping[str, object]: +def modify_max_tokens( + model: object, + args: tuple[object, ...], + kwargs: dict[str, object], # mutable-ok: the call's keyword view, rewritten in place as `@client` does +) -> None: + unclamped: Final = kwargs.get("max_tokens") + clamp_max_tokens(model, args, kwargs, CallTypes.anthropic_messages.value) + if kwargs.get("max_tokens") != unclamped: + cache.remember_unclamped(kwargs.get("litellm_logging_obj"), unclamped) + + +def _body_fields(request: Mapping[str, object]) -> Mapping[str, object]: params: Final = known_params(_BODY_FIELDS.difference(handler.additional_drop_paths(request)), request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: return params return _with_deployment_effort(params, request) +def body_params(request: Mapping[str, object], asynchronous: bool, hooked: bool = False) -> Mapping[str, object]: + fields: Final = _body_fields(request) + if not detectors.cache_control_injection(NativeCall(args=(), kwargs=request, base={})): + return {**fields, "messages": request["messages"]} if hooked else fields + messages, system, tools = _prompt_cached(request, asynchronous, hooked) + return { + **fields, + "messages": messages, + **({"system": system} if "system" in fields else {}), + **({"tools": tools} if tools is not None and "tools" in fields else {}), + } + + def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str, object]) -> Mapping[str, object]: """Python's handler rounds the requested effort to the levels the deployment declares in `model_info.reasoning_effort_levels`, the one place a deployment, not the model, decides.""" @@ -116,12 +212,309 @@ def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str, ) +def cached_function() -> Callable[..., object]: + return cast( # cast-ok: the @client cache read names and checks the legacy handler it wraps + "Callable[..., object]", + handler.anthropic_messages, # noqa: TID251 # the @client cache read names and checks the legacy handler it wraps + ) + + +_ENTRY: Final = cast( # cast-ok: only the untyped handler's signature is read + Callable[..., object], + handler.anthropic_messages, # noqa: TID251 # reads the handler's signature, never calls it +) +_HANDLER: Final = cast( # cast-ok: only the untyped handler's signature is read + Callable[..., object], + handler.anthropic_messages_handler, # noqa: TID251 # reads the handler's signature, never calls it +) +_ENTRY_PARAMS: Final = frozenset(inspect.signature(_ENTRY).parameters) - {"kwargs"} +_HANDLER_PARAMS: Final = frozenset(inspect.signature(_HANDLER).parameters) - {"kwargs"} +_JsonBlocks: TypeAlias = list[dict[str, JsonValue]] +_Prepared: TypeAlias = tuple[_JsonBlocks, str | _JsonBlocks | None, Sequence[object] | None] + + +def _prompt_cached(request: Mapping[str, object], asynchronous: bool, hooked: bool) -> _Prepared: + model: Final = str(request["model"]) + provider: Final = optional_str(request.get("custom_llm_provider")) + api_base: Final = optional_str(request.get("api_base")) + tools: Final = cast(_JsonBlocks | None, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + messages: Final = cast(_JsonBlocks, request["messages"]) # cast-ok: handed on unchecked, as the handler does + system: Final = cast( # cast-ok: handed on unchecked, as the handler does + str | _JsonBlocks | None, request.get("system") + ) + if not asynchronous: + return _handler_prepared( + request, + messages, + system, + model=model, + provider=provider, + tools=tools, + api_base=api_base, + presanitized=False, + ) + entry_kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + if hooked: + return _handler_prepared( + entry_kwargs, + messages, + system, + model=model, + provider=handler.provider_after_pre_request_hooks(model, provider, entry_kwargs), + tools=tools, + api_base=api_base, + presanitized=True, + ) + entry_messages, entry_system = prepare_native_messages( + messages, system, entry_kwargs, model=model, custom_llm_provider=provider, tools=tools, api_base=api_base + ) + return _handler_prepared( + entry_kwargs, + entry_messages, + entry_system, + model=model, + provider=handler.provider_after_pre_request_hooks(model, provider, entry_kwargs), + tools=tools, + api_base=api_base, + presanitized=True, + ) + + +def _handler_prepared( + call: Mapping[str, object], + messages: _JsonBlocks, + system: str | _JsonBlocks | None, + *, + model: str, + provider: str | None, + tools: _JsonBlocks | None, + api_base: str | None, + presanitized: bool, +) -> _Prepared: + handler_kwargs: Final = {name: value for name, value in call.items() if name not in _HANDLER_PARAMS} + prepared_messages, prepared_system = prepare_native_messages( + messages, + system, + handler_kwargs, + model=model, + custom_llm_provider=provider, + tools=tools, + api_base=api_base, + presanitized=presanitized, + ) + return prepared_messages, prepared_system, handler.tools_marked_for_cache_points(tools, handler_kwargs) + + +_execute_pre_request_hooks: Final = cast( # cast-ok: the handler's untyped pre-request hook loop + Callable[..., Awaitable[dict[str, object]]], + handler._execute_pre_request_hooks, # pyright: ignore[reportPrivateUsage] # the handler's own pre-request hook loop +) +_HOOKED_PARAMS: Final = ( + "tools", + "stream", + "metadata", + "stop_sequences", + "system", + "temperature", + "thinking", + "tool_choice", + "top_k", + "top_p", +) + + +def runs_pre_request_hooks() -> bool: + return detectors.overridden("async_pre_request_hook", litellm.callbacks) or any( + isinstance(callback, CustomLogger) and "async_pre_request_hook" in vars(callback) + for callback in litellm.callbacks + ) + + +def _unsupported(provider: str | None, model: str, unsupported: str) -> litellm.UnsupportedParamsError: + return litellm.UnsupportedParamsError( + message=( + f"provider {provider} on /v1/messages runs only on the Rust path, which does not support {unsupported}" + ), + llm_provider=provider, + model=model, + ) + + +async def pre_request_hooks(request: Mapping[str, object]) -> Mapping[str, object]: + model: Final = str(request["model"]) + provider: Final = optional_str(request.get("custom_llm_provider")) + tools: Final = cast(_JsonBlocks | None, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + messages, system = prepare_native_messages( + cast(_JsonBlocks, request["messages"]), # cast-ok: handed on unchecked, as the handler does + cast(str | _JsonBlocks | None, request.get("system")), # cast-ok: handed on unchecked, as the handler does + kwargs, + model=model, + custom_llm_provider=provider, + tools=tools, + api_base=optional_str(request.get("api_base")), + ) + named: Final = {**{name: request.get(name) for name in _ENTRY_PARAMS}, "messages": messages, "system": system} + hooked: Final = await _execute_pre_request_hooks( + model=model, + messages=messages, + tools=tools, + stream=cast(bool | None, named["stream"]), # cast-ok: handed on unchecked, as the handler does + custom_llm_provider=provider, + tool_choice=named["tool_choice"], + **kwargs, + ) + for name in _HOOKED_PARAMS: + named[name] = hooked.pop(name, named[name]) + hooked_provider: Final = handler.provider_after_pre_request_hooks(model, provider, hooked) + hooked.pop("litellm_params", None) + kwargs.update(hooked) + requested_provider: Final = handler.provider_after_pre_request_hooks(model, provider, {}) + functools.partial( + _HANDLER, + **{name: value for name, value in named.items() if name in _HANDLER_PARAMS - {"custom_llm_provider"}}, + custom_llm_provider=hooked_provider, + _litellm_messages_presanitized=True, + **kwargs, + ) + _python_handler_params({**named, **kwargs}, hooked_provider) + if hooked_provider != requested_provider: + raise _unsupported(requested_provider, model, f"a pre-request hook routing the call to {hooked_provider}") + return {**named, **kwargs} + + +def _intercepted_provider(request: Mapping[str, object]) -> str | None: + return handler.provider_after_pre_request_hooks( + str(request["model"]), optional_str(request.get("custom_llm_provider")), {} + ) + + +def has_interceptors() -> bool: + return bool(get_messages_interceptors()) + + +def short_circuits() -> bool: + from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger + + stock: Final[object] = WebSearchInterceptionLogger.try_short_circuit_search # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the logger annotates bare dicts + return any( + isinstance(callback, WebSearchInterceptionLogger) and type(callback).try_short_circuit_search is not stock # pyright: ignore[reportUnknownMemberType] # the logger annotates bare dicts + for callback in litellm.callbacks + ) + + +def interceptor(request: Mapping[str, object]) -> MessagesInterceptor | None: + tools: Final = cast(_JsonBlocks | None, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + provider: Final = _intercepted_provider(request) + return next((owner for owner in get_messages_interceptors() if owner.can_handle(tools, provider)), None) + + +def _entry_messages( + request: Mapping[str, object], + kwargs: dict[str, object], # mutable-ok: the call's own kwargs, rewritten in place by the handler's normalizers + hooked: bool, +) -> tuple[_JsonBlocks, str | _JsonBlocks | None]: + messages: Final = cast(_JsonBlocks, request["messages"]) # cast-ok: handed on unchecked, as the handler does + system: Final = cast( # cast-ok: handed on unchecked, as the handler does + str | _JsonBlocks | None, request.get("system") + ) + if hooked: + return messages, system + return prepare_native_messages( + messages, + system, + kwargs, + model=str(request["model"]), + custom_llm_provider=optional_str(request.get("custom_llm_provider")), + tools=cast(_JsonBlocks | None, request.get("tools")), # cast-ok: handed on unchecked, as the handler does + api_base=optional_str(request.get("api_base")), + ) + + +def short_circuit(request: Mapping[str, object], stream: bool, hooked: bool) -> Awaitable[object]: + kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + messages, _ = _entry_messages(request, kwargs, hooked) + return handler._try_websearch_short_circuit( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType] # the handler's own short-circuit, which annotates bare dicts + model=str(request["model"]), + messages=messages, + tools=cast(_JsonBlocks | None, request.get("tools")), # cast-ok: handed on unchecked, as the handler does + custom_llm_provider=_intercepted_provider(request), + stream=stream or bool(kwargs.get(WEBSEARCH_CONVERTED_STREAM_KEY, False)), + kwargs={**kwargs, "metadata": request.get("metadata")}, + ) + + +def answer(owner: MessagesInterceptor, request: Mapping[str, object], stream: bool, hooked: bool) -> Awaitable[object]: + model: Final = str(request["model"]) + tools: Final = cast(_JsonBlocks | None, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + messages, system = _entry_messages(request, kwargs, hooked) + handle: Final = cast( # cast-ok: the interceptor answers with whatever the caller receives + Callable[..., Awaitable[object]], owner.handle + ) + return handle( + model=model, + messages=messages, + tools=tools, + stream=stream or kwargs.get("_websearch_interception_converted_stream", False), + max_tokens=cast(int, request["max_tokens"]), # cast-ok: admission requires an int max_tokens + custom_llm_provider=_intercepted_provider(request), + api_key=request.get("api_key"), + api_base=request.get("api_base"), + metadata=request.get("metadata"), + stop_sequences=request.get("stop_sequences"), + system=system, + temperature=request.get("temperature"), + thinking=request.get("thinking"), + tool_choice=request.get("tool_choice"), + top_k=request.get("top_k"), + top_p=request.get("top_p"), + **kwargs, + ) + + def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: from litellm.llms.anthropic.messages.native.streaming_iterator import ( anthropic_messages_stream_hidden_params, ) - return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) + return MessagesStreamHead(anthropic_messages_stream_hidden_params(callbacks_legacy_python.wire_headers(headers))) + + +async def _unrequested_stream_delivered( + stream: object, + kwargs: dict[str, object], # mutable-ok: the call's own kwargs, as `@client` hands them on +) -> object: + try: + delivered: Final = await callbacks_legacy_python.after_deployment_success( + kwargs, stream, CallTypes.anthropic_messages.value + ) + function: Final = cached_function() + caching: Final = LLMCachingHandler( + original_function=function, + request_kwargs=kwargs, + start_time=datetime.datetime.now(), # noqa: DTZ005 # @client's naive start time + ) + store: Final = cast( # cast-ok: the handler leaves its original_function untyped + Callable[..., Awaitable[None]], caching.async_set_cache + ) + await store(result=delivered, original_function=function, kwargs=kwargs) + except BaseException: + from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_if_supported + + await aclose_if_supported(stream) + raise + return delivered + + +def unrequested_stream( + head: object, + kwargs: dict[str, object], # mutable-ok: the call's own kwargs, as `@client` hands them on +) -> None: + if not isinstance(head, MessagesStreamHead): + return + deliver: Final = functools.partial(_unrequested_stream_delivered, kwargs=kwargs) + head.deliver = deliver # rebind-ok: the head carries the delivery back _UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) @@ -129,13 +522,17 @@ _UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) _BODY_CUT: Final = TypeAdapter(tuple[int, int | None, bool]) _ERROR_CLASS: TypeAlias = Callable[..., Exception] _validate_metadata: Final = cast( # cast-ok: the handler's untyped metadata validator - Callable[[dict[str, object]], object], handler.validate_anthropic_api_metadata + Callable[[object], object], handler.validate_anthropic_api_metadata +) +_mock_response: Final = cast( # cast-ok: the handler's mock builder takes untyped messages + Callable[..., Mapping[str, object]], native_utils.mock_response ) _exception_type: Final = cast( # cast-ok: the legacy public exception mapper is untyped Callable[..., Exception], exception_mapping_utils.exception_type ) _BODY_READ_TIMEOUT: Final = "Timeout on reading data from socket" # aiohttp's SocketTimeoutError text _UNREACHABLE: Final = "could not reach the provider: " +_TOO_MANY_REDIRECTS: Final = "upstream network error: error following redirect: too many redirects" _SEND_FAILED: Final = "upstream network error: error sending request: client error (SendRequest): " _DNS_FAILURE: Final = "failed to lookup address information: " _HTTPX_DNS_ERRNO: Final = { @@ -161,6 +558,12 @@ def _rejection(error: Exception, request: Mapping[str, object], request_provider _validate_metadata(_METADATA.validate_python(metadata)) except ValidationError as invalid: return invalid + provider_rejection: Final = getattr(error, "messages_provider_rejection", None) + if isinstance(provider_rejection, str): + _log_resolved_call(request, request_provider) + return _public_provider_error( + _config_error(400, provider_rejection, (), request, request_provider), error, request, request_provider + ) if getattr(error, "messages_unsupported_param", False): return litellm.UnsupportedParamsError(message=str(error)) return litellm.BadRequestError( @@ -251,7 +654,12 @@ def _aiohttp_connect_failure(url: httpx.URL, port: int, request: Mapping[str, ob def _transport_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> str | None: """The text Python's HTTP client raises for a call that failed before any response arrived.""" + body_cut: Final = getattr(error, "messages_body_cut", None) + if body_cut is not None: + return str(_body_cut(*_BODY_CUT.validate_python(body_cut))) native: Final = str(error) + if native == _TOO_MANY_REDIRECTS: + return "Exceeded maximum allowed redirects." if native.startswith(_SEND_FAILED): reset: Final = native.endswith(f"Connection reset by peer (os error {errno.ECONNRESET})") closed: Final = native.endswith("connection closed before message completed") @@ -301,6 +709,20 @@ def _provider_error(error: Exception, request: Mapping[str, object], request_pro if failure is None: return None status, text, headers = failure + provider_error: Final = _config_error(status, text, headers, request, request_provider) + if transport is not None and isinstance(provider_error, BaseLLMException): + # Python's `_handle_error` gives a failure without a response the status 500. + provider_error.status_code_is_synthesized = True + return provider_error + + +def _config_error( + status: int, + text: str, + headers: Sequence[tuple[str, str]], + request: Mapping[str, object], + request_provider: str, +) -> Exception: api_base: Final = request.get("api_base") config: Final = native_messages_config( str(request["model"]).removeprefix(f"{request_provider}/"), @@ -308,17 +730,26 @@ def _provider_error(error: Exception, request: Mapping[str, object], request_pro request.get("model_info"), api_base if isinstance(api_base, str) else None, ) - provider_error: Final = ( - cast(_ERROR_CLASS, config.get_error_class)( # cast-ok: the config's untyped headers parameter - error_message=text, status_code=status, headers=dict(headers) - ) - if config is not None - else AnthropicError(status_code=status, message=text, headers=httpx.Headers(headers)) + if config is None: + return AnthropicError(status_code=status, message=text, headers=httpx.Headers(headers)) + return cast(_ERROR_CLASS, config.get_error_class)( # cast-ok: the config's untyped headers parameter + error_message=text, status_code=status, headers=dict(headers) ) - if transport is not None and isinstance(provider_error, BaseLLMException): - # Python's `_handle_error` gives a failure without a response the status 500. - provider_error.status_code_is_synthesized = True - return provider_error + + +def _public_provider_error( + provider_error: Exception, error: Exception, request: Mapping[str, object], request_provider: str +) -> Exception: + try: + return _exception_type( + model=str(request["model"]), + custom_llm_provider=request_provider, + original_exception=provider_error, + extra_kwargs=dict(request), + ) + except Exception as mapped: # noqa: BLE001 # exception_type raises the mapped public error + mapped.__context__ = error + return mapped def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> None: @@ -331,20 +762,22 @@ def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> logger, kwargs={name: value for name, value in request.items() if name != "metadata"}, model=str(request["model"]).removeprefix(f"{request_provider}/"), - optional_params=dict(body_params(request)), - litellm_params={"model_info": request.get("model_info"), **body_params(request)}, + optional_params=dict(_body_fields(request)), + litellm_params={"model_info": request.get("model_info"), **_body_fields(request)}, custom_llm_provider=request_provider, ) -def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: +def map_failure( + error: Exception, request: Mapping[str, object], request_provider: str, asynchronous: bool = False +) -> Exception: """Raises what Python's Messages handler raises for the same failure.""" if getattr(error, "messages_timeout", None) == "Body": # Python reads a stream's body outside its HTTP client's error mapping, so the transport's # own timeout reaches the caller unmapped. - return httpx.ReadTimeout(_BODY_READ_TIMEOUT) + return httpx.ReadTimeout(_BODY_READ_TIMEOUT if http_transport.uses_aiohttp_transport() else "") body_cut: Final = getattr(error, "messages_body_cut", None) - if body_cut is not None: + if body_cut is not None and request.get("stream") is True: # Python reads a stream's body outside its HTTP client's error mapping, so a body the # connection cut short reaches the caller as its transport raised it. return _body_cut(*_BODY_CUT.validate_python(body_cut)) @@ -356,6 +789,11 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide error.__context__ = chained # rebind-ok: Python chains the transport timeout as the exception context if getattr(error, "messages_late_rejection", False): _log_resolved_call(request, request_provider) + if getattr(error, "messages_python_raised", False): + # Python's request transform raises this builtin exception unmapped, e.g. the TypeError + # of comparing a string max_tokens while capping the thinking budget. + _flag_messages_call(request.get("litellm_logging_obj"), asynchronous) + return error if getattr(error, "messages_request_error", False): return _rejection(error, request, request_provider) if getattr(error, "messages_missing_credential", False): @@ -369,28 +807,759 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide return failures.map_native_failure( error, str(request["model"]), request_provider, request, optional_str(request.get("api_base")) ) - try: - return _exception_type( - model=str(request["model"]), - custom_llm_provider=request_provider, - original_exception=provider_error, - extra_kwargs=dict(request), - ) - except Exception as mapped: # noqa: BLE001 # exception_type raises the mapped public error - mapped.__context__ = error - return mapped + return _public_provider_error(provider_error, error, request, request_provider) def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True -def settings(kwargs: Mapping[str, object]) -> dict[str, object]: +def _python_handler_params(request: Mapping[str, object], custom_llm_provider: object) -> GenericLiteLLMParams: + return GenericLiteLLMParams( + **{name: value for name, value in request.items() if name not in _HANDLER_PARAMS}, + api_key=request.get("api_key"), + api_base=request.get("api_base"), + custom_llm_provider=custom_llm_provider, + ) + + +def _resolved_provider(request: Mapping[str, object]) -> str | None: + return optional_str(request.get("custom_llm_provider")) or inferred_provider( + str(request["model"]), optional_str(request.get("api_base")) + ) + + +def _scope_provider_specific_header(request: Mapping[str, object]) -> None: + scoped: Final = request.get("provider_specific_header") + if scoped is None: + return + ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=cast( # cast-ok: the handler hands the caller's value on unvalidated + ProviderSpecificHeader | None, scoped + ), + custom_llm_provider=_resolved_provider(request), + native_anthropic_messages=True, + ) + + +def _python_messages_url(api_base: str | None) -> str: + """`AnthropicMessagesConfig.get_complete_url`, untrimmed, whose `urlsplit` raises a bare + ValueError for a base it cannot split (``Invalid IPv6 URL``).""" + return ensure_api_route(AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com", "v1/messages") + + +def _raise_as_python_does(request: Mapping[str, object], asynchronous: bool) -> None: + iter(cast(Iterable[object], request.get("messages"))) # cast-ok: iterated unchecked, as Python's handler does + tools: Final = request.get("tools") + if tools is not None and not isinstance(tools, list | tuple) and has_interceptors(): + interceptor(request) + _python_handler_params(request, request.get("custom_llm_provider")) + # Python's handler scopes the caller's provider_specific_header before it builds any header, so + # a value of the wrong shape raises its own error (a str: "'str' object has no attribute 'get'"), + # and then builds the URL, unmapped. + logger: Final = request.get("litellm_logging_obj") + try: + _scope_provider_specific_header(request) + except Exception: + _flag_messages_call(logger, asynchronous) + raise + api_base: Final = request.get("api_base") + if not isinstance(api_base, str) or _resolved_provider(request) != "anthropic": + return + try: + _python_messages_url(api_base) + except ValueError: + _log_resolved_call(request, "anthropic") + _flag_messages_call(logger, asynchronous) + raise + + +def settings( + kwargs: Mapping[str, object], request: Mapping[str, object] | None = None, asynchronous: bool = False +) -> dict[str, object]: + if request is not None: + _raise_as_python_does(request, asynchronous) + aiohttp: Final = http_transport.uses_aiohttp_transport() + configured: Final = MessagesSettings( + drop_params=_drop_params(kwargs), + reasoning_auto_summary=is_reasoning_auto_summary_enabled(), + additional_drop_params=handler.additional_drop_paths(kwargs), + request_timeout=get_configured_request_timeout(), + drops_unsupported_params=drops_unsupported_params(kwargs.get("drop_params")), + ascii_response_headers=aiohttp, + joined_reads=aiohttp, + ) return asdict( - MessagesSettings( - drop_params=_drop_params(kwargs), - reasoning_auto_summary=is_reasoning_auto_summary_enabled(), - additional_drop_params=handler.additional_drop_paths(kwargs), - request_timeout=get_configured_request_timeout(), + configured if request is None else _with_call_settings(configured, request, kwargs.get("litellm_logging_obj")) + ) + + +def _flag_messages_call(logger: object, asynchronous: bool) -> None: + if not isinstance(logger, LiteLLMLoggingObj): + return + logger.model_call_details.setdefault("litellm_params", {})[CallTypes.aanthropic_messages.value] = asynchronous + + +@dataclass(frozen=True, slots=True) +class _Answered: + value: object + + def send(self, _value: object) -> execution.Complete: + return execution.Complete(self.value) + + def throw(self, error: BaseException) -> execution.Step: + raise error + + def close(self) -> None: + return None + + +def sync_answer(request: Mapping[str, object]) -> _Answered | None: + logger: Final = request.get("litellm_logging_obj") + mock: Final = _mock(request, logger, asynchronous=False) + if mock is not None: + _flag_messages_call(logger, asynchronous=False) + return _Answered(response(mock)) + served: Final = gateway(request, asynchronous=False) + return None if served is None else _Answered(served) + + +def _handler_params(request: Mapping[str, object]) -> GenericLiteLLMParams: + return GenericLiteLLMParams( + **{name: value for name, value in request.items() if name not in _HANDLER_PARAMS}, + api_key=optional_str(request.get("api_key")), + api_base=optional_str(request.get("api_base")), + custom_llm_provider=optional_str(request.get("custom_llm_provider")), + ) + + +def _requested_model(request: Mapping[str, object]) -> str: + model, _, _, _ = _requested_provider( + str(request["model"]), optional_str(request.get("custom_llm_provider")), _handler_params(request) + ) + return model + + +def _model_without_latest_alias(request: Mapping[str, object], model: str, custom_llm_provider: str | None) -> str: + if not is_model_latest_alias(model, custom_llm_provider): + return model + resolved, provider, _, _ = _requested_provider(model, custom_llm_provider, _handler_params(request)) + prefix, separator, _ = model.partition("/") + return f"{prefix}/{resolved}" if separator and prefix == provider else resolved + + +def _mock(request: Mapping[str, object], logger: object, asynchronous: bool) -> Mapping[str, object] | None: + mock: Final = request.get("mock_response") + if not mock or not isinstance(mock, str): + return None + _validate_metadata(request.get("metadata")) + messages: Final = request.get("messages") + max_tokens: Final = request.get("max_tokens") + model: Final = _requested_model(request) + try: + return dict( + _mock_response( + model=model, + messages=messages if isinstance(messages, list) else [], + max_tokens=max_tokens if isinstance(max_tokens, int) else 0, + mock_response=mock, + ) + ) + except Exception: + _flag_messages_call(logger, asynchronous) + raise + + +def _stream_cost_model(request: Mapping[str, object]) -> str | None: + if request.get("stream") is not True or not stream_cost_injection_requested(proxy_request_headers(request)): + return None + return _requested_model(request) + + +def _affinity_extra_headers( + request: Mapping[str, object], model: str, explicit_provider: str | None +) -> Mapping[str, object] | None: + api_base: Final = optional_str(request.get("api_base")) + anthropic_header: Final = get_provider_affinity_header(AnthropicModelInfo.get_api_base(api_base)) + other_header: Final = get_provider_affinity_header(api_base) + if anthropic_header is None and other_header is None: + return None + provider: Final = explicit_provider or inferred_provider(model, api_base) + default_header: Final = anthropic_header if provider == "anthropic" else other_header + if default_header is None: + return None + headers: Final = as_str_mapping(request.get("headers")) + extra_headers: Final = as_str_mapping(request.get("extra_headers")) + merged: Final = _merged_headers( + headers, + extra_headers, + ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=cast( # cast-ok: the handler hands the caller's value on unvalidated + ProviderSpecificHeader | None, request.get("provider_specific_header") + ), + custom_llm_provider=provider, + native_anthropic_messages=True, + ), + ) + added: Final = add_provider_affinity_header(merged, _handler_params(request), default_header) + affinity: Final = {name: value for name, value in added.items() if name not in merged} + if not affinity: + return None + return dict(_merged_headers(headers, extra_headers, affinity)) + + +def _with_call_settings( + configured: MessagesSettings, request: Mapping[str, object], logger: object +) -> MessagesSettings: + explicit_provider: Final = optional_str(request.get("custom_llm_provider")) + model: Final = _model_without_latest_alias(request, str(request["model"]), explicit_provider) + mock: Final = _mock(request, logger, asynchronous=True) + if mock is not None: + return replace(configured, model=model, mock_response=mock) + return replace( + configured, + model=model, + stream_cost_model=_stream_cost_model(request), + extra_headers=_affinity_extra_headers(request, model, explicit_provider), + mcp_gateway=_takes_mcp_gateway(request), + ) + + +def _takes_mcp_gateway(request: Mapping[str, object]) -> bool: + tools: Final = optional_sequence(request.get("tools")) + return not request.get("_skip_mcp_handler", False) and bool(tools) and may_name_litellm_mcp_gateway(tools) + + +MessagesResult: TypeAlias = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object] + + +def gateway( + request: Mapping[str, object], asynchronous: bool, hooked: bool = False +) -> Coroutine[object, object, MessagesResult] | None: + if not _takes_mcp_gateway(request): + return None + model: Final = str(request["model"]) + provider: Final = optional_str(request.get("custom_llm_provider")) + api_base: Final = optional_str(request.get("api_base")) + tools: Final = cast(_JsonBlocks, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + messages: Final = cast(_JsonBlocks, request["messages"]) # cast-ok: handed on unchecked, as the handler does + system: Final = cast( # cast-ok: handed on unchecked, as the handler does + str | _JsonBlocks | None, request.get("system") + ) + if not asynchronous: + return _gateway(request, messages, system, model=model, provider=provider, tools=tools, presanitized=False) + entry_kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + entry_messages, entry_system = ( + (messages, system) + if hooked + else prepare_native_messages( + messages, system, entry_kwargs, model=model, custom_llm_provider=provider, tools=tools, api_base=api_base ) ) + return _gateway( + {**request, **entry_kwargs, "is_async": True}, + entry_messages, + entry_system, + model=model, + provider=handler.provider_after_pre_request_hooks(model, provider, entry_kwargs), + tools=tools, + presanitized=True, + ) + + +def _gateway( + call: Mapping[str, object], + messages: _JsonBlocks, + system: str | _JsonBlocks | None, + *, + model: str, + provider: str | None, + tools: _JsonBlocks, + presanitized: bool, +) -> Coroutine[object, object, MessagesResult]: + kwargs: Final = {name: value for name, value in call.items() if name not in _HANDLER_PARAMS} + api_key: Final = optional_str(call.get("api_key")) + api_base: Final = optional_str(call.get("api_base")) + prepared_messages, prepared_system = prepare_native_messages( + messages, + system, + kwargs, + model=model, + custom_llm_provider=provider, + tools=tools, + api_base=api_base, + presanitized=presanitized, + ) + metadata: Final = _validate_metadata(call.get("metadata")) + is_async: Final = kwargs.pop("is_async", False) + logging_obj: Final = kwargs.get("litellm_logging_obj") + litellm_params: Final = GenericLiteLLMParams( + **kwargs, api_key=api_key, api_base=api_base, custom_llm_provider=provider + ) + resolved_model, resolved_provider, dynamic_api_key, dynamic_api_base = _requested_provider( + model, provider, litellm_params + ) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.model_call_details["agentic_loop_params"] = { + "model": model, + "custom_llm_provider": resolved_provider, + } + logging_obj.model_call_details.setdefault("litellm_params", {})[CallTypes.aanthropic_messages.value] = is_async + if kwargs.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + bridge: Final = _bridged( + _BridgeTarget(model=model, api_key=api_key, api_base=api_base), + _BridgeTarget( + model=resolved_model, + api_key=None if _own_entra_credential(litellm_params) else dynamic_api_key, + api_base=dynamic_api_base, + ), + resolved_provider, + ) + kwargs.pop("_skip_mcp_handler", False) + from litellm.llms.anthropic.messages.native.mcp_handler import anthropic_messages_with_mcp + + return anthropic_messages_with_mcp( + max_tokens=cast(int, call["max_tokens"]), # cast-ok: handed on unchecked, as the handler does + messages=prepared_messages, + model=resolved_model, + metadata=metadata, + stop_sequences=call.get("stop_sequences"), + stream=call.get("stream", False), + system=prepared_system, + temperature=call.get("temperature"), + thinking=call.get("thinking"), + tool_choice=call.get("tool_choice"), + tools=tools, + top_k=call.get("top_k"), + top_p=call.get("top_p"), + container=call.get("container"), + api_key=bridge.api_key, + api_base=bridge.api_base, + client=call.get("client"), + custom_llm_provider=resolved_provider, + **kwargs, + ) + + +def inject_stream_cost(frames: bytes, model: str, request: Mapping[str, object]) -> object: + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + logger: Final = request.get("litellm_logging_obj") + logging_obj: Final = logger if isinstance(logger, LiteLLMLoggingObj) else None + return ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + frames, + model, + logging_obj, + request_headers=proxy_request_headers( + as_str_mapping(logging_obj.model_call_details.get("litellm_params")) if logging_obj else None + ), + ) + + +def workload_identity_token(params: Mapping[str, object], api_base: str | None) -> str | None: + litellm_params: Final = GenericLiteLLMParams(**params) + return get_anthropic_wif_token(dict(litellm_params), api_base, "") + + +class RawTransportFailure(Exception): + def __init__(self, error: Exception) -> None: + super().__init__(str(error)) + self.error = error + + +@dataclass(frozen=True, slots=True) +class Sent: + status: int + headers: Sequence[tuple[str, str]] + text: str | None = None + chunks: AsyncIterator[bytes] | None = None + rejection: httpx.HTTPStatusError | None = None + + +def _unsendable_api_base(api_base: object) -> bool: + """A caller api_base the native URL parser rejects or reads otherwise than httpx and aiohttp + (no http(s) scheme, whitespace, control or non-ASCII characters): Python's client decides what + such a URL does, from a relative-URL error to a request on the host it finds.""" + if not isinstance(api_base, str) or not api_base: + return False + return not ( + api_base.isascii() + and api_base.isprintable() + and " " not in api_base + and api_base.lower().startswith(("http://", "https://")) + ) + + +def transport_handler(request: Mapping[str, object]) -> AsyncHTTPHandler | None: + client: Final = request.get("client") + if isinstance(client, AsyncHTTPHandler): + return client + if not http_transport.sends_through_python(client, request.get("timeout")) and not _unsendable_api_base( + request.get("api_base") + ): + return None + return get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) + + +def decoding_failure(message: str, request: Mapping[str, object], request_provider: str) -> Exception: + error: Final = httpx.DecodingError(message) + if request.get("stream") is True: + return error + return map_transport_failure(error, request, request_provider) + + +def unencodable_header(value: bytes, received: bool) -> Exception: + text: Final = value.decode("utf-8", "surrogateescape" if received else "surrogatepass") + try: + text.encode("ascii") + except UnicodeEncodeError as unencodable: + return unencodable + return ValueError(text) + + +def _sent_url(request: Mapping[str, object], url: str) -> str: + """The native URL builder trims the caller's api_base; Python sends it as given.""" + api_base: Final = request.get("api_base") + if not isinstance(api_base, str) or api_base == api_base.strip() or not _unsendable_api_base(api_base): + return url + try: + return _python_messages_url(api_base) if url == _python_messages_url(api_base.strip()) else url + except ValueError: # a base urlsplit cannot split keeps the URL the native builder made + return url + + +async def transport_send( + handler: AsyncHTTPHandler, + request: Mapping[str, object], + url: str, + headers: Sequence[tuple[str, str]], + body: object, + stream: bool, +) -> Sent: + try: + response: Final = await handler.post( # pyright: ignore[reportUnknownMemberType] # the handler's untyped dict params + url=_sent_url(request, url), + headers=dict(headers), + data=body if isinstance(body, bytes) else json.dumps(body), + stream=stream, + logging_obj=cast( # cast-ok: the call's logger, handed on as Python's handler hands it + "LiteLLMLoggingObj | None", request.get("litellm_logging_obj") + ), + timeout=cast( # cast-ok: the call's timeout as the handler resolved it for the post + "float | httpx.Timeout | None", request.get("timeout") + ), + ) + response.raise_for_status() + except httpx.HTTPStatusError as rejection: + try: + text: Final = rejection.response.text + except Exception as unreadable: # noqa: BLE001 # Python's retry check raises it from its except block, unmapped + raise RawTransportFailure(unreadable) from None + return Sent(rejection.response.status_code, rejection.response.headers.multi_items(), text, None, rejection) + if stream: + return Sent(response.status_code, response.headers.multi_items(), chunks=response.aiter_bytes()) + return Sent(response.status_code, response.headers.multi_items(), text=response.text) + + +def map_transport_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: + api_base: Final = request.get("api_base") + config: Final = native_messages_config( + str(request["model"]).removeprefix(f"{request_provider}/"), + request_provider, + request.get("model_info"), + api_base if isinstance(api_base, str) else None, + ) + try: + handler.base_llm_http_handler._handle_error(e=error, provider_config=config) # pyright: ignore[reportPrivateUsage, reportArgumentType] # the handler's own error mapping, which always raises + except Exception as provider_error: # noqa: BLE001 # _handle_error raises the provider exception it builds + provider_error.__context__ = error + if not isinstance(provider_error, BaseLLMException): + return provider_error + try: + return _exception_type( + model=str(request["model"]), + custom_llm_provider=request_provider, + original_exception=provider_error, + extra_kwargs=dict(request), + ) + except Exception as mapped: # noqa: BLE001 # exception_type raises the mapped public error + mapped.__context__ = provider_error + return mapped + return error + + +_AGENTIC_UNSENT_PARAMS: Final = frozenset({"model", "messages", "max_tokens"}) +_HANDLER_POPPED: Final = frozenset( + {"is_async", "_litellm_messages_presanitized", "_skip_mcp_handler", "cache_control_injection_points"} +) + + +def _logger(request: Mapping[str, object]) -> LiteLLMLoggingObj | None: + logger: Final = request.get("litellm_logging_obj") + return logger if isinstance(logger, LiteLLMLoggingObj) else None + + +def runs_agentic_hooks(request: Mapping[str, object]) -> bool: + logger: Final = _logger(request) + if logger is None: + return False + return BaseLLMHTTPHandler._has_agentic_completion_hook(logger) or bool( # pyright: ignore[reportPrivateUsage] # the handler's own agentic gate + request.get(WEBSEARCH_CONVERTED_STREAM_KEY, False) + ) + + +def record_agentic_loop(request: Mapping[str, object], custom_llm_provider: str) -> None: + logger: Final = _logger(request) + if logger is None: + return + logger.model_call_details["agentic_loop_params"] = { + "model": request["model"], + "custom_llm_provider": custom_llm_provider, + } + if request.get(WEBSEARCH_CONVERTED_STREAM_KEY, False): + logger.model_call_details["websearch_interception_converted_stream"] = True + + +@dataclass(frozen=True, slots=True) +class _AgenticCall: + model: str + messages: list[dict[str, object]] # mutable-ok: the handler's agentic hooks take a message list + config: BaseAnthropicMessagesConfig + params: dict[str, object] # mutable-ok: the handler's agentic hooks take a params dict + logger: LiteLLMLoggingObj + provider: str + kwargs: dict[str, object] # mutable-ok: the handler's agentic hooks take a kwargs dict + + +def _agentic_params( + sent: Mapping[str, object], requested: Mapping[str, object], request: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the handler's agentic hooks take a params dict + params: Final = {name: value for name, value in sent.items() if name not in _AGENTIC_UNSENT_PARAMS} + if "tools" not in params: + return params + caller_tools: Final = request.get("tools") + requested_tools: Final = requested.get("tools") + return {**params, "tools": caller_tools if caller_tools == requested_tools else requested_tools} + + +def _agentic_call( + request: Mapping[str, object], + sent: Mapping[str, object], + model: str, + custom_llm_provider: str, + requested: Mapping[str, object], + hooked: bool, +) -> _AgenticCall: + logger: Final = _logger(request) + api_base: Final = optional_str(request.get("api_base")) + config: Final = native_messages_config(model, custom_llm_provider, request.get("model_info"), api_base) + if logger is None or config is None: + raise ValueError("an agentic Messages call needs its logger and provider config") + kwargs: Final = { + name: value for name, value in request.items() if name not in _HANDLER_PARAMS and name not in _HANDLER_POPPED + } + messages, _, _ = _prompt_cached(request, asynchronous=True, hooked=hooked) + return _AgenticCall( + model=model, + messages=cast( # cast-ok: the handler hands its prepared messages to the hooks as dicts + list[dict[str, object]], strip_provider_specific_fields_from_anthropic_messages(messages) + ), + config=config, + params=_agentic_params(sent, requested, request), + logger=logger, + provider=custom_llm_provider, + kwargs=BaseLLMHTTPHandler._agentic_hook_kwargs( # pyright: ignore[reportPrivateUsage] # the handler's own follow-up kwargs + kwargs=kwargs, api_key=optional_str(request.get("api_key")), api_base=api_base + ), + ) + + +async def _failed_deployment(request: Mapping[str, object], error: Exception) -> None: + from litellm.utils import async_post_call_failure_deployment_hook + + try: + await async_post_call_failure_deployment_hook( + request_data=request, exception=error, call_type=CallTypes.anthropic_messages.value + ) + except asyncio.CancelledError: + raise + except BaseException: # noqa: BLE001, S110 # a failing hook never replaces the agentic loop's own failure, as in `@client` + pass + + +async def agentic_response( + response: object, + request: Mapping[str, object], + sent: Mapping[str, object], + model: str, + custom_llm_provider: str, + requested: Mapping[str, object], + hooked: bool, +) -> object: + try: + call: Final = _agentic_call(request, sent, model, custom_llm_provider, requested, hooked) + initial: Final = cast(AnthropicMessagesResponse, response) # cast-ok: the public response the handler finishes + return await handler.base_llm_http_handler._finalize_anthropic_messages_response( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType] # the handler's own agentic finish, typed with bare dicts + initial_response=initial, + model=call.model, + messages=call.messages, + anthropic_messages_provider_config=call.config, + anthropic_messages_optional_request_params=call.params, + logging_obj=call.logger, + custom_llm_provider=call.provider, + kwargs=call.kwargs, + ) + except Exception as error: + await _failed_deployment(request, error) + raise + + +def _agentic_stream(stream: AsyncStream, call: _AgenticCall, held: frozenset[str]) -> object: + from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import AgenticAnthropicStreamingIterator + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, + ) + + return AnthropicMessagesStreamingResponse( + completion_stream=AgenticAnthropicStreamingIterator( + completion_stream=stream, + http_handler=handler.base_llm_http_handler, + model=call.model, + messages=call.messages, + anthropic_messages_provider_config=call.config, + anthropic_messages_optional_request_params=call.params, + logging_obj=call.logger, + custom_llm_provider=call.provider, + kwargs=call.kwargs, + hold_back=bool(held), + server_fulfilled_tool_names=held, + ), + hidden_params=cast( # cast-ok: the stream's stamped hidden params, shared as Python's handler shares them + AnthropicMessagesStreamHiddenParams, + stream._hidden_params, # pyright: ignore[reportPrivateUsage] # the native stream's own hidden params + ), + ) + + +def agentic_stream( + head: object, + request: Mapping[str, object], + sent: Mapping[str, object], + model: str, + custom_llm_provider: str, + requested: Mapping[str, object], + hooked: bool, +) -> None: + logger: Final = _logger(request) + if ( + not isinstance(head, MessagesStreamHead) + or logger is None + or not BaseLLMHTTPHandler._has_agentic_completion_hook(logger) # pyright: ignore[reportPrivateUsage] # the handler's own agentic gate + ): + return + call: Final = _agentic_call(request, sent, model, custom_llm_provider, requested, hooked) + held: Final = BaseLLMHTTPHandler._server_fulfilled_tools_in_request( # pyright: ignore[reportPrivateUsage] # the handler's own hold-back gate + logging_obj=logger, tools=call.params.get("tools") + ) + wrap: Final = functools.partial(_agentic_stream, call=call, held=held) + head.agentic = wrap # rebind-ok: the head carries the stream wrapper back + + +class MessagesSafeguards: + def __init__(self, request: Mapping[str, object]) -> None: + self._request: Final = request + self._frames: _StampedFrames | None = None + self._results: object = None + + async def _evaluator(self) -> SafeguardsEvaluator | None: + model, provider = _handler_model(self._request) + if not await handler.proxy_classifies_tool_uses(model, provider, self._request): + return None + messages: Final = normalize_native_messages(optional_sequence(self._request.get("messages")) or ()) + return await asyncio.to_thread( + handler._native_safeguards_evaluator, # pyright: ignore[reportPrivateUsage] # the handler's own evaluator build + self._request, + messages, + ) + + async def results(self, value: Mapping[str, object]) -> object: + evaluator: Final = await self._evaluator() + if evaluator is None: + return None + unstamped: Final = response(value) + stamped: Final = await with_safeguard_results(unstamped, evaluator) + if stamped is unstamped: + return None + self._results = stamped.get("safeguard_results") + return self._results + + def published(self, public: Mapping[str, object]) -> Mapping[str, object]: + return public if self._results is None else {**public, "safeguard_results": self._results} + + async def open_stream(self) -> bool: + evaluator: Final = await self._evaluator() + if evaluator is None: + return False + self._frames = _StampedFrames(StreamedSafeguardResults(evaluator)) + return True + + async def frames(self, chunk: bytes) -> Sequence[bytes]: + return await self._open_frames().feed(chunk) + + async def flush(self) -> Sequence[bytes]: + return await self._open_frames().flush() + + def _open_frames(self) -> _StampedFrames: + if self._frames is None: + raise RuntimeError("safeguards stream frames before the stream opened") + return self._frames + + +class _StampedFrames: + def __init__(self, safeguard_results: StreamedSafeguardResults) -> None: + self._safeguard_results: Final = safeguard_results + self._open_frame: Final = bytearray() + + async def feed(self, chunk: bytes) -> Sequence[bytes]: + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + _SSE_FRAME_END, # pyright: ignore[reportPrivateUsage] # the handler's own SSE framing + _SSE_TERMINATOR_OVERLAP, # pyright: ignore[reportPrivateUsage] # the handler's own SSE framing + _ends_a_frame, # pyright: ignore[reportPrivateUsage] # the handler's own SSE framing + _stamped_sse_frame, # pyright: ignore[reportPrivateUsage] # the handler's own SSE framing + ) + + scan_from: Final = max(len(self._open_frame) - _SSE_TERMINATOR_OVERLAP, 0) + self._open_frame.extend(chunk) + if not _ends_a_frame(self._open_frame, scan_from): + return () + *frames, unterminated = _SSE_FRAME_END.split(self._open_frame) + self._open_frame[:] = unterminated + return tuple([bytes(await _stamped_sse_frame(frame, self._safeguard_results)) for frame in frames]) + + async def flush(self) -> Sequence[bytes]: + if not self._open_frame: + return () + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + _stamped_sse_frame, # pyright: ignore[reportPrivateUsage] # the handler's own SSE framing + ) + + return (bytes(await _stamped_sse_frame(bytes(self._open_frame), self._safeguard_results)),) + + +def _handler_model(request: Mapping[str, object]) -> tuple[str, str | None]: + explicit_provider: Final = optional_str(request.get("custom_llm_provider")) + try: + model, provider, _, _ = litellm.get_llm_provider( + model=str(request["model"]), custom_llm_provider=explicit_provider + ) + except litellm.BadRequestError: + return str(request["model"]), explicit_provider + return model, provider + + +def safeguards(request: Mapping[str, object]) -> MessagesSafeguards | None: + if "safeguards" not in request or not native_route_classifies(str(request["model"]), request["safeguards"]): + return None + return MessagesSafeguards(request) diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py deleted file mode 100644 index 99555bb4852..00000000000 --- a/litellm/rust_bridge/messages/detectors.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Detectors for Python-only Anthropic Messages behavior.""" - -from __future__ import annotations - -from typing import Final - -import litellm -from litellm.anthropic_beta_headers_manager import is_foundry_api_base -from litellm.litellm_core_utils.core_helpers import ( - drops_unsupported_params, - proxy_request_headers, - stream_cost_injection_requested, -) -from litellm.llms.opencode.common_utils import is_opencode_api_base -from litellm.rust_bridge.detectors import api_base, overridden, present, streaming -from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str - - -def provider_affinity(request: NativeCall) -> bool: - return is_opencode_api_base(api_base(request)) - - -def foundry_beta_row(request: NativeCall) -> bool: - from litellm.llms.anthropic.common_utils import AnthropicModelInfo - - return is_foundry_api_base(AnthropicModelInfo.get_api_base(api_base(request))) - - -def anthropic_workload_identity(request: NativeCall) -> bool: - """Python mints an Anthropic federation token when no static credential is configured.""" - from litellm.llms.anthropic.common_utils import AnthropicModelInfo - from litellm.secret_managers.main import get_secret_str - - def configured(param: str, env: str) -> bool: - return present(request, param) or get_secret_str(env) is not None - - return ( - AnthropicModelInfo.get_api_key(optional_str(request.resolved.get("api_key"))) is None - and AnthropicModelInfo.get_auth_token() is None - and configured("anthropic_federation_rule_id", "ANTHROPIC_FEDERATION_RULE_ID") - and configured("anthropic_organization_id", "ANTHROPIC_ORGANIZATION_ID") - ) - - -_THINKING_BLOCKS: Final = frozenset({"thinking", "redacted_thinking"}) - - -def _replays_thinking(message: object) -> bool: - content: Final = optional_sequence((optional_mapping(message) or {}).get("content")) or () - return any((optional_mapping(block) or {}).get("type") in _THINKING_BLOCKS for block in content) - - -def thinking_signature_retry(request: NativeCall) -> bool: - """Python may strip replayed thinking blocks and retry after a rejected signature.""" - return any(_replays_thinking(message) for message in optional_sequence(request.resolved.get("messages")) or ()) - - -def tool_custom_field(request: NativeCall) -> bool: - tools: Final = optional_sequence(request.resolved.get("tools")) or () - return any("custom" in (optional_mapping(tool) or {}) for tool in tools) - - -def strict_tools_retry(request: NativeCall) -> bool: - tools: Final = optional_sequence(request.resolved.get("tools")) or () - return drops_unsupported_params(request.resolved.get("drop_params")) and any( - (optional_mapping(tool) or {}).get("strict") is True for tool in tools - ) - - -def stream_cost_injection(request: NativeCall) -> bool: - """Python adds response cost to streaming usage when the global setting or the request header asks for it.""" - return streaming(request) and stream_cost_injection_requested(proxy_request_headers(request.resolved)) - - -def pre_request_hook(_request: NativeCall) -> bool: - """The async Messages handler runs callbacks that can rewrite the request.""" - return overridden("async_pre_request_hook", litellm.callbacks) - - -def websearch_interception(_request: NativeCall) -> bool: - """The async Messages handler can answer a web-search request without an LLM call.""" - from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger - - return any(isinstance(callback, WebSearchInterceptionLogger) for callback in litellm.callbacks) - - -def _resolved_provider(request: NativeCall) -> str | None: - from litellm.exceptions import BadRequestError - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - provider: Final = optional_str(request.resolved.get("custom_llm_provider")) - try: - return get_llm_provider(str(request.resolved["model"]), provider)[1] - except BadRequestError: - return provider - - -def intercepted(request: NativeCall) -> bool: - """The async Messages handler delegates matching tools to a registered interceptor.""" - from litellm.llms.anthropic.messages.native.interceptors import get_messages_interceptors - - tools: Final = optional_sequence(request.resolved.get("tools")) - if not tools: - return False - mappings: Final = tuple(mapping for mapping in map(optional_mapping, tools) if mapping is not None) - provider: Final = _resolved_provider(request) - return any(interceptor.can_handle(mappings, provider) for interceptor in get_messages_interceptors()) - - -def safeguards_classifier(request: NativeCall) -> bool: - """Python's proxy classifier evaluates dangerous tool use for non-Claude models.""" - from litellm.llms.anthropic.messages.safeguards import native_route_classifies - - return native_route_classifies(str(request.resolved["model"]), request.resolved.get("safeguards")) diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py index 756d62b747a..ca195d99c62 100644 --- a/litellm/rust_bridge/messages/entrypoints.py +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -1,6 +1,7 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables from litellm.rust_bridge import execution @@ -9,6 +10,30 @@ from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +class MessagesStreamHead(dict[str, object]): + agentic: Callable[[execution.AsyncStream], object] | None = None + deliver: Callable[[object], Awaitable[object]] | None = None + + +@dataclass(frozen=True, slots=True) +class _Undelivered: + stream: object + deliver: Callable[[object], Awaitable[object]] + + +def _open_stream(native: execution.Execution, head: object) -> object: + stream: Final = execution.AsyncStream(native, head) + agentic: Final = head.agentic if isinstance(head, MessagesStreamHead) else None + opened: Final = stream if agentic is None else agentic(stream) + deliver: Final = head.deliver if isinstance(head, MessagesStreamHead) else None + return opened if deliver is None else _Undelivered(opened, deliver) + + +async def _delivered(opened: Awaitable[object]) -> object: + result: Final = await opened + return await result.deliver(result.stream) if isinstance(result, _Undelivered) else result + + class NativeMessages(Protocol): def __call__( self, @@ -33,8 +58,12 @@ def _messages_binding(value: object) -> NativeMessages | None: def _amessages_binding(value: object) -> NativeAmessages | None: if not callable(value): return None - native: Final = execution.asynchronous(value) - return cast("NativeAmessages", native) # cast-ok: the execution adapter returns the route's result + native: Final = execution.asynchronous(value, _open_stream) + + def amessages(request: NativeCall) -> Awaitable[object]: + return _delivered(native(request)) + + return cast("NativeAmessages", amessages) # cast-ok: the execution adapter returns the route's result NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding) diff --git a/litellm/rust_bridge/response_metadata.py b/litellm/rust_bridge/response_metadata.py index d2ddf61f7db..b5293bbe291 100644 --- a/litellm/rust_bridge/response_metadata.py +++ b/litellm/rust_bridge/response_metadata.py @@ -4,12 +4,15 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # reuse the proxy's identity-preserving response metadata writer get_hidden_params_dict, ) +from litellm.rust_bridge.catalog import Route ResultT: Final = TypeVar("ResultT") -def mark_rust_response(response: ResultT) -> ResultT: - cache_key: Final = get_hidden_params_dict(response).get("cache_key") +def mark_rust_response(response: ResultT, route: Route) -> ResultT: + if isinstance(response, dict): + return response # pyright: ignore[reportUnknownVariableType] # the caller's own dict, returned untouched + cache_key: Final = None if route is Route.MESSAGES else get_hidden_params_dict(response).get("cache_key") _add_headers_to_response( response, {"x-litellm-rust": "true", **({"x-litellm-cache-key": cache_key} if isinstance(cache_key, str) else {})}, diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index e22ebe73b08..6159d8de551 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -134,7 +134,7 @@ async def _aattempt_native( def _required(result: RustAttempt[ResultT], context: RouteContext) -> ResultT: if isinstance(result, RustHandled): - return mark_rust_response(result.value) + return mark_rust_response(result.value, context.route) _raise_required(result, _error_context(context)) diff --git a/litellm/utils.py b/litellm/utils.py index 2820638cd28..1d035215ca5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2087,6 +2087,39 @@ def _is_litellm_router_call(kwargs: Mapping[str, object]) -> bool: return any(isinstance(bucket, Mapping) and "model_group" in bucket for bucket in metadata_buckets) +def clamp_max_tokens( + model: object, + args: tuple[object, ...], + kwargs: dict[str, object], # mutable-ok: the @client wrapper rewrites the call's own kwargs + call_type: str, +) -> None: + if ( + kwargs.get("max_tokens", None) is not None + and model is not None + and litellm.modify_params is True + and call_type in (CallTypes.acompletion.value, CallTypes.completion.value, CallTypes.anthropic_messages.value) + ): + try: + base_model = model + if kwargs.get("hf_model_name", None) is not None: + base_model = f"huggingface/{kwargs.get('hf_model_name')}" + messages = None + if len(args) > 1: + messages = args[1] + elif kwargs.get("messages", None): + messages = kwargs["messages"] + kwargs["max_tokens"] = _get_modified_max_tokens()( + model=model, + base_model=base_model, + messages=messages, + user_max_tokens=kwargs.get("max_tokens"), + buffer_num=None, + buffer_perc=None, + ) + except Exception as e: + print_verbose(f"Error while checking max token limit: {e}") + + def client(original_function): from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit @@ -2203,38 +2236,7 @@ def client(original_function): verbose_logger.debug("Cache hit!") return caching_handler_response.cached_result - # CHECK MAX TOKENS - if ( - kwargs.get("max_tokens", None) is not None - and model is not None - and litellm.modify_params is True # user is okay with params being modified - and ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.completion.value - or call_type == CallTypes.anthropic_messages.value - ) - ): - try: - base_model = model - if kwargs.get("hf_model_name", None) is not None: - base_model = f"huggingface/{kwargs.get('hf_model_name')}" - messages = None - if len(args) > 1: - messages = args[1] - elif kwargs.get("messages", None): - messages = kwargs["messages"] - user_max_tokens: Final = kwargs.get("max_tokens") - modified_max_tokens: Final = _get_modified_max_tokens()( - model=model, - base_model=base_model, - messages=messages, - user_max_tokens=user_max_tokens, - buffer_num=None, - buffer_perc=None, - ) - kwargs["max_tokens"] = modified_max_tokens - except Exception as e: - print_verbose(f"Error while checking max token limit: {e}") + clamp_max_tokens(model, args, kwargs, call_type) # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() @@ -2499,38 +2501,7 @@ def client(original_function): if _llm_caching_handler.preset_cache_key is not None: logging_obj.litellm_params["preset_cache_key"] = _llm_caching_handler.preset_cache_key - # CHECK MAX TOKENS - if ( - kwargs.get("max_tokens", None) is not None - and model is not None - and litellm.modify_params is True # user is okay with params being modified - and ( - call_type == CallTypes.acompletion.value - or call_type == CallTypes.completion.value - or call_type == CallTypes.anthropic_messages.value - ) - ): - try: - base_model = model - if kwargs.get("hf_model_name", None) is not None: - base_model = f"huggingface/{kwargs.get('hf_model_name')}" - messages = None - if len(args) > 1: - messages = args[1] - elif kwargs.get("messages", None): - messages = kwargs["messages"] - user_max_tokens: Final = kwargs.get("max_tokens") - modified_max_tokens: Final = _get_modified_max_tokens()( - model=model, - base_model=base_model, - messages=messages, - user_max_tokens=user_max_tokens, - buffer_num=None, - buffer_perc=None, - ) - kwargs["max_tokens"] = modified_max_tokens - except Exception as e: - print_verbose(f"Error while checking max token limit: {e}") + clamp_max_tokens(model, args, kwargs, call_type) # MODEL CALL call_kwargs: Final = ( diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py index 41008f0d57f..224af4d5613 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -20,25 +20,68 @@ R request, S response, C callbacks, F covered by a fallback detector today) """ import asyncio +import contextlib import copy +import datetime +import gzip import json +import logging +import math import os import re -from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Iterable, Iterator, Mapping +import time +import uuid +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Generator, + Iterable, + Iterator, + Mapping, + Sequence, +) from contextlib import contextmanager from dataclasses import asdict, dataclass, field, replace from types import MappingProxyType -from typing import Final, cast +from typing import ClassVar, Final, cast +import httpx import pytest from pydantic import JsonValue, TypeAdapter import litellm +from litellm._logging import verbose_logger +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.caching.caching import Cache, CacheMode +from litellm.caching.caching_handler import ( + _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # a primed entry must land before the measured call reads it +) +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger +from litellm.litellm_core_utils import logging_utils from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.messages.native import interceptors as messages_interceptors +from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator +from litellm.llms.anthropic.messages.native.interceptors import MessagesInterceptor +from litellm.llms.base_llm.search.transformation import SearchResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.messages import dispatch +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.hooks.message_thread_ownership import MessageThreadOwnership +from litellm.proxy.utils import InternalUsageCache +from litellm.router_utils.cooldown_handlers import is_message_thread_not_found from litellm.rust_bridge import catalog -from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec @@ -51,8 +94,8 @@ OPUS_FAST: Final = "anthropic/claude-opus-5-5" API_KEY: Final = "sk-ant-parity-test" USER: Final[Mapping[str, JsonValue]] = MappingProxyType({"role": "user", "content": "What is the capital of France?"}) -_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _LOGGED: Final = TypeAdapter(dict[str, object]) +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _TRANSPORT_HEADERS: Final = frozenset( {"accept", "accept-encoding", "connection", "content-length", "host", "user-agent"} ) @@ -82,6 +125,7 @@ _VOLATILE: Final = ( (re.compile(r"Timeout passed=[0-9.]+"), "Timeout passed="), (re.compile(r"127\.0\.0\.1:[0-9]+"), "127.0.0.1:"), ) +_SERVER_TOOL_ID: Final = re.compile(r"srvtoolu_[0-9a-f]{32}") TEXT_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( { @@ -161,6 +205,7 @@ THINKING_REPLAY: Final[tuple[JsonValue, ...]] = ( {"role": "user", "content": "And Germany?"}, ) +_MODIFY_PARAMS: Final[Mapping[str, object]] = MappingProxyType({"modify_params": True}) NO_INPUT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) PNG_PIXEL: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" @@ -283,6 +328,60 @@ SLOW_STREAM: Final = _events( ) +WIDE_NUMBERS_REPLY: Final = ( + b'{"id":"msg_01ParityWide","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929",' + b'"content":[{"type":"tool_use","id":"toolu_01ParityWide","name":"get_ids","input":{"big":18446744073709551616,' + b'"huge":12345678901234567890123,"negative":-12345678901234567890123,"neg_zero":-0.0,"overflow":1e400,' + b'"underflow":-1E+400,"exp":1E2,"tiny":5e-324,"precise":0.30000000000000004,"text":"1e400 18446744073709551616"}}],' + b'"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":12,"output_tokens":9}}' +) +WIDE_INTEGER_REPLY: Final = ( + b'{"id":"msg_01ParityWideCached","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929",' + b'"content":[{"type":"tool_use","id":"toolu_01ParityWideCached","name":"get_ids","input":{"huge":12345678901234567890123}}],' + b'"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":12,"output_tokens":9}}' +) + +MARKED_TEXTS_REPLY: Final = ( + '{"id":"msg_01ParityMarked","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929",' + '"content":[{"type":"text","text":"\\uf8fflitellm-raw-json:42"},{"type":"tool_use","id":"toolu_01ParityMarked",' + '"name":"get_ids","input":{"n":"\uf8fflitellm-raw-json:NaN","s":"\\uF8FFlitellm-raw-json:\\"x\\"",' + '"\\uf8fflitellm-raw-json:1":true,"big":18446744073709551616}}],' + '"stop_reason":"tool_use","stop_sequence":null,"usage":{"input_tokens":12,"output_tokens":9}}' +).encode() + + +def _non_string_tool_result_texts( + *texts: object, +) -> list[dict[str, object]]: # mutable-ok: the messages kwarg is a list of dicts + return [ + dict(USER), + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01Texts", "name": "get_ids", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01Texts", + "content": [{"type": "text", "text": text} for text in texts], + }, + {"type": "text", "text": "go on"}, + ], + }, + ] + + +_BAD_TRAILER_GZIP: Final = gzip.compress(json.dumps(dict(TEXT_RESPONSE)).encode(), mtime=0)[:-8] + bytes(8) +_NOT_GZIP: Final = b"\x1f\x8b\x08\x00not really gzip data at allnot really gzip data at all" + + +def _framed_bytes(events: tuple[tuple[str, object], ...]) -> tuple[bytes, ...]: + return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in events) + + +def _encoded(body: bytes, encoding: str) -> ResponseSpec: + return ResponseSpec(body=body, headers={"content-encoding": encoding}) + + def _upstream_error(kind: str, message: str) -> Mapping[str, JsonValue]: return MappingProxyType({"type": "error", "error": {"type": kind, "message": message}}) @@ -295,6 +394,45 @@ def _stream(events: tuple[tuple[str, object], ...], chunk_delay: float = 0) -> R return ResponseSpec(body={}, events=events, chunk_delay=chunk_delay) +def _burst(events: tuple[tuple[str, object], ...], chunked: bool = True) -> ResponseSpec: + return ResponseSpec(body={}, events=events, burst=True, chunked=chunked) + + +def _chunked(events: tuple[tuple[str, object], ...], chunk_delay: float) -> ResponseSpec: + return ResponseSpec(body={}, events=events, chunk_delay=chunk_delay, chunked=True) + + +_PING: Final[tuple[str, JsonValue]] = ("ping", {"type": "ping"}) + +PINGED_STREAM: Final = _events( + _message_start("msg_01ParityPinged"), + *((_PING,) * 50), + *_text_block(0, "Pa", "ris"), + *_message_end("end_turn", 4), +) + +HEAVILY_PINGED_STREAM: Final = _events( + _message_start("msg_01ParityHeavilyPinged"), + *((_PING,) * 1000), + *_text_block(0, "Pa", "ris"), + *_message_end("end_turn", 4), +) + +LARGE_STREAM: Final = _events( + _message_start("msg_01ParityLarge"), + *_text_block(0, "a" * 15000, "b" * 15000, "c" * 15000), + *_message_end("end_turn", 4), +) + + +def _raw_stream(*frames: str) -> ResponseSpec: + return ResponseSpec(body="".join(f"{frame}\n\n" for frame in frames).encode()) + + +def _framed(events: tuple[tuple[str, object], ...], event_line: str = "event: ") -> tuple[str, ...]: + return tuple(f"{event_line}{event}\ndata: {json.dumps(data)}" for event, data in events) + + def _failure(status: int, kind: str, message: str, **headers: str) -> ResponseSpec: return ResponseSpec(body=dict(_upstream_error(kind, message)), status=status, headers=dict(headers)) @@ -305,6 +443,10 @@ CALLBACK_ASPECTS: Final = ( "logged_response_headers", "logged_api_key", "logged_call", + "logged_request", + "logged_sent_request", + "logged_model_parameters", + "logged_error", ) ASPECTS: Final = ("dispatch", "request", "outcome", "hidden_params", "logged", *CALLBACK_ASPECTS) @@ -319,12 +461,38 @@ class Case: env: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) settings: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) private_keys: bool = False + cache: Callable[[], Cache] | None = None + callbacks: tuple[Callable[[], CustomLogger], ...] = () + classifier: str | None = None + interceptor: Callable[[], MessagesInterceptor] | None = None gaps: Mapping[str, str | None] = field(default_factory=lambda: MappingProxyType({})) + loggers: Callable[[], Sequence[CustomLogger]] | None = None + gateway: bool = False + racy: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + partitioned: bool = False + nones: tuple[str, ...] = () + busy_before_reading: float = 0 + warned: bool = False + + def aspects(self) -> tuple[str, ...]: + return ( + *ASPECTS, + *(("cached",) if self.cache is not None else ()), + *(("warnings",) if self.warned else ()), + *(("hook_inputs",) if self.loggers else ()), + *(("classified",) if self.classifier is not None else ()), + *(("tool_runs",) if self.gateway else ()), + ) def params(self) -> Iterator[object]: - for aspect in ASPECTS: + for aspect in self.aspects(): gap = {**KNOWN_GAPS.get(self.name, {}), **self.gaps}.get(aspect) - marks = () if not gap else (pytest.mark.xfail(reason=gap, strict=True),) + race = self.racy.get(aspect) + marks = ( + (pytest.mark.skip(reason=race),) + if race + else (() if not gap else (pytest.mark.xfail(reason=gap, strict=True),)) + ) yield pytest.param(self, aspect, id=f"{self.name}-{aspect}", marks=marks) @@ -336,54 +504,48 @@ def _kw(**kwargs: object) -> Mapping[str, object]: return MappingProxyType(kwargs) +def _tls(**settings: object) -> Mapping[str, object]: + return MappingProxyType({**settings, "in_memory_llm_clients_cache": LLMClientCache()}) + + +E04_HTTPX_CHUNK_READS: Final = "E04: httpx yields each HTTP chunk whole from a 64 KiB socket read; hyper hands a large chunk over in pieces of its smaller read buffer" A10_NO_MAX_TOKENS: Final = "A10: native admission hands a call without an int max_tokens to Python, whose signature raises the TypeError; only dispatch differs" -S6_PRIVATE_KEYS: Final = "S6: native writes _hidden_params into the response dict" -S15_COMPLETE_DEADLINE: Final = ( - "S15: by design, native bounds a complete response by one overall deadline where Python bounds each read, " - "so a body trickling in slower than the timeout succeeds in Python and times out natively" +C_UNSENT_CALL_FLAG: Final = ( + "C: a call that fails natively before its provider request logs no aanthropic_messages flag, " + "which Python's handler sets before it shapes the request" ) -F_MOCK: Final = "F: mock_response is Python-only (detector mock_response)" -F_THINKING_RETRY: Final = ( - "F: retry without thinking after a rejected signature is Python-only (detector thinking_signature_retry)" + +F46_SEE_OTHER_CONTENT_TYPE: Final = ( + "F46: httpx keeps content-type on the GET that follows a 303; reqwest's redirect follower drops it with the body" +) +REDIRECTED_SENT_REQUEST: Final = ( + "a followed provider redirect: Python logs the redirected request (its URL, its body already streamed) on the " + "logged response, the native call logs the request it sent" ) -F_CACHE_INJECTION: Final = "F: prompt cache_control injection is Python-only (detector cache_control_injection)" -F_TOOL_CUSTOM: Final = "F: hoisting a tool's custom field is Python-only (detector tool_custom_field)" -F_STREAM_COST: Final = "F: cost injection into streamed usage is Python-only (detector stream_cost_injection)" KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType( { - "response_private_keys": _gaps(outcome=S6_PRIVATE_KEYS), "missing_max_tokens": _gaps(dispatch=A10_NO_MAX_TOKENS), - "mock_response": _gaps( - logged_response_headers=F_MOCK, - request=F_MOCK, - outcome=F_MOCK, - hidden_params=F_MOCK, - logged=F_MOCK, - logged_metadata=F_MOCK, - logged_api_key=F_MOCK, - logged_call=F_MOCK, - ), - "mock_response_rate_limit": _gaps( - logged_response_headers=F_MOCK, - request=F_MOCK, - outcome=F_MOCK, - logged=F_MOCK, - logged_api_key=F_MOCK, - logged_call=F_MOCK, - ), - "thinking_signature_rejected_then_retried": _gaps( - request=F_THINKING_RETRY, - outcome=F_THINKING_RETRY, - hidden_params=F_THINKING_RETRY, - logged=F_THINKING_RETRY, - logged_metadata=F_THINKING_RETRY, - logged_response_headers=F_THINKING_RETRY, - ), - "cache_control_injection_points": _gaps(request=F_CACHE_INJECTION), - "tool_custom_field": _gaps(request=F_TOOL_CUSTOM), - "stream_cost_injection": _gaps(outcome=F_STREAM_COST), - "enable_prompt_caching": _gaps(request=F_CACHE_INJECTION), + **{ + name: _gaps(logged_call=C_UNSENT_CALL_FLAG) + for name in ( + "speed_on_model_without_speed", + "temperature_on_adaptive_model", + "top_k_on_adaptive_model", + "legacy_thinking_non_numeric_budget_on_adaptive_model", + "legacy_thinking_null_budget_on_adaptive_model", + "missing_api_key", + "tls_security_level_python_cannot_apply", + "max_tokens_none", + "reasoning_effort_unknown", + "output_config_effort_unknown", + "reasoning_effort_unknown_on_adaptive_model", + "reasoning_effort_xhigh_on_model_without_xhigh", + "int_temperature_on_adaptive_model", + "int_top_p_on_adaptive_model", + "int_speed_on_model_without_speed", + ) + }, } ) @@ -440,6 +602,7 @@ REQUEST_CASES: Final = ( "legacy_thinking_non_numeric_budget_dropped", _kw(thinking={"type": "enabled", "budget_tokens": [1]}, max_tokens=8192, drop_params=True), model=OPUS_ADAPTIVE, + warned=True, ), Case( "legacy_thinking_carries_extra_keys", @@ -657,6 +820,192 @@ RESPONSE_CASES: Final = ( ), ), Case("response_private_keys", private_keys=True), + Case( + "response_gzip_under_a_padded_mixed_case_alias", + upstream=(_encoded(gzip.compress(json.dumps(dict(TEXT_RESPONSE)).encode(), mtime=0), " X-Gzip "),), + ), + Case("response_deflate_body_that_is_not_deflate", upstream=(_encoded(b"not deflate at all", "deflate"),)), + Case("response_gzip_body_with_a_bad_trailer", upstream=(_encoded(_BAD_TRAILER_GZIP, "gzip"),)), + Case("response_gzip_body_that_is_not_gzip", upstream=(_encoded(_NOT_GZIP, "gzip"),)), + Case( + "stream_deflate_body_that_is_not_deflate", + _kw(stream=True), + upstream=(_encoded(b"not deflate at all", "deflate"),), + ), + Case( + "stream_gzip_listed_with_an_unknown_coding", + _kw(stream=True), + upstream=(_encoded(gzip.compress(b"".join(_framed_bytes(TEXT_STREAM)), mtime=0), "gzip, br"),), + ), +) + +INPUT_VALUE_CASES: Final = ( + Case("max_tokens_numeric_string", _kw(max_tokens="1024")), + Case("max_tokens_fraction", _kw(max_tokens=1.5)), + Case("max_tokens_whole_float", _kw(max_tokens=1024.0)), + Case("max_tokens_bool", _kw(max_tokens=True)), + Case("max_tokens_beyond_64_bits", _kw(max_tokens=2**64)), + Case("reasoning_effort_with_fractional_max_tokens", _kw(reasoning_effort="high", max_tokens=1500.5)), + Case("temperature_numeric_string", _kw(temperature="0.5")), + Case("temperature_integer_beyond_float_precision", _kw(temperature=2**60 + 1)), + Case("top_p_integer_beyond_float_precision", _kw(top_p=-9223372036854775807)), + Case( + "reasoning_effort_with_numeric_string_max_tokens", + _kw(reasoning_effort="high", max_tokens="8192"), + upstream=(), + ), + Case("top_k_beyond_64_bits", _kw(top_k=2**70)), + Case("messages_mapping", _kw(messages={"role": "user", "content": "hi"})), + Case("messages_string", _kw(messages="hi")), + Case("messages_with_non_object_entries", _kw(messages=[dict(USER), "stray", 7])), + Case("system_lone_surrogate", _kw(system="sys \ud800 lone")), + Case("metadata_user_id_lone_surrogate", _kw(metadata={"user_id": "user-\udc00"})), + Case("message_text_lone_surrogate", _kw(messages=[{"role": "user", "content": "a\ud800b \u00e9"}])), + Case("text_starting_with_the_raw_leaf_mark", _kw(system="\uf8fflitellm-raw-json:1")), + Case("thinking_budget_beyond_64_bits", _kw(thinking={"type": "enabled", "budget_tokens": 2**70}, max_tokens=2048)), + Case( + "tool_schema_integers_beyond_64_bits", + _kw(tools=[{**WEATHER_TOOL, "input_schema": {"type": "integer", "maximum": 2**70, "minimum": -(2**70)}}]), + ), + Case("temperature_nan", _kw(temperature=float("nan"))), + Case( + "tool_schema_infinities", + _kw( + tools=[ + {**WEATHER_TOOL, "input_schema": {"type": "number", "maximum": float("inf"), "minimum": -float("inf")}} + ] + ), + ), +) + +LOOSE_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"name": "lookup", "description": "lookup", "input_schema": {"type": "object", "properties": {}}} +) + + +class _StringStreamHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + kwargs["stream"] = "true" + return kwargs + + +class _IntegerStreamHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + kwargs["stream"] = 1 + return kwargs + + +class _StringParamsHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + kwargs["litellm_params"] = "not-a-mapping" + return kwargs + + +class _IntegerProviderHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + params: Final = kwargs.setdefault("litellm_params", {}) + assert isinstance(params, dict) + params["custom_llm_provider"] = 5 + return kwargs + + +ARGUMENT_TYPE_CASES: Final = ( + Case("stream_string_true", _kw(stream="true"), upstream=(_stream(TEXT_STREAM),)), + Case("stream_string_false", _kw(stream="false"), upstream=(_stream(TEXT_STREAM),)), + Case("stream_zero", _kw(stream=0)), + Case("stream_one", _kw(stream=1), upstream=(_stream(TEXT_STREAM),)), + Case("stream_one_sync", _kw(stream=1), upstream=(_stream(TEXT_STREAM),), sync=True), + Case( + "pre_request_hook_streams_on_a_string", + _kw(stream=True), + callbacks=(_StringStreamHook,), + upstream=(_stream(TEXT_STREAM),), + ), + Case("pre_request_hook_streams_on_an_integer", callbacks=(_IntegerStreamHook,), upstream=(_stream(TEXT_STREAM),)), + Case("pre_request_hook_sets_string_litellm_params", callbacks=(_StringParamsHook,), model="claude-sonnet-4-5"), + Case("pre_request_hook_sets_integer_provider", callbacks=(_IntegerProviderHook,), model="claude-sonnet-4-5"), + Case("model_info_string", _kw(model_info="not-a-mapping")), + Case("tools_mapping", _kw(tools=dict(LOOSE_TOOL)), upstream=()), + Case("messages_none", _kw(messages=None), nones=("messages",), upstream=()), + Case("max_tokens_none", _kw(max_tokens=None), nones=("max_tokens",), upstream=()), + Case("reasoning_effort_unknown", _kw(reasoning_effort="ultra"), upstream=()), + Case("reasoning_effort_unknown_on_adaptive_model", _kw(reasoning_effort="ultra"), model=OPUS_ADAPTIVE, upstream=()), + Case( + "reasoning_effort_xhigh_on_model_without_xhigh", + _kw(reasoning_effort="xhigh", max_tokens=8192), + model="anthropic/claude-opus-4-6", + upstream=(), + ), + Case("int_temperature_on_adaptive_model", _kw(temperature=0), model=OPUS_ADAPTIVE, upstream=()), + Case("int_top_p_on_adaptive_model", _kw(top_p=1), model=OPUS_ADAPTIVE, upstream=()), + Case("int_speed_on_model_without_speed", _kw(speed=2), upstream=()), + Case("string_output_format_dropped", _kw(output_format="json", drop_params=True)), + Case("string_output_format_kept_without_drop_params", _kw(output_format="json")), + Case("non_string_reasoning_effort_warns", _kw(reasoning_effort=3), warned=True), + Case( + "reasoning_effort_that_cannot_fit_max_tokens_warns", + _kw(reasoning_effort="high", max_tokens=500), + warned=True, + ), + Case( + "adaptive_thinking_that_cannot_fit_max_tokens_warns", + _kw(thinking={"type": "adaptive"}, output_config={"effort": "high"}, max_tokens=1024), + warned=True, + ), + Case( + "adaptive_thinking_on_model_without_reasoning_warns", + _kw(thinking={"type": "adaptive"}), + model="anthropic/claude-3-haiku-20240307", + warned=True, + ), + Case("output_config_effort_unknown", _kw(output_config={"effort": "ultra"})), + Case( + "global_drop_params_string_true", + _kw(tools=[{**LOOSE_TOOL, "custom": {"note": "opaque"}}]), + settings=_kw(drop_params="true"), + ), + Case( + "global_drop_params_string_false", + _kw(tools=[{**LOOSE_TOOL, "custom": {"note": "opaque"}}]), + settings=_kw(drop_params="false"), + ), + Case( + "request_drop_params_string_true", + _kw(drop_params="true", tools=[{**LOOSE_TOOL, "custom": {"note": "opaque"}}]), + ), + Case( + "additional_drop_params_tool_custom", + _kw( + additional_drop_params=["tools[*].custom"], + tools=[{**LOOSE_TOOL, "custom": {"defer_loading": True, "note": "x"}}], + ), + ), + Case( + "additional_drop_params_tool_input_examples", + _kw( + additional_drop_params=["tools[*].input_examples", "metadata.user_id", "fz.nope"], + metadata={"user_id": "u1"}, + tools=[{**LOOSE_TOOL, "input_examples": [{"q": "x"}]}], + ), + ), + Case( + "response_without_stop_reason", + upstream=(_ok({key: value for key, value in TEXT_RESPONSE.items() if key != "stop_reason"}),), + ), + Case( + "response_without_stop_fields", + upstream=( + _ok({key: value for key, value in TEXT_RESPONSE.items() if key not in ("stop_reason", "stop_sequence")}), + ), + ), ) FAILURE_CASES: Final = ( @@ -684,16 +1033,122 @@ FAILURE_CASES: Final = ( "response_trickles_slower_than_timeout", _kw(timeout=1), upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), chunk_delay=0.5, body_chunks=4),), - gaps=_gaps( - outcome=S15_COMPLETE_DEADLINE, - hidden_params=S15_COMPLETE_DEADLINE, - logged=S15_COMPLETE_DEADLINE, - logged_metadata=S15_COMPLETE_DEADLINE, - logged_response_headers=S15_COMPLETE_DEADLINE, + ), + Case( + "response_body_stalls_past_timeout", + _kw(timeout=0.5), + upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), chunk_delay=1.0, body_chunks=2),), + ), + Case( + "upstream_500_body_cut", + upstream=( + ResponseSpec(body=dict(_upstream_error("api_error", "cut")), status=500, body_chunks=2, cut_after=1), ), ), + Case("response_body_cut", upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), body_chunks=2, cut_after=1),)), + Case("gzip_response", upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), gzip=True),)), + Case("upstream_500_gzip", upstream=(replace(_failure(500, "api_error", "internal"), gzip=True),)), Case("missing_max_tokens", _kw(max_tokens=None), upstream=()), Case("missing_api_key", _kw(api_key=None), upstream=()), + Case( + "api_key_with_a_trailing_newline", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_API_KEY": f"{API_KEY}\n"}), + ), + Case( + "api_key_with_a_trailing_crlf_streamed", + _kw(api_key=None, stream=True), + upstream=(_stream(TEXT_STREAM),), + env=MappingProxyType({"ANTHROPIC_API_KEY": f"{API_KEY}\r\n"}), + ), + Case( + "api_key_padded_with_whitespace", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_API_KEY": f" {API_KEY}\t\n"}), + ), + Case( + "api_key_padded_with_separators", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_API_KEY": f"\x1f{API_KEY}\x1f"}), + ), + Case( + "api_key_of_separators_falls_to_the_auth_token", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_API_KEY": "\x1c\x1f\n", "ANTHROPIC_AUTH_TOKEN": "oauth-token"}), + ), + Case( + "auth_token_with_a_trailing_newline", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_AUTH_TOKEN": "oauth-token\n"}), + ), + Case("caller_api_key_padded_with_separators", _kw(api_key=f"\x1f{API_KEY}\x1f")), + Case("caller_api_key_with_a_trailing_newline", _kw(api_key=f"{API_KEY}\n")), + Case( + "auth_token_padded_with_whitespace", + _kw(api_key=None), + env=MappingProxyType({"ANTHROPIC_AUTH_TOKEN": " fz-env-auth-token "}), + ), + Case( + "provider_redirect_is_followed", + gaps=_gaps(logged_sent_request=REDIRECTED_SENT_REQUEST), + upstream=( + ResponseSpec(body=b"moved", status=307, headers={"location": "/redirected/v1/messages"}), + ResponseSpec(body=dict(TEXT_RESPONSE)), + ), + ), + Case( + "provider_redirect_chain_of_eleven_hops", + gaps=_gaps(logged_sent_request=REDIRECTED_SENT_REQUEST), + upstream=( + *( + ResponseSpec(body=b"moved", status=307, headers={"location": f"/hop{hop}/v1/messages"}) + for hop in range(11) + ), + ResponseSpec(body=dict(TEXT_RESPONSE)), + ), + ), + Case( + "provider_redirect_chain_past_twenty_hops", + upstream=tuple( + ResponseSpec(body=b"moved", status=307, headers={"location": f"/hop{hop}/v1/messages"}) for hop in range(21) + ), + ), + Case( + "provider_see_other_redirect", + gaps=_gaps(logged_sent_request=REDIRECTED_SENT_REQUEST, request=F46_SEE_OTHER_CONTENT_TYPE), + upstream=( + ResponseSpec(body=b"see other", status=303, headers={"location": "/other/v1/messages"}), + ResponseSpec(body=dict(TEXT_RESPONSE)), + ), + ), + Case("extra_header_lone_surrogate", _kw(extra_headers={"x-fz-surrogate": "h\ud800"}), upstream=()), + Case( + "provider_specific_header_lone_surrogate", + _kw( + provider_specific_header={ + "custom_llm_provider": "anthropic", + "extra_headers": {"x-fz-surrogate": "h\ud800"}, + } + ), + upstream=(), + ), + Case("extra_header_non_ascii", _kw(extra_headers={"x-fz-unicode": "caf\u00e9"}), upstream=()), + Case("caller_api_key_non_ascii", _kw(api_key="sk-ant-\u200djoined"), upstream=()), + Case("caller_api_key_non_ascii_stream", _kw(api_key="sk-ant-\u200djoined", stream=True), upstream=()), + Case("api_base_with_a_space_and_no_scheme", _kw(api_base="event: message_stop"), upstream=()), + Case("api_base_with_a_space_and_no_scheme_stream", _kw(api_base="data: [DONE]", stream=True), upstream=()), + Case("api_base_with_a_nul", _kw(api_base="a\x00b"), upstream=()), + Case("api_base_non_ascii_without_a_scheme", _kw(api_base="h\u00e9llo"), upstream=()), + Case("api_base_with_an_unclosed_ipv6_bracket", _kw(api_base_suffix="data: [DONE]"), upstream=()), + Case("provider_specific_header_as_a_string", _kw(provider_specific_header="anthropic"), upstream=()), + Case("response_header_latin_1", upstream=(_ok(**{"x-fz-unicode": "caf\u00e9"}),)), + Case("response_header_utf_8", upstream=(_ok(**{"x-fz-unicode": "caf\u00c3\u00a9"}),)), + Case( + "response_header_utf_8_httpx_transport", + upstream=(_ok(**{"x-fz-unicode": "caf\u00c3\u00a9"}),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + ), ) STREAM_CASES: Final = ( @@ -705,6 +1160,75 @@ STREAM_CASES: Final = ( upstream=(_stream(THINKING_STREAM),), ), Case("stream_error_event_mid_stream", _kw(stream=True), upstream=(_stream(ERROR_MID_STREAM),)), + Case("stream_events_in_one_write", _kw(stream=True), upstream=(_burst(TEXT_STREAM),), partitioned=True), + Case( + "stream_body_in_one_write", + _kw(stream=True), + upstream=(_burst(TEXT_STREAM, chunked=False),), + partitioned=True, + ), + Case( + "stream_events_in_chunks_written_apart", + _kw(stream=True), + upstream=(_chunked(TEXT_STREAM, chunk_delay=0.05),), + partitioned=True, + ), + Case( + "stream_events_written_apart", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM, chunk_delay=0.05),), + partitioned=True, + ), + Case( + "stream_pings_after_message_start_in_one_write", + _kw(stream=True), + upstream=(_burst(PINGED_STREAM),), + partitioned=True, + ), + Case( + "stream_pings_after_message_start_written_apart", + _kw(stream=True), + upstream=(_stream(PINGED_STREAM, chunk_delay=0.01),), + partitioned=True, + ), + Case("stream_burst_past_one_read", _kw(stream=True), upstream=(_burst(LARGE_STREAM),), partitioned=True), + Case( + "stream_thousand_pings_in_one_write", + _kw(stream=True), + upstream=(_burst(HEAVILY_PINGED_STREAM),), + partitioned=True, + ), + Case( + "stream_error_event_in_one_write", + _kw(stream=True), + upstream=(_burst(ERROR_MID_STREAM),), + partitioned=True, + ), + Case( + "stream_events_in_one_write_httpx_transport", + _kw(stream=True), + upstream=(_burst(TEXT_STREAM),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + partitioned=True, + ), + Case( + "stream_burst_past_one_read_httpx_transport", + _kw(stream=True), + upstream=(_burst(LARGE_STREAM),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + partitioned=True, + gaps=_gaps(outcome=E04_HTTPX_CHUNK_READS), + ), + Case( + "stream_events_written_apart_httpx_transport", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM, chunk_delay=0.05),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + partitioned=True, + ), Case( "stream_longer_than_timeout", _kw(stream=True, timeout=1), @@ -721,6 +1245,81 @@ STREAM_CASES: Final = ( upstream=(_failure(429, "rate_limit_error", "slow down", **{"retry-after": "3"}),), ), Case("stream_false_explicit", _kw(stream=False)), + Case( + "stream_upstream_500_body_cut", + _kw(stream=True), + upstream=( + ResponseSpec(body=dict(_upstream_error("api_error", "cut")), status=500, body_chunks=2, cut_after=1), + ), + ), + Case( + "stream_cut_after_four_events", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM), cut_after=4, cut_delay=0.1),), + ), + Case( + "stream_cut_after_message_start", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM), cut_after=1, cut_delay=0.1),), + ), + Case( + "stream_cut_after_events_written_apart", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM, chunk_delay=0.05), cut_after=6, cut_delay=0.1),), + ), + Case( + "stream_chunked_cut_after_four_events", + _kw(stream=True), + upstream=(replace(_chunked(TEXT_STREAM, chunk_delay=0), cut_after=4, cut_delay=0.1),), + ), + Case( + "stream_cut_while_the_caller_is_busy", + _kw(stream=True), + upstream=(replace(_chunked(TEXT_STREAM, chunk_delay=0), cut_after=4, cut_delay=0.1),), + busy_before_reading=0.4, + ), + Case( + "stream_cut_after_message_start_while_the_caller_is_busy", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM), cut_after=1, cut_delay=0.1),), + busy_before_reading=0.4, + ), + Case( + "stream_upstream_500_body_stalls", + _kw(stream=True, timeout=8), + upstream=( + ResponseSpec(body=dict(_upstream_error("api_error", "stalled")), status=500, chunk_delay=6, body_chunks=2), + ), + ), + Case( + "stream_upstream_429_body_stalls_past_timeout", + _kw(stream=True, timeout=1), + upstream=( + ResponseSpec( + body=dict(_upstream_error("rate_limit_error", "stalled")), status=429, chunk_delay=2, body_chunks=2 + ), + ), + ), + Case("stream_gzip_response", _kw(stream=True), upstream=(replace(_stream(TEXT_STREAM), gzip=True),)), + Case( + "stream_read_stalls_past_timeout_httpx_transport", + _kw(stream=True, timeout=0.5), + upstream=(_stream(SLOW_STREAM, chunk_delay=1.0),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + ), + Case( + "stream_response_header_utf_8", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM), headers={"x-fz-unicode": "caf\u00c3\u00a9"}),), + ), + Case( + "stream_response_header_utf_8_httpx_transport", + _kw(stream=True), + upstream=(replace(_stream(TEXT_STREAM), headers={"x-fz-unicode": "caf\u00c3\u00a9"}),), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + ), ) SYNC_CASES: Final = ( @@ -729,47 +1328,2022 @@ SYNC_CASES: Final = ( Case("sync_upstream_429", upstream=(_failure(429, "rate_limit_error", "slow down"),), sync=True), ) -FALLBACK_CASES: Final = ( - Case("mock_response", _kw(mock_response="Hi from the mock"), upstream=()), - Case("mock_response_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=()), - Case( - "thinking_replay", - _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), - ), - Case( - "thinking_signature_rejected_then_retried", - _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), - upstream=( - _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), - _ok(), - ), - ), +SYSTEM_POINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"location": "message", "role": "system"}) +LAST_MESSAGE_POINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"location": "message", "index": -1}) +SYSTEM_BLOCKS: Final[tuple[JsonValue, ...]] = ( + {"type": "text", "text": "Answer in one word."}, + {"type": "text", "text": "Be terse."}, +) +MARKED_TURNS: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": [{"type": "text", "text": "One.", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Two.", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "Three.", "cache_control": {"type": "ephemeral"}}]}, +) + +PROMPT_CACHING_CASES: Final = ( Case( "cache_control_injection_points", - _kw(system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "system"}]), + _kw(system="Answer in one word.", cache_control_injection_points=[dict(SYSTEM_POINT)]), ), - Case("tool_custom_field", _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True}}])), + Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), Case( - "strict_tools_rejected_then_retried", - _kw(tools=[{**WEATHER_TOOL, "strict": True}], drop_params=True), - upstream=( - _failure(400, "invalid_request_error", "tools.0.strict: too many strict tools"), - _ok(), + "enable_prompt_caching_without_system", + _kw(messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)], enable_prompt_caching=True), + ), + Case( + "global_prompt_caching_with_ttl", + _kw(system=list(SYSTEM_BLOCKS)), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True, "anthropic_prompt_caching_ttl": "1h"}), + ), + Case( + "global_prompt_caching_stands_down_on_client_marks", + _kw(system=[{"type": "text", "text": "Answer in one word.", "cache_control": {"type": "ephemeral"}}]), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "global_prompt_caching_stands_down_on_marked_tools", + _kw(tools=[{**WEATHER_TOOL, "cache_control": {"type": "ephemeral"}}]), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "enable_prompt_caching_false_overrides_nothing", + _kw(system="Answer in one word.", enable_prompt_caching=False), + ), + Case( + "injection_point_negative_index", + _kw( + messages=list(TOOL_TURNS), + cache_control_injection_points=[ + {"location": "message", "index": -2, "control": {"type": "ephemeral", "ttl": "1h"}} + ], ), ), - Case("modify_params", _kw(max_tokens=1), settings=MappingProxyType({"modify_params": True})), + Case( + "injection_point_string_index", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "index": "1"}]), + ), + Case( + "injection_point_out_of_range_index", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "index": 7}]), + ), + Case( + "injection_point_by_role", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "role": "user"}]), + ), + Case( + "injection_points_capped_by_client_marks", + _kw( + messages=list(MARKED_TURNS), + system=list(SYSTEM_BLOCKS), + tools=[{**WEATHER_TOOL, "cache_control": {"type": "ephemeral"}}], + cache_control_injection_points=[dict(SYSTEM_POINT), dict(LAST_MESSAGE_POINT)], + ), + ), + Case( + "injection_points_beside_client_marks", + _kw( + messages=[*MARKED_TURNS[:2], {"role": "user", "content": "Three."}], + cache_control_injection_points=[dict(SYSTEM_POINT), dict(LAST_MESSAGE_POINT)], + system="Answer in one word.", + ), + ), + Case( + "tool_config_point_marks_the_last_tool", + _kw( + tools=[dict(WEATHER_TOOL), {**WEATHER_TOOL, "name": "get_time"}], + cache_control_injection_points=[{"location": "tool_config"}, dict(LAST_MESSAGE_POINT)], + ), + ), + Case( + "tool_config_point_without_tools", + _kw(cache_control_injection_points=[{"location": "tool_config"}]), + ), + Case( + "injection_marks_the_gateway_deployment", + _kw( + system="Answer in one word.", + enable_prompt_caching=True, + litellm_metadata={"user_api_key": "hashed"}, + model_info={"id": "deployment-1"}, + ), + ), + Case( + "injection_into_whitespace_only_turn", + _kw( + messages=[dict(USER), {"role": "assistant", "content": "Paris"}, {"role": "user", "content": " "}], + cache_control_injection_points=[dict(LAST_MESSAGE_POINT)], + ), + upstream=(_failure(400, "invalid_request_error", "messages: text content blocks must be non-empty"),), + ), + Case( + "injection_after_empty_blocks_are_stripped", + _kw( + messages=[ + dict(USER), + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "And Germany?"}, + ], + cache_control_injection_points=[{"location": "message", "index": 1}], + ), + ), + Case( + "stream_with_prompt_caching", + _kw(stream=True, system="Answer in one word.", enable_prompt_caching=True), + upstream=(_stream(TEXT_STREAM),), + ), + Case( + "global_prompt_caching_beside_an_unmatched_point", + _kw( + system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "assistant"}] + ), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), +) + +SYNC_PROMPT_CACHING_CASES: Final = tuple( + replace(case, name=f"sync_{case.name}", sync=True) for case in PROMPT_CACHING_CASES +) + +LATEST_ALIAS: Final = "anthropic/claude-opus-latest" +CUSTOM_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType({**WEATHER_TOOL, "custom": {"defer_loading": True}}) +COST_HEADER: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}} +) + +SERVED_NATIVELY_CASES: Final = ( + Case("mock_response", _kw(mock_response="Hi from the mock"), upstream=()), + Case("mock_response_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=()), + Case("mock_response_context_window", _kw(mock_response="litellm.ContextWindowExceededError"), upstream=()), + Case("mock_response_empty_calls_the_provider", _kw(mock_response="")), + Case("mock_response_bad_metadata", _kw(mock_response="Hi", metadata={"user_id": 123}), upstream=()), + Case("mock_response_stream", _kw(mock_response="Hi from the mock", stream=True), upstream=()), + Case("mock_response_stream_rate_limit", _kw(mock_response="litellm.RateLimitError", stream=True), upstream=()), + Case( + "mock_response_stream_empty_calls_the_provider", + _kw(mock_response="", stream=True), + upstream=(_stream(TEXT_STREAM),), + ), + Case("mock_response_cached", _kw(mock_response="Hi from the mock"), upstream=(), cache=Cache), + Case("mock_response_stream_cached", _kw(mock_response="Hi from the mock", stream=True), upstream=(), cache=Cache), + Case("mock_response_sync", _kw(mock_response="Hi from the mock"), upstream=(), sync=True), + Case("mock_response_sync_stream", _kw(mock_response="Hi from the mock", stream=True), upstream=(), sync=True), + Case("mock_response_sync_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=(), sync=True), + Case("mock_response_sync_internal_error", _kw(mock_response="litellm.InternalServerError"), upstream=(), sync=True), + Case("mock_response_sync_bad_metadata", _kw(mock_response="Hi", metadata={"user_id": 123}), upstream=(), sync=True), + Case("mock_response_sync_empty_calls_the_provider", _kw(mock_response=""), sync=True), + Case("latest_alias", model=LATEST_ALIAS), + Case("latest_alias_bare_with_provider", _kw(custom_llm_provider="anthropic"), model="claude-opus-latest"), + Case("latest_alias_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), model=LATEST_ALIAS), + Case( + "latest_alias_upstream_400", + upstream=(_failure(400, "invalid_request_error", "max_tokens: too large"),), + model=LATEST_ALIAS, + ), + Case("tool_custom_field", _kw(tools=[dict(CUSTOM_TOOL)])), + Case( + "tool_custom_field_keeps_other_custom_keys", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": False, "note": "kept"}}]), + ), + Case( + "tool_custom_field_dropped_with_drop_params", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True, "note": "dropped"}}], drop_params=True), + ), + Case("tool_custom_field_not_an_object", _kw(tools=[{**WEATHER_TOOL, "custom": "opaque"}])), + Case( + "tool_custom_field_under_own_defer_loading", + _kw(tools=[{**WEATHER_TOOL, "defer_loading": False, "custom": {"defer_loading": True}}]), + ), Case( "stream_cost_injection", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), settings=MappingProxyType({"include_cost_in_streaming_usage": True}), ), - Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), + Case( + "stream_cost_injection_from_proxy_header", + _kw(stream=True, proxy_server_request=dict(COST_HEADER)), + upstream=(_stream(TOOL_STREAM, chunk_delay=0.01),), + ), + Case( + "cache_off_by_default", + settings=MappingProxyType({"cache": Cache(mode=CacheMode.default_off)}), + ), + Case( + "cache_skips_messages_call_type", + settings=MappingProxyType({"cache": Cache(supported_call_types=["acompletion"])}), + ), +) + +FALLBACK_CASES: Final = ( + Case("modify_params", _kw(max_tokens=1), settings=_MODIFY_PARAMS), + Case("modify_params_clamps_max_tokens", _kw(max_tokens=200_000), settings=_MODIFY_PARAMS), + Case( + "modify_params_clamps_streamed_max_tokens", + _kw(stream=True, max_tokens=200_000), + upstream=(_stream(TEXT_STREAM),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_clamps_tool_turns", + _kw(messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)], max_tokens=200_000), + upstream=(_ok(TOOL_USE_RESPONSE),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_clamps_thinking", + _kw(thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=200_000), + upstream=(_ok(THINKING_RESPONSE),), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_image_content", + _kw( + messages=[ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_PIXEL}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + max_tokens=200_000, + ), + settings=_MODIFY_PARAMS, + ), + Case( + "modify_params_unknown_model", + _kw(max_tokens=200_000), + model="anthropic/claude-parity-unknown", + settings=_MODIFY_PARAMS, + ), + Case("modify_params_hf_model_name", _kw(max_tokens=200_000, hf_model_name="gpt2"), settings=_MODIFY_PARAMS), + Case( + "modify_params_cache_write_without_read", + _kw(max_tokens=200_000, cache={"no-cache": True}), + upstream=(_ok(), _ok()), + settings=_MODIFY_PARAMS, + cache=Cache, + ), + Case( + "modify_params_clamped_call_never_hits_the_cache", + _kw(max_tokens=200_000), + upstream=(_ok(), _ok()), + settings=_MODIFY_PARAMS, + cache=Cache, + ), + Case( + "modify_params_unclamped_call_hits_the_cache", + _kw(max_tokens=10), + upstream=(_ok(), _ok()), + settings=_MODIFY_PARAMS, + cache=Cache, + ), + Case("cache_write_without_read", _kw(cache={"no-cache": True}), upstream=(_ok(), _ok()), cache=Cache), + Case("reply_with_wide_numbers", upstream=(ResponseSpec(body=WIDE_NUMBERS_REPLY),)), + Case("reply_with_a_wide_integer_cached", upstream=(ResponseSpec(body=WIDE_INTEGER_REPLY),), cache=Cache), + Case("reply_with_raw_leaf_marked_texts", upstream=(ResponseSpec(body=MARKED_TEXTS_REPLY),)), + Case("reply_with_raw_leaf_marked_texts_cached", upstream=(ResponseSpec(body=MARKED_TEXTS_REPLY),), cache=Cache), + Case( + "modify_params_clamped_stream_hits_the_cache", + _kw(max_tokens=200_000, stream=True), + upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), + settings=_MODIFY_PARAMS, + cache=Cache, + ), + Case( + "drop_params_drops_non_string_tool_result_texts", + _kw( + messages=_non_string_tool_result_texts(5, None, {"a": 1}, 8.5, 2**70, float("nan"), "kept"), + drop_params=True, + ), + ), + Case( + "global_drop_params_drops_every_tool_result_text", + _kw(messages=_non_string_tool_result_texts(5, None)), + settings=MappingProxyType({"drop_params": True}), + ), + Case( + "drop_params_keeps_string_tool_result_texts", + _kw(messages=_non_string_tool_result_texts("\uf8fflitellm-raw-json:1", "kept"), drop_params=True), + ), + Case("modify_params_sync", _kw(max_tokens=200_000), sync=True, settings=_MODIFY_PARAMS), Case("request_timeout_kwarg", _kw(request_timeout=30)), Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) -CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES) +INVALID_SIGNATURE: Final = "messages.1.content.0: Invalid `signature` in `thinking` block" +BOUND_ELSEWHERE: Final = f"{INVALID_SIGNATURE}. The block is bound to a different conversation." +THINKING: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "enabled", "budget_tokens": 2048}) +BINDING_ERRORS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}} +) +THINKING_ONLY_REPLAY: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "What is the capital of France?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "EmwKAhgBEgy"}, + {"type": "thinking", "thinking": "Paris.", "signature": "EqQBCkYIARgCKkB"}, + ], + }, + {"role": "assistant", "content": []}, + {"role": "user", "content": "And Germany?"}, +) + + +def _replay(**kwargs: object) -> Mapping[str, object]: + return _kw(messages=list(THINKING_REPLAY), thinking=dict(THINKING), max_tokens=4096, **kwargs) + + +def _rejected(message: str = INVALID_SIGNATURE, status: int = 400) -> ResponseSpec: + return _failure(status, "invalid_request_error", message) + + +THINKING_RETRY_CASES: Final = ( + Case("thinking_replay", _replay()), + Case("thinking_signature_rejected_then_retried", _replay(), upstream=(_rejected(), _ok())), + Case( + "thinking_signature_rejected_then_retried_stream", + _replay(stream=True), + upstream=(_rejected(), _stream(TEXT_STREAM)), + ), + Case( + "thinking_signature_missing_then_retried", + _replay(), + upstream=(_rejected("messages.1.content.0.thinking.signature.str: Input should be a valid string"), _ok()), + ), + Case( + "empty_thinking_rejected_then_retried_without_thinking_only_turns", + _kw(messages=list(THINKING_ONLY_REPLAY), thinking=dict(THINKING), max_tokens=4096), + upstream=(_rejected("messages.1.content.1.thinking: each thinking block must contain thinking"), _ok()), + ), + Case("thinking_signature_rejected_twice", _replay(), upstream=(_rejected(), _rejected())), + Case( + "thinking_signature_rejected_twice_stream", + _replay(stream=True), + upstream=(_rejected(), _rejected()), + ), + Case("thinking_signature_text_on_a_429_is_not_retried", _replay(), upstream=(_rejected(status=429),)), + Case( + "unrelated_400_is_not_retried", + _replay(), + upstream=(_rejected("max_tokens: 4096 > 64000, which is the maximum allowed"),), + ), + Case( + "binding_mismatch_retried_when_the_caller_did_not_pin_it", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "adaptive"}, max_tokens=4096), + model=OPUS_ADAPTIVE, + upstream=(_rejected(BOUND_ELSEWHERE), _ok()), + ), + Case( + "binding_mismatch_raised_when_the_caller_pinned_it", + _kw(messages=list(THINKING_REPLAY), thinking=dict(BINDING_ERRORS), max_tokens=4096), + model=OPUS_ADAPTIVE, + upstream=(_rejected(BOUND_ELSEWHERE),), + ), + Case( + "invalid_signature_retried_even_when_the_caller_pinned_binding", + _kw(messages=list(THINKING_REPLAY), thinking=dict(BINDING_ERRORS), max_tokens=4096), + model=OPUS_ADAPTIVE, + upstream=(_rejected(), _ok()), + ), +) + + +def _default_off_cache() -> Cache: + return Cache(mode=CacheMode.default_off) + + +def _completion_only_cache() -> Cache: + return Cache(supported_call_types=["acompletion"]) + + +CACHE_CASES: Final = ( + Case("cache_hit", upstream=(_ok(), _ok()), cache=Cache), + Case("cache_hit_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), cache=Cache), + Case("cache_no_cache", _kw(cache={"no-cache": True}), upstream=(_ok(), _ok()), cache=Cache), + Case("cache_no_store", _kw(cache={"no-store": True}), upstream=(_ok(), _ok()), cache=Cache), + Case("cache_caching_false", _kw(caching=False), upstream=(_ok(), _ok()), cache=Cache), + Case( + "cache_ttl_and_namespace", _kw(cache={"ttl": 60, "namespace": "parity"}), upstream=(_ok(), _ok()), cache=Cache + ), + Case("cache_default_off", upstream=(_ok(), _ok()), cache=_default_off_cache), + Case( + "cache_default_off_use_cache", + _kw(cache={"use-cache": True}), + upstream=(_ok(), _ok()), + cache=_default_off_cache, + ), + Case("cache_unsupported_call_type", upstream=(_ok(), _ok()), cache=_completion_only_cache), + Case("cache_hit_with_metadata", _kw(metadata={"user_id": "user-123"}), upstream=(_ok(), _ok()), cache=Cache), + Case("cache_sync", sync=True, upstream=(_ok(), _ok()), cache=Cache), + *( + Case(name, _kw(stream=True), upstream=(_raw_stream(*frames), _raw_stream(*frames)), cache=Cache) + for name, frames in ( + ("cache_stream_data_only_frames", tuple(f"data: {json.dumps(data)}" for _, data in TEXT_STREAM)), + ("cache_stream_event_without_space", _framed(TEXT_STREAM, "event:")), + ("cache_stream_done_sentinel", (*_framed(TEXT_STREAM), "data: [DONE]")), + ( + "cache_stream_without_output", + _framed( + _events(_message_start("msg_01ParityEmpty"), *_text_block(0, ""), *_message_end("end_turn", 1)) + ), + ), + ( + "cache_stream_error_after_stop", + ( + *_framed(TEXT_STREAM), + *_framed((("error", dict(_upstream_error("overloaded_error", "Overloaded"))),)), + ), + ), + ) + ), +) + +STRICT_REJECTION: Final = "Too many strict tools (21). The maximum number of strict tools is 20." +STRICT_TOOLS: Final[tuple[JsonValue, ...]] = ( + {**WEATHER_TOOL, "strict": True}, + {**WEATHER_TOOL, "name": "get_time", "strict": False}, + dict(WEATHER_TOOL, name="get_tide"), +) + + +def _strict(**kwargs: object) -> Mapping[str, object]: + return _kw(tools=list(STRICT_TOOLS), **kwargs) + + +STRICT_TOOLS_RETRY_CASES: Final = ( + Case( + "strict_tools_rejected_then_retried", _strict(drop_params=True), upstream=(_rejected(STRICT_REJECTION), _ok()) + ), + Case( + "strict_tools_rejected_then_retried_stream", + _strict(drop_params=True, stream=True), + upstream=(_rejected(STRICT_REJECTION), _stream(TEXT_STREAM)), + ), + Case( + "strict_tools_rejected_then_retried_under_global_drop_params", + _strict(), + upstream=(_rejected("The compiled grammar is too large"), _ok()), + settings=MappingProxyType({"drop_params": True}), + ), + Case( + "strict_tools_custom_field_rejected_then_retried", + _strict(drop_params="true"), + upstream=(_rejected("tools.0.custom.strict: Extra inputs are not permitted"), _ok()), + ), + Case( + "strict_tools_rejected_without_drop_params", + _strict(), + upstream=(_rejected(STRICT_REJECTION),), + settings=MappingProxyType({"drop_params": False}), + ), + Case( + "strict_tools_rejected_under_a_truthy_global_drop_params_that_is_not_true", + _strict(), + upstream=(_rejected(STRICT_REJECTION),), + settings=MappingProxyType({"drop_params": 1}), + ), + Case( + "strict_rejection_without_a_strict_tool", + _kw(tools=[{**WEATHER_TOOL, "strict": False}], drop_params=True), + upstream=(_rejected(STRICT_REJECTION),), + ), + Case( + "strict_rejection_text_on_a_429_is_not_retried", + _strict(drop_params=True), + upstream=(_rejected(STRICT_REJECTION, status=429),), + ), + Case( + "strict_tools_rejected_twice", + _strict(drop_params=True), + upstream=(_rejected(STRICT_REJECTION), _rejected("Grammar compilation timed out")), + ), + Case( + "strict_tools_rejected_twice_stream", + _strict(drop_params=True, stream=True), + upstream=(_rejected(STRICT_REJECTION), _rejected(STRICT_REJECTION)), + ), + Case( + "strict_tools_and_thinking_strict_rejected_then_retried", + _replay(tools=list(STRICT_TOOLS), drop_params=True), + upstream=(_rejected(STRICT_REJECTION), _ok()), + ), + Case( + "strict_tools_and_thinking_signature_rejected_then_retried", + _replay(tools=list(STRICT_TOOLS), drop_params=True), + upstream=(_rejected(), _ok()), + ), + Case( + "strict_tools_and_thinking_each_rejected_spends_the_one_retry", + _replay(tools=list(STRICT_TOOLS), drop_params=True), + upstream=(_rejected(STRICT_REJECTION), _rejected()), + ), + Case( + "strict_tools_and_thinking_signature_then_strict_rejected", + _replay(tools=list(STRICT_TOOLS), drop_params=True), + upstream=(_rejected(), _rejected(STRICT_REJECTION), _ok()), + ), + Case( + "one_error_naming_both_drops_thinking_and_strict_without_drop_params", + _replay(tools=list(STRICT_TOOLS)), + upstream=(_rejected(f"{INVALID_SIGNATURE}. {STRICT_REJECTION}"), _ok()), + ), +) + +THREADS_BETA: Final = "message-threads-2026-08-12" +ADVISOR_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-8"} +) +THREAD_CREATE: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "create"}) +THREAD_CONTINUE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "continue", "previous_message_id": "msg_01ParityText"} +) +ADVISOR_TURNS: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "Build a worker pool."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "server_tool_use", "id": "srvtoolu_01Parity", "name": "advisor", "input": {}}, + { + "type": "advisor_tool_result", + "tool_use_id": "srvtoolu_01Parity", + "content": {"type": "advisor_result", "text": "Use channels."}, + }, + {"type": "text", "text": "Here is the implementation."}, + ], + }, + {"role": "user", "content": "Now add a timeout.", "output_config": {"effort": "low"}}, + {"role": "system", "content": [{"type": "tool_addition", "tool": {"type": "tool_reference", "name": "ping"}}]}, + {"role": "user", "content": "Go on."}, +) +THREAD_NOT_FOUND: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "error", + "error": { + "type": "not_found_error", + "message": "No thread state was found for the requested `previous_message_id`.", + "details": {"error_code": "thread_not_found"}, + }, + "request_id": "req_01Parity", + } +) + + +def _thread_not_found() -> ResponseSpec: + return ResponseSpec(body=dict(THREAD_NOT_FOUND), status=404) + + +THREAD_CASES: Final = ( + Case("thread_create", _kw(thread=dict(THREAD_CREATE), system="Answer in one word.")), + Case("thread_create_with_advisor_tool", _kw(thread=dict(THREAD_CREATE), tools=[dict(ADVISOR_TOOL)])), + Case("thread_create_scans_history", _kw(thread=dict(THREAD_CREATE), messages=list(ADVISOR_TURNS))), + Case("thread_continue", _kw(thread=dict(THREAD_CONTINUE))), + Case( + "thread_continue_keeps_history_and_skips_its_betas", + _kw(thread=dict(THREAD_CONTINUE), messages=list(ADVISOR_TURNS)), + ), + Case( + "thread_continue_keeps_param_betas", + _kw( + thread=dict(THREAD_CONTINUE), + messages=list(ADVISOR_TURNS), + compaction={}, + context_management={"edits": [{"type": "compact_20260112"}]}, + ), + ), + Case("thread_continue_without_previous_message_id", _kw(thread={"type": "continue"}, messages=list(ADVISOR_TURNS))), + Case( + "thread_continue_with_a_malformed_previous_message_id", + _kw(thread={"type": "continue", "previous_message_id": 7}, messages=list(ADVISOR_TURNS)), + ), + Case("thread_of_unknown_shape", _kw(thread="opaque", messages=list(ADVISOR_TURNS))), + Case("thread_none_is_left_out", _kw(thread=None)), + Case( + "thread_beside_a_client_beta", + _kw(thread=dict(THREAD_CONTINUE), extra_headers={"anthropic-beta": f"advisor-tool-2026-03-01,{THREADS_BETA}"}), + ), + Case( + "thread_dropped_by_additional_drop_params", _kw(thread=dict(THREAD_CREATE), additional_drop_params=["thread"]) + ), + Case("thread_create_stream", _kw(thread=dict(THREAD_CREATE), stream=True), upstream=(_stream(TEXT_STREAM),)), + Case( + "thread_continue_stream", + _kw(thread=dict(THREAD_CONTINUE), stream=True, messages=list(ADVISOR_TURNS)), + upstream=(_stream(TEXT_STREAM),), + ), + Case("thread_create_sync", _kw(thread=dict(THREAD_CREATE)), sync=True), + Case( + "thread_continue_sync_stream", + _kw(thread=dict(THREAD_CONTINUE), stream=True), + upstream=(_stream(TEXT_STREAM),), + sync=True, + ), + Case("thread_not_found", _kw(thread=dict(THREAD_CONTINUE)), upstream=(_thread_not_found(),)), + Case("thread_not_found_stream", _kw(thread=dict(THREAD_CONTINUE), stream=True), upstream=(_thread_not_found(),)), + Case("thread_not_found_sync", _kw(thread=dict(THREAD_CONTINUE)), upstream=(_thread_not_found(),), sync=True), + Case( + "thread_continue_retried_after_a_rejected_signature", + _replay(thread=dict(THREAD_CONTINUE)), + upstream=(_rejected(), _ok()), + ), +) + + +class _CallerHandler(AsyncHTTPHandler): + def __deepcopy__(self, memo: dict[int, object]) -> "_CallerHandler": + return _CallerHandler() + + +class _UncheckedHandler(_CallerHandler): + def __deepcopy__(self, memo: dict[int, object]) -> "_UncheckedHandler": + return _UncheckedHandler() + + async def post(self, url: str, **kwargs: object) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the test double ignores the rest of the post + return httpx.Response(503, text="unavailable") + + +class _UnreadHandler(_CallerHandler): + def __deepcopy__(self, memo: dict[int, object]) -> "_UnreadHandler": + return _UnreadHandler() + + async def post(self, url: str, **kwargs: object) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the test double ignores the rest of the post + return httpx.Response(500, stream=httpx.ByteStream(b"unread"), request=httpx.Request("POST", url)) + + +class _CallerHttpxClient(httpx.AsyncClient): + def __deepcopy__(self, memo: dict[int, object]) -> "_CallerHttpxClient": + return _CallerHttpxClient() + + +class _CallerSyncHandler(HTTPHandler): + def __deepcopy__(self, memo: dict[int, object]) -> "_CallerSyncHandler": + return _CallerSyncHandler() + + +def _caller(**kwargs: object) -> Mapping[str, object]: + return _kw(client=_CallerHandler(), **kwargs) + + +_TLS_LEVEL: Final = _tls(ssl_security_level="DEFAULT@SECLEVEL=1") +_TLS_CURVE: Final = _tls(ssl_ecdh_curve="secp521r1") + +TRANSPORT_CASES: Final = ( + Case("caller_client", _caller()), + Case("caller_client_stream", _caller(stream=True), upstream=(_stream(TEXT_STREAM),)), + Case( + "caller_client_upstream_429", + _caller(), + upstream=( + _failure(429, "rate_limit_error", "slow down", **{"retry-after": "7", "request-id": "req_01Parity"}), + ), + ), + Case( + "caller_client_stream_upstream_429", + _caller(stream=True), + upstream=(_failure(429, "rate_limit_error", "slow down", **{"retry-after": "3"}),), + ), + Case( + "caller_client_non_json_error", + _caller(), + upstream=(ResponseSpec(body=b"bad gateway", status=502),), + ), + Case("caller_client_invalid_json_success", _caller(), upstream=(ResponseSpec(body=b"not json", status=200),)), + Case("caller_client_timeout", _caller(timeout=0.5), upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), delay=1.5),)), + Case( + "caller_client_stream_read_stalls_past_timeout", + _caller(stream=True, timeout=0.5), + upstream=(_stream(SLOW_STREAM, chunk_delay=1.0),), + ), + Case( + "caller_client_thinking_signature_rejected_then_retried", + _kw(**_replay(), client=_CallerHandler()), + upstream=(_rejected(), _ok()), + ), + Case( + "caller_client_strict_tools_rejected_then_retried_stream", + _strict(drop_params=True, stream=True, client=_CallerHandler()), + upstream=(_rejected(STRICT_REJECTION), _stream(TEXT_STREAM)), + ), + Case("caller_client_sync", _caller(), sync=True), + Case("caller_client_post_returns_unchecked_failure", _kw(client=_UncheckedHandler()), upstream=()), + Case("caller_client_post_returns_unread_failure", _kw(client=_UnreadHandler(), stream=True), upstream=()), + Case("caller_httpx_client_is_ignored", _kw(client=_CallerHttpxClient())), + Case("caller_sync_handler_is_ignored", _kw(client=_CallerSyncHandler())), + Case("tls_security_level", settings=_TLS_LEVEL), + Case("tls_security_level_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), settings=_TLS_LEVEL), + Case("tls_ecdh_curve", settings=_TLS_CURVE), + Case("tls_ecdh_curve_x448", settings=_tls(ssl_ecdh_curve="X448")), + Case("tls_ecdh_curve_upstream_500", upstream=(_failure(500, "api_error", "internal"),), settings=_TLS_CURVE), + Case( + "tls_security_level_from_environment", + env=MappingProxyType({"SSL_SECURITY_LEVEL": "DEFAULT@SECLEVEL=1"}), + settings=_tls(), + ), + Case("tls_security_level_python_cannot_apply", upstream=(), settings=_tls(ssl_security_level="NOT-A-CIPHER")), + Case("tls_ecdh_curve_nist_name", settings=_tls(ssl_ecdh_curve="P-256")), + Case("tls_ecdh_curve_padded_lowercase", settings=_tls(ssl_ecdh_curve=" x25519 ")), + Case("negative_timeout", _kw(timeout=-1), upstream=(ResponseSpec(body=dict(TEXT_RESPONSE), delay=0.3),)), + Case( + "negative_timeout_stream", + _kw(stream=True, timeout=-0.5), + upstream=(replace(_stream(TEXT_STREAM), delay=0.3),), + ), + Case( + "zero_timeout_httpx_transport", + _kw(timeout=0), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + ), + Case( + "negative_timeout_httpx_transport", + _kw(timeout=-1), + env=MappingProxyType({"DISABLE_AIOHTTP_TRANSPORT": "true"}), + settings=_tls(), + ), +) + + +class _RewritingHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + kwargs["tools"] = [dict(WEATHER_TOOL)] + kwargs["tool_choice"] = {"type": "tool", "name": "get_weather"} + kwargs["temperature"] = 0.25 + kwargs["top_k"] = 7 + kwargs["system"] = "Answer as a forecaster." + kwargs["stop_sequences"] = ["STOP"] + kwargs["metadata"] = {"user_id": "hooked-user"} + kwargs["extra_headers"] = {"x-hooked": "1"} + return kwargs + + +class _InPlaceHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + messages.append({"role": "assistant", "content": [{"type": "text", "text": ""}]}) + messages.append({"role": "user", "content": "And Germany?"}) + return None + + +class _UnstreamingHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "stream": not kwargs["stream"]} + + +class _MockingHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "mock_response": "Hi from the hook"} + + +class _RaisingHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + raise ValueError("the hook refused the request") + + +class _ModelHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "model": OPUS_ADAPTIVE} + + +class _CredentialHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "api_key": "sk-ant-hooked"} + + +def _instance_hook() -> CustomLogger: + async def hook(model: str, messages: list[dict[str, object]], kwargs: dict[str, object]) -> dict[str, object]: + return {**kwargs, "temperature": 0.25, "metadata": {"user_id": "instance-hook"}} + + logger: Final = CustomLogger() + logger.async_pre_request_hook = hook # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: an instance-level override, as a caller may install one + return logger + + +PRE_REQUEST_HOOK_CASES: Final = ( + Case("pre_request_hook_rewrites_the_request", callbacks=(_RewritingHook,), upstream=(_ok(TOOL_USE_RESPONSE),)), + Case("pre_request_hook_edits_messages_in_place", callbacks=(_InPlaceHook,)), + Case( + "pre_request_hook_after_prompt_caching", + _kw(cache_control_injection_points=[{"location": "message", "role": "user"}]), + callbacks=(_InPlaceHook,), + ), + Case( + "pre_request_hook_unstreams", + _kw(stream=True), + callbacks=(_UnstreamingHook,), + ), + Case( + "pre_request_hook_unstreams_behind_a_cache", + _kw(stream=True), + callbacks=(_UnstreamingHook,), + upstream=(_ok(), _ok()), + cache=Cache, + ), + Case("pre_request_hook_streams", callbacks=(_UnstreamingHook,), upstream=(_stream(TEXT_STREAM),)), + Case("pre_request_hook_sets_a_mock", callbacks=(_MockingHook,), upstream=()), + Case("pre_request_hook_raises", callbacks=(_RaisingHook,), upstream=()), + Case("pre_request_hook_set_on_the_instance", callbacks=(_instance_hook,)), + Case("pre_request_hook_names_the_model", callbacks=(_ModelHook,), upstream=()), + Case("pre_request_hook_names_a_credential", callbacks=(_CredentialHook,), upstream=()), + Case( + "pre_request_hook_two_hooks_chain", callbacks=(_RewritingHook, _InPlaceHook), upstream=(_ok(TOOL_USE_RESPONSE),) + ), + Case( + "pre_request_hook_behind_a_cache_hit", + callbacks=(_RewritingHook,), + upstream=(_ok(TOOL_USE_RESPONSE), _ok(TOOL_USE_RESPONSE)), + cache=Cache, + ), +) + +_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) +_VOLATILE_HOOK_KWARGS: Final = frozenset({"litellm_call_id", "litellm_logging_obj"}) + + +def _content(response: object) -> list[dict[str, object]]: + if not isinstance(response, dict): + return [] + return _OBJECTS.validate_python(_OBJECT.validate_python(response).get("content") or []) + + +def _tool_calls(response: object) -> list[dict[str, object]]: + return [block for block in _content(response) if block.get("type") == "tool_use"] + + +class _FollowUpOnToolUse(CustomLogger): + """Answers each tool call with a canned result in one follow-up turn, recording what the hooks were handed.""" + + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger takes untyped kwargs + self.seen: list[Mapping[str, object]] = [] + + async def async_should_run_agentic_loop( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + response: object, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict[str, object], + ) -> tuple[bool, dict[str, object]]: + calls: Final = _tool_calls(response) + self.seen.append( + { + "hook": "should_run", + "response": _OBJECT.validate_python({**response, "id": ""}) if isinstance(response, dict) else None, + "model": model, + "messages": _OBJECTS.validate_python(messages), + "tools": tools, + "stream": stream, + "custom_llm_provider": custom_llm_provider, + "kwargs": sorted(name for name in kwargs if name not in _VOLATILE_HOOK_KWARGS), + "api_key": kwargs.get("api_key"), + } + ) + return bool(calls), {"tool_calls": calls} + + async def async_build_agentic_loop_plan( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + tools: dict[str, object], + model: str, + messages: list[dict[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: dict[str, object], + logging_obj: Logging, + stream: bool, + kwargs: dict[str, object], + ) -> AgenticLoopPlan: + self.seen.append( + { + "hook": "build_plan", + "optional_params": anthropic_messages_optional_request_params, + "config": type(anthropic_messages_provider_config).__name__, + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params"), + "kwargs": sorted(name for name in kwargs if name not in _VOLATILE_HOOK_KWARGS), + } + ) + return self.plan(messages, response) + + def plan(self, messages: list[dict[str, object]], response: object) -> AgenticLoopPlan: + answers: Final = [ + {"type": "tool_result", "tool_use_id": call["id"], "content": "18C, sunny"} + for call in _tool_calls(response) + ] + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[ + *messages, + {"role": "assistant", "content": _content(response)}, + {"role": "user", "content": answers}, + ] + ), + ) + + +class _HeldBack(_FollowUpOnToolUse): + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({"get_weather"}) + + +class _Terminates(_FollowUpOnToolUse): + def plan(self, messages: list[dict[str, object]], response: object) -> AgenticLoopPlan: + return AgenticLoopPlan(terminate=True, stop_reason="answered") + + +class _Overrides(_FollowUpOnToolUse): + def plan(self, messages: list[dict[str, object]], response: object) -> AgenticLoopPlan: + return AgenticLoopPlan(response_override={**TEXT_RESPONSE, "id": "msg_01Override"}) + + +class _Fails(_FollowUpOnToolUse): + async def async_build_agentic_loop_plan( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + tools: dict[str, object], + model: str, + messages: list[dict[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: dict[str, object], + logging_obj: Logging, + stream: bool, + kwargs: dict[str, object], + ) -> AgenticLoopPlan: + raise ValueError("the agentic plan failed") + + +class _WebSearch(WebSearchInterceptionLogger): + """The stock web search interception, recording what its pre-request hook leaves for its agentic hook.""" + + def __init__(self, enabled_providers: list[str] | None = None, max_agentic_loops: int | None = None) -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # the logger annotates bare containers + enabled_providers=list(enabled_providers) if enabled_providers is not None else None, + max_agentic_loops=max_agentic_loops, + ) + self.seen: list[Mapping[str, object]] = [] + + async def async_should_run_agentic_loop( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + response: object, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict[str, object], + ) -> tuple[bool, dict[str, object]]: + self.seen.append( + { + "model": model, + "messages": _OBJECTS.validate_python(messages), + "tools": tools, + "stream": stream, + "custom_llm_provider": custom_llm_provider, + "kwargs": { + name: kwargs[name] + for name in sorted(kwargs) + if name == "max_agentic_loops" or name.startswith("_websearch_interception") + }, + "kwarg_names": sorted(name for name in kwargs if name not in _VOLATILE_HOOK_KWARGS), + } + ) + return await super().async_should_run_agentic_loop( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the logger annotates bare containers + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + + +class _DeploymentHooks(CustomLogger): + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger takes untyped kwargs + self.seen: list[Mapping[str, object]] = [] + + async def async_post_call_success_deployment_hook( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, request_data: dict[str, object], response: object, call_type: object + ) -> object: + self.seen.append({"hook": "success", "call_type": str(call_type)}) + return None + + async def async_post_call_failure_deployment_hook( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, request_data: Mapping[str, object], exception: Exception, call_type: object + ) -> None: + self.seen.append({"hook": "failure", "error": type(exception).__name__, "call_type": str(call_type)}) + + +def _loggers(*kinds: Callable[[], CustomLogger]) -> Callable[[], Sequence[CustomLogger]]: + return lambda: tuple(kind() for kind in kinds) + + +def _web_search() -> CustomLogger: + return _WebSearch(enabled_providers=["anthropic"], max_agentic_loops=2) + + +async def _search(*args: object, **kwargs: object) -> SearchResponse: # kwargs-ok: litellm.asearch's call shape + return SearchResponse.model_validate( + {"results": [{"title": "Paris", "url": "https://example.com/paris", "snippet": "Paris is the capital."}]} + ) + + +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 2} +) +WEB_SEARCH_CALL_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "id": "msg_01ParitySearch", + "content": [ + {"type": "tool_use", "id": "toolu_01Search", "name": "litellm_web_search", "input": {"query": "capital"}} + ], + "stop_reason": "tool_use", + } +) +_WEB_SEARCH_SETTINGS: Final[Mapping[str, object]] = MappingProxyType({"asearch": _search}) +_WEATHER: Final = _kw(tools=[dict(WEATHER_TOOL)]) + +AGENTIC_CASES: Final = ( + Case("agentic_loop_without_tool_use", _WEATHER, loggers=_loggers(_FollowUpOnToolUse)), + Case( + "agentic_loop_follow_up", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE), _ok()), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_follow_up_with_request_params", + _kw( + tools=[dict(WEATHER_TOOL)], + system="Answer in one word.", + temperature=0.3, + metadata={"user_id": "user-123"}, + thinking={"type": "enabled", "budget_tokens": 2048}, + max_tokens=4096, + messages=[{"role": "user", "content": "Weather?", "provider_specific_fields": {"x": 1}}], + ), + upstream=(_ok(TOOL_USE_RESPONSE), _ok()), + loggers=_loggers(_FollowUpOnToolUse), + ), + *( + Case( + f"agentic_loop_follow_up_keeps_{name}", + _kw(**{"tools": [dict(WEATHER_TOOL)], **params}), + upstream=(_ok(TOOL_USE_RESPONSE), _ok()), + model=model, + loggers=_loggers(_FollowUpOnToolUse), + ) + for name, params, model in ( + ( + "output_format", + { + "output_format": { + "type": "json_schema", + "schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + }, + SONNET, + ), + ("reasoning_effort", {"reasoning_effort": "high", "max_tokens": 8192}, SONNET), + ("output_config_effort", {"output_config": {"effort": "medium"}}, OPUS_ADAPTIVE), + ( + "injected_cache_points", + {"cache_control_injection_points": [{"location": "message", "role": "user"}]}, + SONNET, + ), + ("context_management", {"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}, SONNET), + ("speed", {"speed": "fast"}, OPUS_FAST), + ("service_tier", {"service_tier": "auto"}, SONNET), + ("inference_geo", {"inference_geo": "us"}, SONNET), + ( + "container", + {"container": {"skills": [{"type": "anthropic", "skill_id": "pptx", "version": "latest"}]}}, + SONNET, + ), + ("custom_tool", {"tools": [{**WEATHER_TOOL, "type": "custom"}]}, SONNET), + ("metadata_extras", {"metadata": {"user_id": "user-123", "session": "dropped"}}, SONNET), + ) + ), + Case( + "agentic_loop_prices_the_final_round", + _kw(tools=[dict(WEATHER_TOOL)]), + upstream=(_ok(TOOL_USE_RESPONSE), _ok({**TEXT_RESPONSE, "usage": {"input_tokens": 900, "output_tokens": 300}})), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_prices_the_final_round_stream", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=( + _stream(TOOL_STREAM), + _stream( + _events(_message_start("msg_01ParityLate"), *_text_block(0, "Paris"), *_message_end("end_turn", 300)) + ), + ), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_follow_up_stream", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_stream(TOOL_STREAM), _stream(TEXT_STREAM)), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_follow_up_stream_in_one_write", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_burst(TOOL_STREAM), _burst(TEXT_STREAM)), + loggers=_loggers(_FollowUpOnToolUse), + partitioned=True, + ), + Case( + "agentic_loop_stream_without_tool_use_in_one_write", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_burst(TEXT_STREAM),), + loggers=_loggers(_FollowUpOnToolUse), + partitioned=True, + ), + Case( + "agentic_loop_stream_without_tool_use", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_holds_back_a_server_fulfilled_tool", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_stream(TOOL_STREAM), _stream(TEXT_STREAM)), + loggers=_loggers(_HeldBack), + ), + Case( + "agentic_loop_terminates", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE),), + loggers=_loggers(_Terminates), + ), + Case( + "agentic_loop_response_override", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE),), + loggers=_loggers(_Overrides), + ), + Case( + "agentic_loop_plan_fails", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE),), + loggers=_loggers(_Fails), + ), + Case( + "agentic_loop_follow_up_fails", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE), _failure(529, "overloaded_error", "Overloaded")), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_plan_fails_runs_the_failure_deployment_hook_once", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE),), + loggers=_loggers(_Fails, _DeploymentHooks), + ), + Case( + "agentic_loop_follow_up_fails_runs_the_failure_deployment_hook_per_round", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE), _failure(529, "overloaded_error", "Overloaded")), + loggers=_loggers(_FollowUpOnToolUse, _DeploymentHooks), + ), + Case( + "agentic_loop_follow_up_runs_the_success_deployment_hook_per_round", + _WEATHER, + upstream=(_ok(TOOL_USE_RESPONSE), _ok()), + loggers=_loggers(_FollowUpOnToolUse, _DeploymentHooks), + ), + Case( + "pre_request_hook_raises_runs_the_failure_deployment_hook_once", + callbacks=(_RaisingHook,), + upstream=(), + loggers=_loggers(_DeploymentHooks), + ), + Case( + "agentic_loop_refused_past_max_loops", + _kw(tools=[dict(WEATHER_TOOL)], max_agentic_loops=1), + upstream=(_ok(TOOL_USE_RESPONSE), _ok(TOOL_USE_RESPONSE)), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "agentic_loop_upstream_error_skips_the_hooks", + _WEATHER, + upstream=(_failure(400, "invalid_request_error", "bad"),), + loggers=_loggers(_FollowUpOnToolUse), + ), + Case( + "web_search_interception", + _kw(tools=[dict(WEB_SEARCH_TOOL)]), + upstream=(_ok(WEB_SEARCH_CALL_RESPONSE), _ok()), + settings=_WEB_SEARCH_SETTINGS, + loggers=_loggers(_web_search), + ), + Case( + "web_search_interception_converted_stream", + _kw(stream=True, tools=[dict(WEB_SEARCH_TOOL)]), + upstream=(_ok(WEB_SEARCH_CALL_RESPONSE), _ok()), + settings=_WEB_SEARCH_SETTINGS, + loggers=_loggers(_web_search), + ), + Case( + "web_search_interception_without_a_search_call", + _kw(stream=True, tools=[dict(WEB_SEARCH_TOOL)]), + settings=_WEB_SEARCH_SETTINGS, + loggers=_loggers(_web_search), + ), + Case( + "web_search_interception_for_a_provider_it_skips", + _kw(tools=[dict(WEB_SEARCH_TOOL)]), + loggers=_loggers(_WebSearch), + ), + Case( + "web_search_interception_rewrites_the_messages_it_hands_the_loop", + _kw( + stream=True, + tools=[dict(WEB_SEARCH_TOOL)], + messages=[ + {"role": "user", "content": [{"type": "text", "text": ""}, {"type": "text", "text": "Capital?"}]} + ], + ), + upstream=(_ok(WEB_SEARCH_CALL_RESPONSE), _ok()), + settings=_WEB_SEARCH_SETTINGS, + loggers=_loggers(_web_search), + ), +) + +CLASSIFIER_MODEL: Final = "openai/parity-classifier" +CLASSIFIER_CALL_TYPE: Final = "acompletion" +NON_CLAUDE: Final = "anthropic/parity-tool-model" +DANGEROUS_TOOL_USE: Final[tuple[JsonValue, ...]] = ( + { + "type": "dangerous_tool_use", + "classifier_context": { + "permission_mode": "auto", + "platform": "linux", + "trusted_directories": ["/repo"], + "git_state": {"cwd": "/repo", "branch": "main", "dirty": True}, + "ignored": "not rendered", + }, + }, +) +FLAGGED: Final = json.dumps({"verdicts": {"toolu_01Parity": {"flagged": True, "explanation": "Sends the weather."}}}) +CLEARED: Final = json.dumps({"verdicts": {"toolu_01Parity": {"flagged": False}}}) +CUT_OFF_TOOL_STREAM: Final = (*TOOL_STREAM[:-2], *_message_end("max_tokens", 22)) +UPSTREAM_ANSWERED: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TOOL_USE_RESPONSE, + "safeguard_results": [{"type": "dangerous_tool_use", "status": {"type": "unsupported"}}], + } +) + + +def _safeguarded(**kwargs: object) -> Mapping[str, object]: + return _kw(safeguards=list(DANGEROUS_TOOL_USE), messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)], **kwargs) + + +SAFEGUARDS_CASES: Final = ( + Case( + "safeguards_classifier_flags_tool_use", + _safeguarded(litellm_metadata={"user_api_key_end_user_id": "end-user-1", "user_api_key": "hashed-key"}), + upstream=(_ok(TOOL_USE_RESPONSE),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_classifier_clears_tool_use", + _safeguarded(), + upstream=(_ok(TOOL_USE_RESPONSE),), + model=NON_CLAUDE, + classifier=CLEARED, + ), + Case( + "safeguards_classifier_unparsable_reply", + _safeguarded(), + upstream=(_ok(TOOL_USE_RESPONSE),), + model=NON_CLAUDE, + classifier="I cannot tell.", + ), + Case( + "safeguards_classifier_rejects_long_input", + _safeguarded(), + upstream=(_ok(TOOL_USE_RESPONSE),), + model=NON_CLAUDE, + classifier="litellm.ContextWindowExceededError", + ), + Case("safeguards_classifier_without_tool_use", _safeguarded(), model=NON_CLAUDE, classifier=FLAGGED), + Case( + "safeguards_classifier_upstream_answered", + _safeguarded(), + upstream=(_ok(UPSTREAM_ANSWERED),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_classifier_stream_tool_use", + _safeguarded(stream=True), + upstream=(_stream(TOOL_STREAM),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_classifier_stream_cut_off", + _safeguarded(stream=True), + upstream=(_stream(CUT_OFF_TOOL_STREAM),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_classifier_stream_text", + _safeguarded(stream=True), + upstream=(_stream(TEXT_STREAM),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_classifier_upstream_failure", + _safeguarded(), + upstream=(_failure(529, "overloaded_error", "Overloaded"),), + model=NON_CLAUDE, + classifier=FLAGGED, + ), + Case( + "safeguards_claude_model_is_not_classified", + _safeguarded(), + upstream=(_ok(TOOL_USE_RESPONSE),), + classifier=FLAGGED, + ), + Case("safeguards_classifier_sync", _safeguarded(), model=NON_CLAUDE, sync=True, classifier=FLAGGED), +) + +INTERCEPTED_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "parity_intercept", "name": "intercept"}) +_ECHOED_KWARGS: Final = ( + "model", + "messages", + "tools", + "stream", + "max_tokens", + "custom_llm_provider", + "api_key", + "metadata", + "system", + "temperature", + "tool_choice", +) + + +class _Intercepting(MessagesInterceptor): + def can_handle(self, tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "anthropic" and any( + tool.get("type") == INTERCEPTED_TOOL["type"] for tool in tools or () + ) + + +class _Answering(_Intercepting): + async def handle(self, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the interceptor contract's keyword call + echoed: Final = {name: kwargs.get(name) for name in _ECHOED_KWARGS} + answer: Final = { + **TEXT_RESPONSE, + "id": "msg_01ParityIntercepted", + "content": [{"type": "text", "text": json.dumps(echoed, sort_keys=True, default=repr)}], + "extra_keys": sorted(name for name in kwargs if name not in _ECHOED_KWARGS), + } + if kwargs.get("stream"): + return FakeAnthropicMessagesStreamIterator( + cast(AnthropicMessagesResponse, answer) + ) # cast-ok: a synthetic Messages response + return answer + + +class _Delegating(_Intercepting): + async def handle(self, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the interceptor contract's keyword call + from litellm.messages import anthropic_messages + + tools: Final = cast(list[Mapping[str, object]], kwargs["tools"]) # cast-ok: the caller's tools + return await anthropic_messages( + model=kwargs["model"], + messages=kwargs["messages"], + max_tokens=kwargs["max_tokens"], + tools=[tool for tool in tools if tool.get("type") != INTERCEPTED_TOOL["type"]], + api_key=kwargs["api_key"], + api_base=kwargs["api_base"], + stream=kwargs["stream"], + ) + + +class _Refusing(_Intercepting): + async def handle(self, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the interceptor contract's keyword call + raise ValueError("the interceptor refused the call") + + +def _intercepted(**kwargs: object) -> Mapping[str, object]: + return _kw(tools=[dict(INTERCEPTED_TOOL), dict(WEATHER_TOOL)], **kwargs) + + +class _InterceptedToolHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "tools": [dict(INTERCEPTED_TOOL), *cast(list[object], kwargs.get("tools") or [])]} + + +class _UninterceptedToolHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "tools": [dict(WEATHER_TOOL)]} + + +class _DeploymentHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger takes untyped kwargs + self.seen: list[Mapping[str, object]] = [] + + async def async_pre_call_deployment_hook(self, kwargs: dict[str, object], call_type: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self.seen.append({"logging_obj": type(kwargs.get("litellm_logging_obj")).__name__, "call_type": str(call_type)}) + + +class _DeploymentResponseRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger takes untyped kwargs + self.seen: list[Mapping[str, object]] = [] + + async def async_post_call_success_deployment_hook( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, request_data: dict[str, object], response: object, call_type: object + ) -> object: + self.seen.append({"response": type(response).__name__, "call_type": str(call_type)}) + return None + + +DEPLOYMENT_HOOK_CASES: Final = ( + Case( + "deployment_hook_sees_the_stream_of_a_string_flag", + _kw(stream="true"), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentResponseRecorder), + ), + Case( + "deployment_hook_sees_the_stream_of_an_integer_flag", + _kw(stream=1), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentResponseRecorder), + ), + Case( + "deployment_hook_sees_the_stream_of_a_bool_flag", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentResponseRecorder), + ), + Case( + "deployment_hook_sees_a_stream_a_hook_made", + upstream=(_stream(TEXT_STREAM),), + callbacks=(_UnstreamingHook,), + loggers=_loggers(_DeploymentResponseRecorder), + ), + Case( + "deployment_hook_sees_a_stream_a_hook_made_on_an_integer", + upstream=(_stream(TEXT_STREAM),), + callbacks=(_IntegerStreamHook,), + loggers=_loggers(_DeploymentResponseRecorder), + ), + Case("deployment_hook_sees_the_logging_object", loggers=_loggers(_DeploymentHookRecorder)), + Case( + "deployment_hook_sees_the_logging_object_stream", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentHookRecorder), + ), +) + +INTERCEPTOR_CASES: Final = ( + Case("interceptor_answers", _intercepted(), upstream=(), interceptor=_Answering), + Case( + "pre_request_hook_adds_a_tool_an_interceptor_owns", + _kw(tools=[dict(WEATHER_TOOL)]), + upstream=(), + callbacks=(_InterceptedToolHook,), + interceptor=_Answering, + ), + Case( + "pre_request_hook_adds_a_tool_an_interceptor_owns_stream", + _kw(tools=[dict(WEATHER_TOOL)], stream=True), + upstream=(), + callbacks=(_InterceptedToolHook,), + interceptor=_Answering, + ), + Case( + "pre_request_hook_unstreams_a_call_an_interceptor_owns", + _intercepted(stream=True), + upstream=(), + callbacks=(_UnstreamingHook,), + interceptor=_Answering, + ), + Case( + "pre_request_hook_removes_the_tool_an_interceptor_owns", + _intercepted(), + callbacks=(_UninterceptedToolHook,), + interceptor=_Answering, + ), + Case( + "interceptor_answers_with_system_and_metadata", + _intercepted(system="Be brief.", metadata={"user_id": "user-123"}, temperature=0.5), + upstream=(), + interceptor=_Answering, + ), + Case("interceptor_answers_stream", _intercepted(stream=True), upstream=(), interceptor=_Answering), + Case( + "interceptor_wins_over_mock_response", + _intercepted(mock_response="Hi from the mock"), + upstream=(), + interceptor=_Answering, + ), + Case( + "interceptor_answers_bare_model", + _intercepted(custom_llm_provider="anthropic"), + upstream=(), + model="claude-sonnet-4-5", + interceptor=_Answering, + ), + Case("interceptor_delegates_to_the_public_api", _intercepted(), interceptor=_Delegating), + Case( + "interceptor_delegates_a_stream", + _intercepted(stream=True), + upstream=(_stream(TEXT_STREAM),), + interceptor=_Delegating, + ), + Case("interceptor_raises", _intercepted(), upstream=(), interceptor=_Refusing), + Case("interceptor_is_not_consulted_on_a_sync_call", _intercepted(), upstream=(), sync=True, interceptor=_Answering), + Case( + "pre_request_hook_is_not_run_on_a_sync_call_an_interceptor_owns", + _intercepted(), + upstream=(), + sync=True, + callbacks=(_RaisingHook,), + interceptor=_Answering, + ), + Case("pre_request_hook_is_not_run_on_a_sync_call", upstream=(), sync=True, callbacks=(_RaisingHook,)), + Case( + "interceptor_raises_runs_the_failure_deployment_hook_once", + _intercepted(), + upstream=(), + interceptor=_Refusing, + loggers=_loggers(_DeploymentHooks), + ), + Case("interceptor_answer_is_cached", _intercepted(), upstream=(), cache=Cache, interceptor=_Answering), + Case( + "interceptor_stream_answer_with_cache", + _intercepted(stream=True), + upstream=(), + cache=Cache, + interceptor=_Answering, + ), + Case("interceptor_ignores_other_calls", _kw(tools=[dict(WEATHER_TOOL)]), interceptor=_Answering), +) + +SHORT_CIRCUITED: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "id": "msg_01ParityShortCircuit", + "content": [{"type": "text", "text": "Paris is the capital."}], + } +) + + +class _ShortCircuit(WebSearchInterceptionLogger): + """A web search interception whose replaced short-circuit records its inputs and answers every call.""" + + answer: Mapping[str, JsonValue] | None = SHORT_CIRCUITED + + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # the logger annotates bare containers + self.seen: list[Mapping[str, object]] = [] + + async def try_short_circuit_search( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + custom_llm_provider: str | None, + kwargs: Mapping[str, object] | None = None, + ) -> dict[str, object] | None: + passed: Final = kwargs or {} + self.seen.append( + { + "model": model, + "messages": _OBJECTS.validate_python(messages), + "tools": tools, + "custom_llm_provider": custom_llm_provider, + "kwargs": sorted(name for name in passed if name not in _VOLATILE_HOOK_KWARGS), + "metadata": passed.get("metadata"), + "safeguards": passed.get("safeguards"), + } + ) + return None if self.answer is None else _OBJECT.validate_python(copy.deepcopy(dict(self.answer))) + + +class _ShortCircuitPassing(_ShortCircuit): + answer = None + + +class _ShortCircuitToolUse(_ShortCircuit): + answer = TOOL_USE_RESPONSE + + +class _ShortCircuitRefusing(_ShortCircuit): + async def try_short_circuit_search( # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + self, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + custom_llm_provider: str | None, + kwargs: Mapping[str, object] | None = None, + ) -> dict[str, object] | None: + await super().try_short_circuit_search(model, messages, tools, custom_llm_provider, kwargs) + raise ValueError("the short-circuit search failed") + + +def _searched(**kwargs: object) -> Mapping[str, object]: + return _kw(tools=[dict(WEB_SEARCH_TOOL)], metadata={"user_id": "parity-user"}, **kwargs) + + +SHORT_CIRCUIT_CASES: Final = ( + Case("short_circuit_answers", _searched(), upstream=(), loggers=_loggers(_ShortCircuit)), + Case("short_circuit_answers_stream", _searched(stream=True), upstream=(), loggers=_loggers(_ShortCircuit)), + Case("short_circuit_raises", _searched(), upstream=(), loggers=_loggers(_ShortCircuitRefusing)), + Case("short_circuit_raises_stream", _searched(stream=True), upstream=(), loggers=_loggers(_ShortCircuitRefusing)), + Case("short_circuit_passes_the_call_on", _searched(), loggers=_loggers(_ShortCircuitPassing)), + Case( + "short_circuit_passes_a_stream_on", + _searched(stream=True), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_ShortCircuitPassing), + ), + Case( + "short_circuit_first_answer_wins", + _searched(), + upstream=(), + loggers=_loggers(_ShortCircuitPassing, _ShortCircuit, _ShortCircuitRefusing), + ), + Case( + "short_circuit_answers_before_the_interceptor", + _intercepted(), + upstream=(), + loggers=_loggers(_ShortCircuit), + interceptor=_Answering, + ), + Case( + "short_circuit_passes_on_to_the_interceptor", + _intercepted(), + upstream=(), + loggers=_loggers(_ShortCircuitPassing), + interceptor=_Answering, + ), + Case( + "short_circuit_answer_with_cache", + _searched(), + upstream=(), + cache=Cache, + loggers=_loggers(_ShortCircuit), + ), + Case( + "short_circuit_answer_is_classified", + _safeguarded(), + upstream=(), + model=NON_CLAUDE, + classifier=FLAGGED, + loggers=_loggers(_ShortCircuitToolUse), + ), + Case( + "short_circuit_stream_answer_is_classified", + _safeguarded(stream=True), + upstream=(), + model=NON_CLAUDE, + classifier=FLAGGED, + loggers=_loggers(_ShortCircuitToolUse), + ), + Case("short_circuit_is_not_consulted_on_a_sync_call", _searched(), sync=True, loggers=_loggers(_ShortCircuit)), +) + +MCP_REFERENCE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy/mcp/deepwiki", + "require_approval": "never", + } +) +MCP_NEEDS_APPROVAL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {key: value for key, value in MCP_REFERENCE.items() if key != "require_approval"} +) +WIKI_TOOL_CALL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "tool_use", "id": "toolu_01Wiki", "name": "read_wiki", "input": {"repo": "a/b"}} +) +WIKI_TOOL_USE_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + **TEXT_RESPONSE, + "id": "msg_01ParityWiki", + "content": [{"type": "text", "text": "Reading the wiki."}, dict(WIKI_TOOL_CALL)], + "stop_reason": "tool_use", + } +) +WIKI_TOOL_STREAM: Final = _events( + _message_start("msg_01ParityWikiStream"), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_01Wiki", "name": "read_wiki", "input": {}}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"repo": "a/b"}'}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + *_message_end("tool_use", 9), +) + + +def _wiki(**kwargs: object) -> Mapping[str, object]: + return _kw(**{"tools": [dict(MCP_REFERENCE)], **kwargs}) + + +class _McpToolHook(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + return {**kwargs, "tools": [dict(MCP_REFERENCE), *cast(list[object], kwargs.get("tools") or [])]} + + +_SHARED_LOGGER: Final = ( + "the gateway's follow-up call rewrites the shared Logging object while the first call's async success callback " + "is still queued, on Python and Rust alike, so the callback logs whichever call it saw" +) +SHARED_LOGGER_RACE: Final = MappingProxyType( + { + "logged_request": _SHARED_LOGGER, + "logged_metadata": _SHARED_LOGGER, + "logged_response_headers": _SHARED_LOGGER, + "logged_sent_request": _SHARED_LOGGER, + } +) + + +MCP_CASES: Final = ( + Case( + "mcp_tool_runs_then_answers", + _wiki(), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "mcp_tool_runs_then_answers_stream", + _wiki(stream=True), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case("mcp_answers_without_a_tool", _wiki(), gateway=True), + Case( + "mcp_tool_beside_a_native_tool", + _wiki(tools=[dict(MCP_REFERENCE), dict(WEATHER_TOOL)], system="Be terse.", temperature=0.2), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "mcp_whole_gateway", + _wiki(tools=[{**MCP_REFERENCE, "server_url": "litellm_proxy"}]), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "mcp_tool_needs_approval", + _wiki(tools=[dict(MCP_NEEDS_APPROVAL)]), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE),), + gateway=True, + ), + Case( + "mcp_tool_needs_approval_stream", + _wiki(tools=[dict(MCP_NEEDS_APPROVAL)], stream=True), + upstream=(_stream(WIKI_TOOL_STREAM),), + gateway=True, + ), + Case( + "mcp_follow_up_fails", + _wiki(), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _failure(429, "rate_limit_error", "slow down")), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "mcp_first_call_fails", + _wiki(), + upstream=(_failure(400, "invalid_request_error", "max_tokens: too large"),), + gateway=True, + ), + Case( + "mcp_metadata_and_litellm_metadata", + _wiki(metadata={"user_id": "user-123"}, litellm_metadata={"user_api_key": "hashed", "tags": ["t1"]}), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "mcp_unserved_proxy_path_passes_through", + _wiki(tools=[{**MCP_REFERENCE, "server_url": "https://gateway.example/mcp/unknown"}]), + gateway=True, + ), + Case( + "mcp_skip_flag_passes_the_reference_through", + _wiki(_skip_mcp_handler=True), + gateway=True, + ), + Case( + "mcp_tool_execution_raises", + _wiki(), + upstream=(_ok({**WIKI_TOOL_USE_RESPONSE, "content": [{**WIKI_TOOL_CALL, "input": {"repo": "boom"}}]}),), + gateway=True, + ), + Case( + "mcp_cache_hit", + _wiki(), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok(), _ok(WIKI_TOOL_USE_RESPONSE), _ok()), + cache=Cache, + gateway=True, + ), + Case("mcp_sync", _wiki(), upstream=(), sync=True, gateway=True), + Case("mcp_mock_answers_first", _wiki(mock_response="Hi from the mock"), upstream=(), gateway=True), + Case("mcp_sync_mock_answers_first", _wiki(mock_response="Hi from the mock"), upstream=(), sync=True, gateway=True), + Case( + "pre_request_hook_adds_an_mcp_tool", + _kw(tools=[dict(WEATHER_TOOL)]), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok()), + callbacks=(_McpToolHook,), + gateway=True, + racy=SHARED_LOGGER_RACE, + ), + Case( + "pre_request_hook_removes_the_mcp_tool", + _wiki(), + callbacks=(_UninterceptedToolHook,), + gateway=True, + ), + Case( + "interceptor_answers_before_the_mcp_gateway", + _kw(tools=[dict(INTERCEPTED_TOOL), dict(MCP_REFERENCE)]), + upstream=(), + interceptor=_Answering, + gateway=True, + ), +) + + +class _OrderedMemory(InMemoryCache): + def __init__(self) -> None: + super().__init__() + self.ops: list[str] = [] + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: records the backend writes in order + self.ops.append(f"set {key}") + await super().async_set_cache(key, value, **kwargs) # pyright: ignore[reportUnknownMemberType] # InMemoryCache leaves its signature untyped + + async def async_get_cache(self, key: str, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: records the backend reads in order + self.ops.append(f"get {key}") + return await super().async_get_cache(key, **kwargs) # pyright: ignore[reportUnknownMemberType] # InMemoryCache leaves its signature untyped + + +def _ordered_cache() -> Cache: + cache: Final = Cache() + cache.cache = _OrderedMemory() + return cache + + +EMPTY_CONTENT_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType({**TEXT_RESPONSE, "content": []}) +EMPTY_TEXT_RESPONSE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {**TEXT_RESPONSE, "content": [{"type": "text", "text": ""}]} +) + +HOOKED_STREAM_AND_CACHE_CASES: Final = ( + Case( + "pre_request_hook_streams_runs_the_deployment_success_hook", + callbacks=(_UnstreamingHook,), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentHooks), + ), + Case( + "pre_request_hook_streams_behind_a_cache", + callbacks=(_UnstreamingHook,), + upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), + cache=Cache, + ), + Case( + "stream_one_runs_the_deployment_success_hook", + _kw(stream=1), + upstream=(_stream(TEXT_STREAM),), + loggers=_loggers(_DeploymentHooks), + ), + Case( + "stream_one_behind_a_cache", _kw(stream=1), upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), cache=Cache + ), + Case( + "stream_one_sync_behind_a_cache", + _kw(stream=1), + upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), + cache=Cache, + sync=True, + ), + Case( + "stream_string_true_behind_a_cache", + _kw(stream="true"), + upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), + cache=Cache, + ), + Case( + "pre_request_hook_streams_on_an_integer_behind_a_cache", + callbacks=(_IntegerStreamHook,), + upstream=(_stream(TEXT_STREAM), _stream(TEXT_STREAM)), + cache=Cache, + ), + Case( + "cache_empty_content_is_not_stored", + upstream=(_ok(EMPTY_CONTENT_RESPONSE), _ok(EMPTY_CONTENT_RESPONSE)), + cache=Cache, + ), + Case( + "cache_empty_text_block_is_not_stored", + upstream=(_ok(EMPTY_TEXT_RESPONSE), _ok(EMPTY_TEXT_RESPONSE)), + cache=Cache, + ), + Case( + "mcp_cache_hit_stream", + _wiki(stream=True), + upstream=(_ok(WIKI_TOOL_USE_RESPONSE), _ok(), _ok(WIKI_TOOL_USE_RESPONSE), _ok()), + cache=Cache, + gateway=True, + ), + Case( + "safeguards_classifier_stream_cached", + _safeguarded(stream=True), + upstream=(_stream(TOOL_STREAM), _stream(TOOL_STREAM)), + model=NON_CLAUDE, + classifier=FLAGGED, + cache=_ordered_cache, + ), +) + +CASES: Final = ( + *REQUEST_CASES, + *INPUT_VALUE_CASES, + *ARGUMENT_TYPE_CASES, + *AGENTIC_CASES, + *PROMPT_CACHING_CASES, + *SYNC_PROMPT_CACHING_CASES, + *RESPONSE_CASES, + *FAILURE_CASES, + *STREAM_CASES, + *SYNC_CASES, + *SERVED_NATIVELY_CASES, + *THINKING_RETRY_CASES, + *STRICT_TOOLS_RETRY_CASES, + *FALLBACK_CASES, + *CACHE_CASES, + *THREAD_CASES, + *TRANSPORT_CASES, + *PRE_REQUEST_HOOK_CASES, + *SAFEGUARDS_CASES, + *INTERCEPTOR_CASES, + *SHORT_CIRCUIT_CASES, + *MCP_CASES, + *DEPLOYMENT_HOOK_CASES, + *HOOKED_STREAM_AND_CACHE_CASES, +) @dataclass(frozen=True, slots=True) @@ -780,6 +3354,11 @@ class Leg: hidden_params: Mapping[str, object] logged: Mapping[str, object] callback_kwargs: Mapping[str, object] + cached: tuple[str, ...] = () + warnings: tuple[str, ...] = () + hook_inputs: tuple[object, ...] = () + classified: tuple[Mapping[str, object], ...] = () + tool_runs: tuple[Mapping[str, object], ...] = () @dataclass(slots=True) @@ -806,7 +3385,6 @@ def native_messages() -> Iterator[FallbackProbe]: return _PYTHON_MESSAGES(*args, **kwargs) with ( - rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT), *catalog.RULES)), rebound(dispatch, "_PYTHON_AMESSAGES", python_afallback), rebound(dispatch, "_PYTHON_MESSAGES", python_fallback), ): @@ -820,6 +3398,12 @@ def _environment(env: Mapping[str, str]) -> Generator[None]: *env, "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + "ANTHROPIC_IDENTITY_TOKEN", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_SOURCE", + "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "LITELLM_REASONING_AUTO_SUMMARY", "LITELLM_RUST", ) @@ -843,6 +3427,48 @@ def _settings(settings: Mapping[str, object]) -> Generator[None]: setattr(litellm, name, value) +@contextmanager +def _classifier(reply: str | None, calls: list[Mapping[str, object]]) -> Generator[None]: + if reply is None: + yield + return + acompletion: Final = litellm.acompletion + + async def classify(**kwargs: object) -> object: # kwargs-ok: stands in for litellm.acompletion + calls.append( + { + "names": sorted(kwargs), + **{name: kwargs.get(name) for name in ("model", "messages", "max_tokens", "num_retries", "user")}, + "litellm_metadata": _mapping(kwargs.get("litellm_metadata")), + } + ) + return await acompletion(**kwargs, mock_response=reply) + + with ( + rebound(proxy_server, "general_settings", {"safeguards_classifier_model": CLASSIFIER_MODEL}), + rebound(litellm, "acompletion", classify), + ): + yield + + +@contextmanager +def _interceptor(interceptor: Callable[[], MessagesInterceptor] | None) -> Generator[None]: + if interceptor is None: + yield + return + with rebound(messages_interceptors, "_interceptors", [interceptor(), *messages_interceptors._interceptors]): # pyright: ignore[reportPrivateUsage] # registers a test interceptor ahead of the shipped ones + yield + + +@contextmanager +def _cache(cache: Callable[[], Cache] | None) -> Generator[None]: + if cache is None: + yield + return + with rebound(litellm, "cache", cache()): + yield + + def _kwargs(case: Case, server: RecordingServer) -> dict[str, object]: # Deep: Python writes into the caller's own `litellm_metadata`, which must not leak into the # other leg. @@ -856,7 +3482,7 @@ def _kwargs(case: Case, server: RecordingServer) -> dict[str, object]: "api_base": f"{server.base_url}{suffix}", **supplied, } - return {name: value for name, value in kwargs.items() if value is not None or name in ("safeguards",)} + return {name: value for name, value in kwargs.items() if value is not None or name in ("safeguards", *case.nones)} def _mapping(value: object) -> Mapping[str, object]: @@ -881,15 +3507,60 @@ def _error(error: BaseException, phase: str, chunks: Iterable[object] = ()) -> M "llm_provider": getattr(error, "llm_provider", None), "model": getattr(error, "model", None), "message": _normalized(str(error)), + "details": AnthropicExceptionMapping.embedded_error_details(str(error)), + "keeps_the_deployment_warm": is_message_thread_not_found(error), "headers": {name: header_items.get(name) for name in ("retry-after", "request-id") if name in header_items}, "partial": _stream_text(chunks), } +def _sent_request(response: object) -> Mapping[str, object] | None: + if not isinstance(response, httpx.Response): + return None + try: + sent: Final = response.request + except RuntimeError: + return {"status": response.status_code, "request": None} + return { + "status": response.status_code, + "request": { + "method": sent.method, + "path": sent.url.path, + "anthropic-beta": sent.headers.get("anthropic-beta"), + "body": _sent_body(sent), + }, + } + + +def _sent_body(sent: httpx.Request) -> object: + try: + sent_body: Final = sent.read() + except (httpx.RequestNotRead, httpx.StreamClosed): + return None + return json.loads(sent_body) if sent_body else None + + +def _partition(chunks: Iterable[object]) -> tuple[int, ...]: + return tuple(len(chunk if isinstance(chunk, bytes) else str(chunk).encode()) for chunk in chunks) + + def _stream_text(chunks: Iterable[object]) -> str: return b"".join(chunk if isinstance(chunk, bytes) else str(chunk).encode() for chunk in chunks).decode() +def _scrubbed(value: object) -> object: + """Masks the random server tool ids inside a response, leaving every other value as it was.""" + match value: + case str(): + return _SERVER_TOOL_ID.sub("srvtoolu_", value) + case dict(): + return {name: _scrubbed(item) for name, item in cast(dict[object, object], value).items()} + case list(): + return [_scrubbed(item) for item in cast(list[object], value)] + case _: + return value + + def _response(result: object, private_keys: bool) -> Mapping[str, object]: assert isinstance(result, dict), f"expected a dict response, got {type(result).__name__}" body: Final = { @@ -897,13 +3568,13 @@ def _response(result: object, private_keys: bool) -> Mapping[str, object]: for name, value in cast(dict[str, object], result).items() if private_keys or not name.startswith("_") } - return {"kind": "response", "body": _OBJECT.validate_python(body)} + return {"kind": "response", "body": _LOGGED.validate_python(_scrubbed(body))} def _hidden_params(result: object) -> Mapping[str, object]: hidden: Final = dict(_mapping(get_hidden_params_dict(result))) headers: Final = _mapping(hidden.pop("additional_headers", None)) - header_names: Final = sorted(name for name in headers if name not in ("x-litellm-rust", "x-litellm-cache-key")) + header_names: Final = sorted(name for name in headers if name != "x-litellm-rust") return {"keys": sorted(hidden), "additional_headers": header_names} @@ -949,32 +3620,42 @@ async def _invoke( if isinstance(result, dict): body = cast(Mapping[str, object], result) return _response(body, case.private_keys), _hidden_params(body) + time.sleep(case.busy_before_reading) # noqa: ASYNC251 # holds the event loop, as a caller's own work between the call and its first read does chunks, failure = await _consume_async(cast(AsyncIterator[object], result)) hidden: Final = _hidden_params(result) + partition: Final = {"partition": _partition(chunks)} if case.partitioned else {} if failure is not None: - return _error(failure, "stream", chunks), hidden + return {**_error(failure, "stream", chunks), **partition}, hidden return { "kind": "stream", "chunk_types": sorted({type(chunk).__name__ for chunk in chunks}), - "text": _stream_text(chunks), + "text": _SERVER_TOOL_ID.sub("srvtoolu_", _stream_text(chunks)), + **partition, }, hidden -async def _terminal_events(recorder: RecordingLogger) -> tuple[str, ...]: +async def _terminal_events(recorder: RecordingLogger, since: int = 0) -> tuple[str, ...]: for _ in range(100): await drain_logging() - if any(name in _TERMINAL_EVENTS for name in recorder.names): + if any(name in _TERMINAL_EVENTS for name in recorder.names[since:]): break await asyncio.sleep(0.05) await asyncio.sleep(0.2) await drain_logging() - return tuple(name for name in recorder.names if name in _TERMINAL_EVENTS) + return tuple(name for name in recorder.names[since:] if name in _TERMINAL_EVENTS) -def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping[str, object], Mapping[str, object]]: - terminal: Final = [event for event in recorder.events if event.name in _TERMINAL_EVENTS] +def _logged( + recorder: RecordingLogger, events: tuple[str, ...], since: int = 0 +) -> tuple[Mapping[str, object], Mapping[str, object]]: + terminal: Final = [ + event + for event in recorder.events[since:] + if event.name in _TERMINAL_EVENTS and event.call_type != CLASSIFIER_CALL_TYPE + ] + hooks: Final = tuple(sorted(set(recorder.names[since:]))) if not terminal: - return {"events": events}, {} + return {"events": events, "hooks": hooks}, {} details: Final = _mapping(terminal[-1].kwargs) standard: Final = _mapping(details.get("standard_logging_object")) error_class: Final = _mapping(standard.get("error_information")).get("error_class") @@ -983,35 +3664,139 @@ def _logged(recorder: RecordingLogger, events: tuple[str, ...]) -> tuple[Mapping metadata_keys: Final = _mapping(metadata).keys() logged: Final[Mapping[str, object]] = { "events": events, + "hooks": hooks, **{name: standard.get(name) for name in _COMPARED_LOG_FIELDS}, "error_class": error_class, } callback_kwargs: Final[Mapping[str, object]] = { "logged_metadata": sorted(metadata_keys) if isinstance(metadata, dict) else metadata, "logged_model_info": litellm_params.get("model_info"), - "logged_response_headers": bool(details.get("response_headers")), + "logged_response_headers": { + name: value for name, value in _mapping(details.get("response_headers")).items() if name.lower() != "date" + }, "logged_api_key": details.get("api_key"), "logged_call": { "custom_llm_provider": details.get("custom_llm_provider"), "model": details.get("model"), "stream": details.get("stream"), + "aanthropic_messages": litellm_params.get("aanthropic_messages"), }, + "logged_request": { + "messages": details.get("messages"), + "thinking": details.get("thinking"), + "tools": details.get("tools"), + "complete_input_dict": _mapping(details.get("additional_args")).get("complete_input_dict"), + }, + "logged_sent_request": _sent_request(details.get("httpx_response")), + "logged_model_parameters": standard.get("model_parameters"), + "logged_error": _logged_error(details, standard), } return logged, callback_kwargs +def _logged_error(details: Mapping[str, object], standard: Mapping[str, object]) -> Mapping[str, object]: + exception: Final = details.get("exception") + information: Final = _mapping(standard.get("error_information")) + return { + "error_str": _normalized(str(standard.get("error_str"))), + "error_message": _normalized(str(information.get("error_message"))), + "error_code": information.get("error_code"), + "llm_provider": information.get("llm_provider"), + "exception": None if exception is None else _normalized(str(exception)), + "exception_model": getattr(exception, "model", None), + "exception_provider": getattr(exception, "llm_provider", None), + } + + +async def _prime(case: Case, server: RecordingServer, recorder: RecordingLogger, native: bool) -> int: + if case.cache is None: + return 0 + outcome: Final = await _invoke(case, _kwargs(case, server), native=native) + assert case.sync or outcome[0]["kind"] != "error", outcome + await _terminal_events(recorder) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + return len(recorder.events) + + +WIKI_MCP_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "name": "read_wiki", + "description": "Read a repository wiki", + "inputSchema": {"type": "object", "properties": {"repo": {"type": "string"}}, "required": ["repo"]}, + } +) + + +@contextmanager +def _gateway(case: Case, runs: list[Mapping[str, object]]) -> Generator[None]: # mutable-ok: records the tool runs + if not case.gateway: + yield + return + from mcp.types import Tool as MCPTool + + from litellm.responses.mcp import tool_precheck + from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + split_mcp_tools: Final = LiteLLM_Proxy_MCP_Handler.split_mcp_tools + + async def served(names: Iterable[str]) -> frozenset[str]: + return frozenset({"deepwiki"}) & frozenset(names) + + async def split(tools: Iterable[Mapping[str, object]] | None) -> object: + return await split_mcp_tools(tools, served) + + async def listed(**_: object) -> tuple[list[MCPTool], list[str]]: + return [MCPTool.model_validate(dict(WIKI_MCP_TOOL))], ["deepwiki"] + + async def executed(**kwargs: object) -> list[dict[str, object]]: + calls: Final = cast(list[Mapping[str, object]], kwargs["tool_calls"]) + runs.append( + { + "tool_calls": [dict(call) for call in calls], + "tool_server_map": kwargs["tool_server_map"], + "request_tags": kwargs["request_tags"], + "guardrail_context": json.loads(json.dumps(kwargs["guardrail_context"], default=repr)), + } + ) + if any(call.get("input") == {"repo": "boom"} for call in calls): + raise RuntimeError("the wiki server went away") + return [{"tool_call_id": call.get("id"), "result": "9 sections", "name": call.get("name")} for call in calls] + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(tool_precheck, "_gateway_registry_loads", lambda: True) + patch.setattr(LiteLLM_Proxy_MCP_Handler, "split_mcp_tools", staticmethod(split)) + patch.setattr(LiteLLM_Proxy_MCP_Handler, "_get_mcp_tools_from_manager", staticmethod(listed)) + patch.setattr(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", staticmethod(executed)) + yield + + async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg: - with isolated_callback_registries(), _environment(case.env), _settings(case.settings): - recorder: Final = RecordingLogger() - litellm.callbacks.append(recorder) + classified: Final[list[Mapping[str, object]]] = [] + runs: Final[list[Mapping[str, object]]] = [] + with ( + isolated_callback_registries(), + _environment(case.env), + _settings(case.settings), + _cache(case.cache), + _classifier(case.classifier, classified), + _interceptor(case.interceptor), + _gateway(case, runs), + ): for spec in case.upstream: server.enqueue(replace(spec)) request_count: Final = len(server.requests) + recorder: Final = RecordingLogger() + loggers: Final = () if case.loggers is None else case.loggers() + litellm.callbacks.extend(loggers) + litellm.callbacks.append(recorder) + litellm.callbacks.extend(callback() for callback in case.callbacks) fallbacks_before: Final = 0 if probe is None else probe.calls - outcome, hidden = await _invoke(case, _kwargs(case, server), native=probe is not None) - events: Final = await _terminal_events(recorder) + primed: Final = await _prime(case, server, recorder, native=probe is not None) + with _warnings() as warned: + outcome, hidden = await _invoke(case, _kwargs(case, server), native=probe is not None) + events: Final = await _terminal_events(recorder, primed) server.responses.clear() - logged, callback_kwargs = _logged(recorder, events) + logged, callback_kwargs = _logged(recorder, events, primed) return Leg( fell_back=probe is not None and probe.calls > fallbacks_before, requests=tuple( @@ -1029,9 +3814,40 @@ async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) hidden_params=hidden, logged=logged, callback_kwargs=callback_kwargs, + cached=() if litellm.cache is None else _cache_entries(litellm.cache), + warnings=tuple(warned), + hook_inputs=tuple(getattr(logger, "seen", ()) for logger in loggers), + classified=tuple(classified), + tool_runs=tuple(runs), ) +class _WarningLines(logging.Handler): + def __init__(self, lines: list[str]) -> None: # mutable-ok: the handler appends each warning as it is logged + super().__init__(logging.WARNING) + self.lines: Final = lines + + def emit(self, record: logging.LogRecord) -> None: + self.lines.append(record.getMessage()) + + +@contextmanager +def _warnings() -> Generator[list[str]]: # mutable-ok: filled while the call runs + lines: Final[list[str]] = [] + handler: Final = _WarningLines(lines) + verbose_logger.addHandler(handler) + try: + yield lines + finally: + verbose_logger.removeHandler(handler) + + +def _cache_entries(cache: Cache) -> tuple[str, ...]: + backend: Final = cache.cache + keys: Final = tuple(sorted(backend.cache_dict)) if isinstance(backend, InMemoryCache) else () + return (*keys, *backend.ops) if isinstance(backend, _OrderedMemory) else keys + + @pytest.fixture def parity_server(recording_server: RecordingServer) -> RecordingServer: recording_server.expected_requests = None @@ -1068,12 +3884,521 @@ async def test_native_messages_match_python( case "dispatch": assert not native.fell_back case "request": - assert native.requests == python.requests + assert _comparable(native.requests) == _comparable(python.requests) case "outcome": assert native.outcome == python.outcome case "hidden_params": assert native.hidden_params == python.hidden_params case "logged": assert native.logged == python.logged + case "cached": + assert native.cached == python.cached + case "warnings": + assert native.warnings == python.warnings + case "hook_inputs": + assert native.hook_inputs == python.hook_inputs + case "classified": + assert native.classified == python.classified + case "tool_runs": + assert native.tool_runs == python.tool_runs case _: - assert native.callback_kwargs.get(aspect) == python.callback_kwargs.get(aspect) + assert _comparable(native.callback_kwargs.get(aspect)) == _comparable(python.callback_kwargs.get(aspect)) + + +def _comparable(value: object) -> object: + match value: + case float() if math.isnan(value): + return "" + case Mapping(): + return {name: _comparable(item) for name, item in cast(Mapping[object, object], value).items()} + case list() | tuple(): + return [_comparable(item) for item in cast(Sequence[object], value)] + case _: + return value + + +@pytest.mark.parametrize( + "case", + [ + pytest.param(case, id=case.name) + for case in (*SERVED_NATIVELY_CASES, *AGENTIC_CASES, *SAFEGUARDS_CASES, *SHORT_CIRCUIT_CASES) + ], +) +def test_catalog_serves_the_case_natively(case: Case, parity_server: RecordingServer) -> None: + with ( + isolated_callback_registries(), + _environment(case.env), + _settings(case.settings), + _classifier(case.classifier, []), + ): + litellm.callbacks.extend(() if case.loggers is None else case.loggers()) + binding: Final = dispatch._MESSAGES if case.sync else dispatch._AMESSAGES # pyright: ignore[reportPrivateUsage] # the public binding the entry routes by + features: Final = dispatch._FEATURES # pyright: ignore[reportPrivateUsage] # the entry's own detectors + request: Final = dispatch._public_request(binding, (), _kwargs(case, parity_server)) # pyright: ignore[reportPrivateUsage] # the public binding the entry routes by + assert request is not None + context: Final = dispatch._context(request, features) # pyright: ignore[reportPrivateUsage] # the entry's own route context + assert catalog.rollout(context) is Rollout.RUST_REQUIRED, context.features + + +@pytest.mark.parametrize("case", [pytest.param(case, id=case.name) for case in TRANSPORT_CASES]) +def test_catalog_serves_the_transport_case_natively(case: Case, parity_server: RecordingServer) -> None: + with isolated_callback_registries(), _environment(case.env), _settings(case.settings): + request: Final = dispatch._public_request(dispatch._AMESSAGES, (), _kwargs(case, parity_server)) # pyright: ignore[reportPrivateUsage] # the public binding the async entry routes by + assert request is not None + context: Final = dispatch._context(request, dispatch._FEATURES) # pyright: ignore[reportPrivateUsage] # the async entry's own route context + assert catalog.rollout(context) is Rollout.RUST_REQUIRED, context.features + + +@dataclass(slots=True) +class _Exchanges: + posts: int = 0 + windows: int = 0 + + +async def _exchanges( + case: Case, server: RecordingServer, probe: FallbackProbe | None, monkeypatch: pytest.MonkeyPatch +) -> _Exchanges: + counted: Final = _Exchanges() + post: Final = AsyncHTTPHandler.post # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the handler's untyped dict params + record: Final = logging_utils._set_duration_in_model_call_details # pyright: ignore[reportPrivateUsage] # the one recorder of an exchange window + + async def posted( + handler: AsyncHTTPHandler, *args: object, **kwargs: object + ) -> object: # kwargs-ok: the handler's own post signature + counted.posts += 1 + return await post(handler, *args, **kwargs) # pyright: ignore[reportArgumentType] # forwarded unchanged + + def windowed(**kwargs: object) -> None: # kwargs-ok: the recorder's own keywords + counted.windows += 1 + record(**kwargs) # pyright: ignore[reportArgumentType] # forwarded unchanged + + with monkeypatch.context() as patch: + patch.setattr(AsyncHTTPHandler, "post", posted) + patch.setattr(logging_utils, "_set_duration_in_model_call_details", windowed) + await _leg(case, server, probe) + return counted + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", [pytest.param(case, id=case.name) for case in TRANSPORT_CASES if not case.sync]) +async def test_transport_cases_post_through_the_handler_python_posts_through( + case: Case, parity_server: RecordingServer, native_messages: FallbackProbe, monkeypatch: pytest.MonkeyPatch +) -> None: + python: Final = await _exchanges(case, parity_server, None, monkeypatch) + native: Final = await _exchanges(case, parity_server, native_messages, monkeypatch) + ignored: Final = case.name.endswith("_is_ignored") + assert python.posts == len(case.upstream) + assert native.posts == (0 if ignored else python.posts) + assert native.windows == python.windows + assert native_messages.calls == 0 + + +_COST_CASES: Final = tuple( + pytest.param(case, id=case.name) for case in SERVED_NATIVELY_CASES if case.name.startswith("stream_cost_injection") +) + + +def _proxy_relayed(text: str, request_headers: Mapping[str, object]) -> str: + return "".join( + str( + ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + frame + "\n\n", "proxy-model-alias", None, request_headers=request_headers + ) + ) + for frame in text.split("\n\n") + if frame + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", _COST_CASES) +async def test_both_lanes_carry_one_cost_per_usage_through_the_proxy_relay( + case: Case, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python, native = await _legs(case, parity_server, native_messages) + opted_in: Final = {"x-litellm-include-cost-in-streaming-usage": "true"} + python_text: Final = str(python.outcome["text"]) + native_text: Final = str(native.outcome["text"]) + assert '"cost"' in python_text + assert native_text == python_text + relayed: Final = _proxy_relayed(native_text, opted_in) + assert relayed == _proxy_relayed(python_text, opted_in) + assert all(frame.count('"cost"') <= 1 for frame in relayed.split("\n\n")) + + +def _proxied(kwargs: Mapping[str, object]) -> dict[str, object]: + logging_obj, data = litellm.utils.function_setup( + original_function="anthropic_messages", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime.now(), + is_async_call=True, + litellm_call_id=str(uuid.uuid4()), + **kwargs, + ) + return {**data, "litellm_logging_obj": logging_obj} + + +async def _through_the_ownership_hook( + hook: MessageThreadOwnership, owner: UserAPIKeyAuth, data: dict[str, object], native: bool +) -> None: + acall: Final = _PUBLIC_AMESSAGES if native else _PYTHON_AMESSAGES + result: Final = await acall(**data) + if isinstance(result, dict): + await hook.async_post_call_success_hook(data, owner, cast(AnthropicMessagesResponse, result)) + return + async for _ in hook.async_post_call_streaming_iterator_hook(owner, cast(AsyncIterator[object], result), data): + pass + + +async def _thread_betas(server: RecordingServer, stream: bool, native: bool) -> tuple[str | None, str | None]: + hook: Final = MessageThreadOwnership(InternalUsageCache(dual_cache=DualCache())) + owner: Final = UserAPIKeyAuth(token="hashed-owner") + base: Final[dict[str, object]] = { + "model": SONNET, + "max_tokens": 64, + "messages": [dict(USER)], + "api_key": API_KEY, + "api_base": server.base_url, + } + server.enqueue(_stream(TEXT_STREAM) if stream else _ok()) + server.enqueue(_ok()) + sent: Final = len(server.requests) + with isolated_callback_registries(), _environment({}): + await _through_the_ownership_hook( + hook, + owner, + _proxied({**base, "thread": dict(THREAD_CREATE), "tools": [dict(ADVISOR_TOOL)], "stream": stream}), + native, + ) + replied: Final = "msg_01ParityStream" if stream else str(TEXT_RESPONSE["id"]) + continued: Final = _proxied({**base, "thread": {"type": "continue", "previous_message_id": replied}}) + hooked: Final = await hook.async_pre_call_hook(owner, DualCache(), continued, "anthropic_messages") + await _through_the_ownership_hook(hook, owner, continued if hooked is None else hooked, native) + await drain_logging() + server.responses.clear() + create, resumed = server.requests[sent:] + return create.headers.get("anthropic-beta"), resumed.headers.get("anthropic-beta") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("json", "streamed")) +async def test_both_lanes_resend_a_thread_s_create_betas_through_the_ownership_hook( + stream: bool, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python: Final = await _thread_betas(parity_server, stream, native=False) + native: Final = await _thread_betas(parity_server, stream, native=True) + + assert native_messages.calls == 0 + assert native == python + assert python == ("advisor-tool-2026-03-01,message-threads-2026-08-12",) * 2 + + +_WEB_SEARCH_PRE_HOOK_CASES: Final = tuple( + pytest.param(case, id=case.name) + for case in AGENTIC_CASES + if case.name in {"web_search_interception", "web_search_interception_converted_stream"} +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", _WEB_SEARCH_PRE_HOOK_CASES) +async def test_both_lanes_hand_the_web_search_pre_hook_kwargs_to_the_agentic_loop( + case: Case, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python, native = await _legs(case, parity_server, native_messages) + + assert native_messages.calls == 0 + assert native.hook_inputs == python.hook_inputs + [seen] = python.hook_inputs + first: Final = cast(list[Mapping[str, object]], seen)[0] + converted: Final = case.name.endswith("_converted_stream") + assert first["kwargs"] == { + "max_agentic_loops": 2, + "_websearch_interception_emit_native_blocks": True, + **({"_websearch_interception_converted_stream": True} if converted else {}), + } + assert [tool.get("name") for tool in cast(list[Mapping[str, object]], first["tools"])] == ["litellm_web_search"] + assert first["stream"] is False + + +class _CountingAnswering(_Answering): + def __init__(self) -> None: + self.asked = 0 + + def can_handle(self, tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None) -> bool: + self.asked += 1 + return super().can_handle(tools, custom_llm_provider) + + +async def _interceptor_asks(server: RecordingServer, hooked: bool, native: bool) -> tuple[int, tuple[object, ...]]: + counting: Final = _CountingAnswering() + acall: Final = _PUBLIC_AMESSAGES if native else _PYTHON_AMESSAGES + kwargs: Final[dict[str, object]] = { + "model": SONNET, + "max_tokens": 64, + "messages": [dict(USER)], + "api_key": API_KEY, + "api_base": server.base_url, + "tools": [dict(WEATHER_TOOL)] if hooked else [dict(INTERCEPTED_TOOL), dict(WEATHER_TOOL)], + } + answers: Final[list[object]] = [] + with ( + isolated_callback_registries(), + _environment({}), + _cache(Cache), + _interceptor(lambda: counting), + ): + if hooked: + litellm.callbacks.append(_InterceptedToolHook()) + for _ in range(2): + answers.append(_response(await acall(**copy.deepcopy(kwargs)), private_keys=False)) + await drain_logging() + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + return counting.asked, tuple(answers) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hooked", (False, True), ids=("plain", "pre_request_hooked")) +async def test_both_lanes_ask_the_interceptor_only_on_a_cache_miss( + hooked: bool, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python: Final = await _interceptor_asks(parity_server, hooked, native=False) + native: Final = await _interceptor_asks(parity_server, hooked, native=True) + + assert native_messages.calls == 0 + assert not parity_server.requests + assert native == python + assert python[0] == 1 + + +class _ServiceRecorder: + def __init__(self) -> None: + self.events: list[Mapping[str, object]] = [] # mutable-ok: appended per service hook + + async def async_service_success_hook(self, **kwargs: object) -> None: # kwargs-ok: the service hook's keyword call + self.events.append( + { + "service": kwargs.get("service"), + "call_type": kwargs.get("call_type"), + "timed": isinstance(kwargs.get("duration"), float), + "parent_otel_span": kwargs.get("parent_otel_span"), + } + ) + + +async def _llm_api_timings(server: RecordingServer, stream: bool, native: bool) -> tuple[Mapping[str, object], ...]: + recorder: Final = _ServiceRecorder() + server.enqueue(_stream(TEXT_STREAM) if stream else _ok()) + acall: Final = _PUBLIC_AMESSAGES if native else _PYTHON_AMESSAGES + with isolated_callback_registries(), _environment({}), rebound(logging_utils, "_service_logger", recorder): + result: Final = await acall( + model=SONNET, + max_tokens=64, + messages=[dict(USER)], + api_key=API_KEY, + api_base=server.base_url, + stream=stream, + ) + if not isinstance(result, dict): + await _consume_async(cast(AsyncIterator[object], result)) + await drain_logging() + await asyncio.sleep(0) + server.responses.clear() + return tuple(recorder.events) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("json", "streamed")) +async def test_both_lanes_time_the_provider_exchange_as_one_llm_service_event( + stream: bool, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python: Final = await _llm_api_timings(parity_server, stream, native=False) + native: Final = await _llm_api_timings(parity_server, stream, native=True) + + assert native_messages.calls == 0 + assert native == python + assert [event["call_type"] for event in python] == ["post <- track_llm_api_timing"] + + +class _AttemptRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger takes untyped kwargs + self.seen: list[Mapping[str, object]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + logged: Final = _LOGGED.validate_python(kwargs.get("standard_logging_object") or {}) + self.seen.append( + { + "event": "success", + "model_id": logged.get("model_id"), + "completion_tokens": logged.get("completion_tokens"), + } + ) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the hook's types + logged: Final = _LOGGED.validate_python(kwargs.get("standard_logging_object") or {}) + self.seen.append( + { + "event": "failure", + "model_id": logged.get("model_id"), + "exception": type(kwargs.get("exception")).__name__, + } + ) + + +async def _stream_fallback_attempts( + server: RecordingServer, native: bool +) -> tuple[tuple[Mapping[str, object], ...], str]: + overloaded: Final = _events( + _message_start("msg_01ParityOverloaded"), ("error", dict(_upstream_error("overloaded_error", "Overloaded"))) + ) + server.enqueue(_stream(overloaded)) + server.enqueue(_stream(TEXT_STREAM)) + recorder: Final = _AttemptRecorder() + deployment: Final[Mapping[str, object]] = {"model": SONNET, "api_base": server.base_url, "api_key": API_KEY} + with contextlib.ExitStack() as stack: + stack.enter_context(isolated_callback_registries()) + stack.enter_context(_environment({})) + if not native: + stack.enter_context(rebound(litellm, "anthropic_messages", _PYTHON_AMESSAGES)) + litellm.callbacks.append(recorder) + router: Final = litellm.Router( + model_list=[ + {"model_name": "primary", "litellm_params": dict(deployment), "model_info": {"id": "dep-primary"}}, + {"model_name": "fallback", "litellm_params": dict(deployment), "model_info": {"id": "dep-fallback"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + result: Final = await router.aanthropic_messages( + model="primary", max_tokens=64, messages=[dict(USER)], stream=True + ) + chunks, failure = await _consume_async(cast(AsyncIterator[object], result)) + assert failure is None + for _ in range(100): + await drain_logging() + if len(recorder.seen) >= 2: + break + await asyncio.sleep(0.05) + server.responses.clear() + return tuple(recorder.seen), _stream_text(chunks) + + +@pytest.mark.asyncio +async def test_both_lanes_log_a_stream_superseded_by_a_router_fallback_as_its_frame_s_failure( + parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python: Final = await _stream_fallback_attempts(parity_server, native=False) + native: Final = await _stream_fallback_attempts(parity_server, native=True) + + assert native_messages.calls == 0 + assert native == python + assert sorted(str(event["event"]) for event in python[0]) == ["failure", "success"] + + +async def _router_stream( + server: RecordingServer, native: bool, upstream: tuple[ResponseSpec, ...] +) -> tuple[tuple[str, ...], tuple[Mapping[str, object], ...]]: + for response in upstream: + server.enqueue(response) + recorder: Final = _AttemptRecorder() + deployment: Final[Mapping[str, object]] = {"model": SONNET, "api_base": server.base_url, "api_key": API_KEY} + with contextlib.ExitStack() as stack: + stack.enter_context(isolated_callback_registries()) + stack.enter_context(_environment({})) + if not native: + stack.enter_context(rebound(litellm, "anthropic_messages", _PYTHON_AMESSAGES)) + litellm.callbacks.append(recorder) + router: Final = litellm.Router( + model_list=[ + {"model_name": "primary", "litellm_params": dict(deployment), "model_info": {"id": "dep-primary"}}, + {"model_name": "fallback", "litellm_params": dict(deployment), "model_info": {"id": "dep-fallback"}}, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + result: Final = await router.aanthropic_messages( + model="primary", max_tokens=64, messages=[dict(USER)], stream=True + ) + chunks, failure = await _consume_async(cast(AsyncIterator[object], result)) + assert failure is None + for _ in range(100): + await drain_logging() + if recorder.seen: + break + await asyncio.sleep(0.05) + server.responses.clear() + return tuple(_stream_text((chunk,)) for chunk in chunks), tuple(recorder.seen) + + +OVERLOADED_AFTER_PINGS: Final = _events( + _message_start("msg_01ParityPingedOverloaded"), + *((_PING,) * 50), + ("error", dict(_upstream_error("overloaded_error", "Overloaded"))), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upstream", + [ + pytest.param((_burst(PINGED_STREAM),), id="pings_after_message_start"), + pytest.param((_burst(HEAVILY_PINGED_STREAM),), id="a_thousand_pings_after_message_start"), + pytest.param((_burst(OVERLOADED_AFTER_PINGS), _burst(TEXT_STREAM)), id="pings_then_overloaded_then_fallback"), + ], +) +async def test_both_lanes_stream_a_router_burst_in_the_same_chunks( + parity_server: RecordingServer, native_messages: FallbackProbe, upstream: tuple[ResponseSpec, ...] +) -> None: + python: Final = await _router_stream(parity_server, native=False, upstream=upstream) + native: Final = await _router_stream(parity_server, native=True, upstream=upstream) + + assert native_messages.calls == 0 + assert native == python + + +async def _abandoned_stream( + server: RecordingServer, native: bool, upstream: ResponseSpec +) -> tuple[tuple[str, ...], tuple[Mapping[str, object], ...]]: + server.enqueue(upstream) + recorder: Final = _AttemptRecorder() + acall: Final = _PUBLIC_AMESSAGES if native else _PYTHON_AMESSAGES + with contextlib.ExitStack() as stack: + stack.enter_context(isolated_callback_registries()) + stack.enter_context(_environment({})) + litellm.callbacks.append(recorder) + stream: Final = cast( + AsyncGenerator[object], + await acall( + model=SONNET, + max_tokens=64, + messages=[dict(USER)], + api_key=API_KEY, + api_base=server.base_url, + stream=True, + ), + ) + first: Final = await anext(stream) + await stream.aclose() + for _ in range(100): + await drain_logging() + if recorder.seen: + break + await asyncio.sleep(0.05) + server.responses.clear() + return (_stream_text((first,)),), tuple(recorder.seen) + + +@pytest.mark.asyncio +async def test_both_lanes_bill_a_stream_abandoned_after_its_first_chunk_alike( + parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python: Final = await _abandoned_stream(parity_server, native=False, upstream=_burst(TEXT_STREAM)) + native: Final = await _abandoned_stream(parity_server, native=True, upstream=_burst(TEXT_STREAM)) + + assert native_messages.calls == 0 + assert native == python diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index d88feec7d5e..4bf0b389356 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -6,6 +6,7 @@ from pydantic import TypeAdapter import litellm from litellm.caching.caching import Cache, CacheMode +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) @@ -33,6 +34,14 @@ from tests.test_litellm_rust.support.routing import native_route pytestmark = pytest.mark.requires_rust_extension +_STORED: Final = TypeAdapter(dict[str, object]) + + +def _entries(storage: InMemoryCache) -> dict[str, object]: + return _STORED.validate_python( + storage.cache_dict # pyright: ignore[reportUnknownMemberType] # InMemoryCache leaves its dict untyped + ) + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -110,7 +119,9 @@ async def test_legacy_cache_controls_and_backend_credential_key_semantics(record recording_server.expected_requests = 3 litellm.cache = Cache() await invoke("messages", recording_server, {"cache": {"no-store": True}}) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) await invoke("messages", recording_server, {}) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) await invoke("messages", recording_server, {}) assert len(recording_server.requests) == 2 await invoke("messages", recording_server, {"cache": {"no-cache": True}}) @@ -137,9 +148,11 @@ async def test_legacy_messages_replays_a_completed_stream(recording_server: Reco assert cache_key(first_stream) is None first: Final = await collect(first_stream) await recorder.wait_for_async("async_log_success_event") + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) second_stream: Final = await litellm.anthropic_messages(**parameters) - assert isinstance(cache_key(second_stream), str) - assert cache_key(second_stream) == get_hidden_params_dict(second_stream)["cache_key"] + assert cache_key(second_stream) is None + assert get_hidden_params_dict(second_stream)["cache_hit"] is True + assert isinstance(get_hidden_params_dict(second_stream)["cache_key"], str) second: Final = await collect(second_stream) assert payload(first) == payload(second) assert first == b"".join(recording_server.default_response.payloads()) @@ -174,13 +187,18 @@ async def test_rust_messages_uses_a_legacy_cache_without_python_inference( "api_base": recording_server.base_url, } first: Final = await invoke("messages", recording_server, parameters) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) second: Final = await invoke("messages", recording_server, parameters) - assert cache_key(second) - assert cache_key(first) is None assert payload(first) == payload(second) assert len(recording_server.requests) == 1 + assert cache_key(first) is None + assert cache_key(second) is None - await logger.wait_for_async("async_log_success_event", count=2) + successes: Final = await logger.wait_for_async("async_log_success_event", count=2) + assert [TypeAdapter(dict[str, object]).validate_python(event.kwargs).get("cache_hit") for event in successes] == [ + None, + True, + ] assert logger.names.count("async_log_success_event") == 2 assert "log_failure_event" not in logger.names assert "async_log_failure_event" not in logger.names @@ -210,7 +228,9 @@ async def test_legacy_cache_default_off_requires_opt_in_even_for_existing_entrie @pytest.mark.asyncio -async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recording_server: RecordingServer) -> None: +async def test_a_legacy_cache_hit_skips_the_pre_call_callback_that_would_rewrite_the_request( + recording_server: RecordingServer, +) -> None: class Rewrite(RecordingLogger): temperature = 0.1 @@ -222,15 +242,16 @@ async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recor litellm.cache = Cache() recording_server.expected_requests = 1 await invoke("messages", recording_server, {"callbacks": [logger]}) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) first_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]}) logger.temperature = 0.8 await invoke("messages", recording_server, {"callbacks": [logger]}) second_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]}) - assert logger.names.count("log_pre_api_call") == 4 + assert logger.names.count("log_pre_api_call") == 1 + assert cache_key(first_hit) is None + assert cache_key(second_hit) is None assert len(recording_server.requests) == recording_server.expected_requests assert recording_server.requests[0].body["temperature"] == 0.1 - assert isinstance(cache_key(first_hit), str) - assert isinstance(cache_key(second_hit), str) @pytest.mark.asyncio @@ -258,7 +279,7 @@ async def test_rust_responses_rejects_a_configured_python_cache( @pytest.mark.asyncio @pytest.mark.parametrize("cancel_lookup", (False, True)) @pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) -async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_unary_store( +async def test_python_cache_lookups_stay_in_the_rust_callers_task_and_stores_run_detached( recording_server: RecordingServer, cancel_lookup: bool, stream: bool, @@ -287,11 +308,8 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_ async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: writer: Final = asyncio.current_task() - if stream: - assert writer is caller - else: - assert writer is not caller - assert writer in _PENDING_CACHE_WRITES + assert writer is not caller + assert writer in _PENDING_CACHE_WRITES await asyncio.sleep(0) storage.set_cache(key, value, **kwargs) @@ -314,7 +332,10 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_ second: Final = await invoke("messages", recording_server, {"stream": stream}) second_payload: Final = await collect(second) if stream else payload(second) assert first_payload == second_payload - assert cache_key(second) + assert cache_key(first) is None + assert cache_key(second) is None + if stream: + assert get_hidden_params_dict(second)["cache_key"] in _entries(storage) assert len(recording_server.requests) == 1 @@ -343,12 +364,12 @@ async def test_rust_unary_messages_store_keeps_the_hidden_params_stamped_after_t "model_id": "deployment", } assert get_hidden_params_dict(first)["additional_headers"] == { - "x-litellm-rust": "true", "x-litellm-model-group": "group", "x-litellm-attempted-retries": 0, } hit: Final = await invoke("messages", recording_server, {}) - assert get_hidden_params_dict(hit)["cache_key"] == cache_key(hit) + assert isinstance(get_hidden_params_dict(hit)["cache_key"], str) + assert cache_key(hit) is None assert len(recording_server.requests) == 1 @@ -389,6 +410,8 @@ async def test_messages_cache_stores_the_response_a_yielding_deployment_hook_ret async def test_rust_messages_legacy_cache_honors_request_namespaces( recording_server: RecordingServer, namespace_source: str ) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + litellm.cache = Cache() recording_server.expected_requests = 2 first_options: Final = ( @@ -399,13 +422,20 @@ async def test_rust_messages_legacy_cache_honors_request_namespaces( if namespace_source == "cache" else {"metadata": {"redis_namespace": "second"}} ) + storage: Final = litellm.cache.cache + assert isinstance(storage, InMemoryCache) await invoke("messages", recording_server, first_options) - second: Final = await invoke("messages", recording_server, second_options) - first_hit: Final = await invoke("messages", recording_server, first_options) - second_hit: Final = await invoke("messages", recording_server, second_options) - assert cache_key(second) is None - assert cache_key(first_hit) is not None - assert cache_key(second_hit) is not None + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + first_keys: Final = set(_entries(storage)) + await invoke("messages", recording_server, second_options) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + second_keys: Final = set(_entries(storage)) - first_keys + assert len(recording_server.requests) == 2 + assert [key.split(":", 1)[0] for key in first_keys] == ["first"] + assert [key.split(":", 1)[0] for key in second_keys] == ["second"] + assert {key.split(":", 1)[1] for key in first_keys} == {key.split(":", 1)[1] for key in second_keys} + await invoke("messages", recording_server, first_options) + await invoke("messages", recording_server, second_options) assert len(recording_server.requests) == 2 @@ -420,9 +450,8 @@ async def test_rust_messages_legacy_semantic_cache_preserves_python_scope( first_options: Final = {"messages": [{"role": "user", "content": "hello"}]} second_options: Final = {"messages": [{"role": "user", "content": "hi"}]} first: Final = await invoke("messages", recording_server, first_options) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) second: Final = await invoke("messages", recording_server, second_options) - assert cache_key(first) is None - assert cache_key(second) is not None assert payload(first) == payload(second) assert len(recording_server.requests) == 1 @@ -443,3 +472,62 @@ async def test_legacy_cache_keeps_public_messages_responses_compatible( second_payload: Final = await collect(second) if stream else payload(second) assert second_payload == first_payload assert len(recording_server.requests) == 1 + + +_PER_CALL_HIDDEN_PARAMS: Final = frozenset( + {"_response_ms", "litellm_overhead_time_ms", "callback_duration_ms", "litellm_call_id"} +) + + +async def _write_then_read( + recording_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + writer: bool, + reader: bool, + stream: bool, +) -> tuple[object, object, object, int]: + from litellm.caching.in_memory_cache import InMemoryCache + + storage: Final = InMemoryCache() + monkeypatch.setattr(litellm, "cache", Cache(_backend=storage)) + requests_before: Final = len(recording_server.requests) + monkeypatch.setenv("LITELLM_RUST", "1" if writer else "0") + first: Final = await invoke("messages", recording_server, {"stream": stream}, native=writer) + if stream: + await collect(first) + else: + add_retry_headers_to_response( + _add_headers_to_response(first, {"x-litellm-model-group": "group"}), attempted_retries=0 + ) + get_hidden_params_dict(first)["model_id"] = "deployment" + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + (entry,) = _entries(storage).values() + stored: Final = TypeAdapter(dict[str, object]).validate_python(entry)["response"] + monkeypatch.setenv("LITELLM_RUST", "1" if reader else "0") + hit: Final = await invoke("messages", recording_server, {"stream": stream}, native=reader) + hidden: Final = { + name: value for name, value in get_hidden_params_dict(hit).items() if name not in _PER_CALL_HIDDEN_PARAMS + } + headers: Final = TypeAdapter(dict[str, object]).validate_python(hidden.pop("additional_headers", {})) + kept: Final = { + name: value for name, value in headers.items() if name not in {"x-litellm-rust", "x-litellm-cache-key"} + } + served: Final = await collect(hit) if stream else payload(hit) + return stored, served, {**hidden, "additional_headers": kept}, len(recording_server.requests) - requests_before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +@pytest.mark.parametrize(("writer", "reader"), ((True, False), (False, True)), ids=("rust_to_python", "python_to_rust")) +async def test_a_cache_entry_written_by_one_leg_is_stored_and_served_like_the_readers_own( + recording_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + writer: bool, + reader: bool, + stream: bool, +) -> None: + recording_server.expected_requests = 2 + crossed: Final = await _write_then_read(recording_server, monkeypatch, writer, reader, stream) + own: Final = await _write_then_read(recording_server, monkeypatch, reader, reader, stream) + assert crossed == own + assert crossed[3] == 1 diff --git a/tests/test_litellm_rust/messages/data.py b/tests/test_litellm_rust/messages/data.py index 46127c9237a..d11c901d7d9 100644 --- a/tests/test_litellm_rust/messages/data.py +++ b/tests/test_litellm_rust/messages/data.py @@ -1,7 +1,6 @@ from typing import Final import litellm -from tests.test_litellm_rust.support.gaps import gap from tests.test_litellm_rust.support.recording_server import RecordingServer from tests.test_litellm_rust.support.routing import call_public @@ -57,6 +56,3 @@ async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs: asynchronous=asynchronous, **kwargs, ) - - -BODY_HOOKS: Final = gap("native Messages does not run the handler body hooks") diff --git a/tests/test_litellm_rust/messages/test_body_hooks.py b/tests/test_litellm_rust/messages/test_body_hooks.py index 59f27d5dc04..bf333a17725 100644 --- a/tests/test_litellm_rust/messages/test_body_hooks.py +++ b/tests/test_litellm_rust/messages/test_body_hooks.py @@ -1,16 +1,22 @@ +import asyncio from collections.abc import Iterator from datetime import datetime from typing import Final import pytest +from pydantic import TypeAdapter import litellm +from litellm.caching.caching import Cache +from litellm.caching.caching_handler import ( + _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # the first call's entry must land before the second reads it +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.messages.native import interceptors from litellm.llms.anthropic.messages.native.interceptors.base import MessagesInterceptor -from tests.test_litellm_rust.messages.data import BODY_HOOKS, MESSAGES_RESPONSE, REQUEST, call +from tests.test_litellm_rust.messages.data import MESSAGES_RESPONSE, REQUEST, call from tests.test_litellm_rust.support.callback_recorder import drain_logging from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks from tests.test_litellm_rust.support.isolation import rebound @@ -20,6 +26,8 @@ from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_nativ pytestmark = pytest.mark.requires_rust_extension +_SENT: Final = TypeAdapter(dict[str, object]) + @pytest.fixture def server(server: RecordingServer) -> Iterator[RecordingServer]: @@ -27,7 +35,6 @@ def server(server: RecordingServer) -> Iterator[RecordingServer]: @pytest.mark.asyncio -@BODY_HOOKS async def test_body_pre_request_hook_rewrites_the_provider_request(server: RecordingServer) -> None: class Rewrite(CustomLogger): async def async_pre_request_hook( @@ -44,7 +51,27 @@ async def test_body_pre_request_hook_rewrites_the_provider_request(server: Recor @pytest.mark.asyncio -@BODY_HOOKS +async def test_body_pre_request_hook_runs_only_after_a_response_cache_miss(server: RecordingServer) -> None: + seen: Final[list[object]] = [] + + class Rewrite(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object]: + seen.append(kwargs.get("stop_sequences")) + return {**kwargs, "stop_sequences": ["HOOKED"]} + + with rebound(litellm, "callbacks", [Rewrite()]), rebound(litellm, "cache", Cache()): + await call(server, stop_sequences=["CALLER"]) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + await call(server, stop_sequences=["CALLER"]) + + assert seen == [None] + assert len(server.requests) == 1 + assert _SENT.validate_python(server.requests[0].body).get("stop_sequences") == ["HOOKED"] + + +@pytest.mark.asyncio async def test_body_websearch_short_circuit_answers_without_the_provider(server: RecordingServer) -> None: class ShortCircuit(WebSearchInterceptionLogger): async def try_short_circuit_search( @@ -65,7 +92,6 @@ async def test_body_websearch_short_circuit_answers_without_the_provider(server: @pytest.mark.asyncio -@BODY_HOOKS async def test_body_interceptor_replaces_the_provider_call(server: RecordingServer) -> None: class Replace(MessagesInterceptor): def can_handle(self, tools: list[dict[str, object]] | None, custom_llm_provider: str | None) -> bool: @@ -82,7 +108,6 @@ async def test_body_interceptor_replaces_the_provider_call(server: RecordingServ @pytest.mark.asyncio -@BODY_HOOKS async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_response(server: RecordingServer) -> None: loop: Final = PlannedAgenticLoop() server.enqueue(ResponseSpec(body=MESSAGES_RESPONSE)) @@ -105,7 +130,6 @@ async def test_body_planned_agentic_loop_sends_the_followup_and_replaces_the_res @pytest.mark.asyncio -@BODY_HOOKS async def test_body_legacy_agentic_loop_replaces_the_response(server: RecordingServer) -> None: loop: Final = LegacyAgenticLoop() diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index b580eca2e42..e26aafa5e55 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -1,5 +1,5 @@ import inspect -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect import httpx @@ -8,22 +8,29 @@ import pytest import litellm from litellm.caching.caching import Cache from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS +from litellm.exceptions import UnsupportedParamsError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.anthropic.messages.native import handler as python_messages +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.anthropic.messages.native import interceptors as messages_interceptors +from litellm.llms.anthropic.messages.native.interceptors import MessagesInterceptor from litellm.messages import dispatch from litellm.messages.dispatch import ( _ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch _DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch ) from litellm.proxy import proxy_server +from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target +from litellm.rust_bridge.catalog import MESSAGES_PYTHON_FEATURES, RULES, Feature, Route, RouteRule, Rules, Target from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages from litellm.rust_bridge.public_call import NativeCall from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.routing import routed MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () @@ -251,50 +258,6 @@ DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context" CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"} -@pytest.mark.parametrize( - ("model", "kwargs", "general_settings", "serves_python"), - ( - ("deepseek/deepseek-v4-flash", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, True), - ( - "deepseek-v4-flash", - {"safeguards": DANGEROUS_TOOL_USE, "custom_llm_provider": "deepseek"}, - CLASSIFIER_SETTINGS, - True, - ), - ("deepseek/deepseek-v4-flash", {"safeguards": DANGEROUS_TOOL_USE}, {}, False), - ("deepseek/deepseek-v4-flash", {}, CLASSIFIER_SETTINGS, False), - ("deepseek/deepseek-v4-flash", {"safeguards": [{"type": "other"}]}, CLASSIFIER_SETTINGS, False), - ("anthropic/claude-sonnet-4-5", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, False), - ), -) -def test_proxy_classified_safeguards_reach_the_feature_rule( - monkeypatch: pytest.MonkeyPatch, - model: str, - kwargs: Mapping[str, object], - general_settings: Mapping[str, object], - serves_python: bool, -) -> None: - monkeypatch.setattr(proxy_server, "general_settings", general_settings) - python_result: Final = response(model) - native_result: Final = response(model) - - def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback - return python_result - - result: Final = _DISPATCH.run( - (16, MESSAGES, model), - kwargs, - python=python, - binding=messages_binding(lambda request: native_result), - native=lambda hook, request, call_args, call_kwargs: hook(request), - rules=( - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), - RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), - ), - ) - assert result is (python_result if serves_python else native_result) - - @pytest.mark.parametrize( ("args", "kwargs"), ( @@ -460,92 +423,15 @@ class _Agentic(CustomLogger): return False, {} +class _ShortCircuit(WebSearchInterceptionLogger): + async def try_short_circuit_search(self, *args: object, **kwargs: object) -> dict[str, object] | None: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return None + + ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": "claude-opus-5-5"} STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": True} -@pytest.mark.parametrize( - ("feature", "model", "kwargs", "settings"), - ( - (Feature.MOCK_RESPONSE, "anthropic/claude-sonnet-4-5", {"mock_response": "hi"}, {}), - ( - Feature.PROXY_MCP_TOOLS, - "anthropic/claude-sonnet-4-5", - {"tools": [{"type": "mcp", "server_label": "gateway", "server_url": "litellm_proxy"}]}, - {}, - ), - (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), - (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), - (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), - (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), - (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), - ( - Feature.WORKLOAD_IDENTITY, - "anthropic/claude-sonnet-4-5", - {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, - {}, - ), - (Feature.CACHE_CONTROL_INJECTION, "anthropic/claude-sonnet-4-5", {"enable_prompt_caching": True}, {}), - ( - Feature.THINKING_SIGNATURE_RETRY, - "anthropic/claude-sonnet-4-5", - {"messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]}, - {}, - ), - ( - Feature.TOOL_CUSTOM_FIELD, - "anthropic/claude-sonnet-4-5", - {"tools": [{"name": "a", "input_schema": {"type": "object"}, "custom": {"eager_input_streaming": True}}]}, - {}, - ), - ( - Feature.STREAM_COST_INJECTION, - "anthropic/claude-sonnet-4-5", - {"stream": True}, - {"include_cost_in_streaming_usage": True}, - ), - ), -) -def test_each_detected_feature_reaches_its_rule( - monkeypatch: pytest.MonkeyPatch, - feature: Feature, - model: str, - kwargs: Mapping[str, object], - settings: Mapping[str, object], -) -> None: - rules: Final[Rules] = ( - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({feature})), - RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), - ) - python_result: Final = response(model) - native_result: Final = response(model) - - def served(call_kwargs: Mapping[str, object]) -> AnthropicMessagesResponse: - messages: Final = call_kwargs.get("messages", MESSAGES) - return _DISPATCH.run( - (16, messages, model), - {name: value for name, value in call_kwargs.items() if name != "messages"}, - python=lambda *call_args, **python_kwargs: python_result, - binding=messages_binding(lambda request: native_result), - native=lambda hook, request, call_args, native_kwargs: hook(request), - rules=rules, - ) - - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "cache", None) - monkeypatch.setattr(litellm, "ssl_security_level", None) - for setting in ("modify_params", "enable_anthropic_prompt_caching", "include_cost_in_streaming_usage"): - monkeypatch.setattr(litellm, setting, False) - for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE", "ANTHROPIC_FEDERATION_RULE_ID", "ANTHROPIC_ORGANIZATION_ID"): - monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("ANTHROPIC_API_KEY", "") - monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) - assert served({}) is native_result - for name, value in settings.items(): - monkeypatch.setattr(litellm, name, value) - assert served(kwargs) is python_result - - @pytest.mark.parametrize( "kwargs", ( @@ -580,27 +466,103 @@ def test_proxy_request_shape_stays_native(kwargs: Mapping[str, object]) -> None: @pytest.mark.asyncio @pytest.mark.parametrize( - "kwargs", + "block", ( - pytest.param({"tools": [STRICT_TOOL]}, id="drop-params-unset"), - pytest.param({"drop_params": False, "tools": [STRICT_TOOL]}, id="drop-params-off"), - pytest.param( - {"drop_params": True, "tools": [{"name": "a", "input_schema": {"type": "object"}}]}, id="no-strict" - ), + pytest.param({"type": "thinking", "thinking": "x", "signature": "s"}, id="thinking"), + pytest.param({"type": "redacted_thinking", "data": "x"}, id="redacted-thinking"), ), ) -async def test_strict_tools_stay_native_unless_python_could_drop_strict( +async def test_replayed_thinking_is_served_natively_under_the_shipped_rules( + monkeypatch: pytest.MonkeyPatch, block: Mapping[str, object] +) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "cache", None) + replay: Final = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [block, {"type": "text", "text": "hello"}]}, + {"role": "user", "content": "again"}, + ] + native_result: Final = response() + python_result: Final = response() + + async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + served: Final = await _ADISPATCH.arun( + (16, replay, "anthropic/claude-sonnet-4-5"), + {"api_key": "sk-ant-test"}, + python=python, + binding=amessages_binding(native), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=RULES, + ) + + assert served is native_result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("kwargs", "global_drop_params"), + ( + pytest.param({"drop_params": True, "tools": [STRICT_TOOL]}, False, id="drop-params-kwarg"), + pytest.param({"tools": [STRICT_TOOL]}, True, id="drop-params-global"), + pytest.param({"tools": [STRICT_TOOL]}, False, id="drop-params-unset"), + ), +) +async def test_strict_tools_are_served_natively_under_the_shipped_rules( + monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object], global_drop_params: bool +) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "cache", None) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + native_result: Final = response() + + async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + pytest.fail("strict tools stay on the native path") + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + served: Final = await _ADISPATCH.arun( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=python, + binding=amessages_binding(native), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=RULES, + ) + + assert served is native_result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kwargs", + ( + pytest.param({}, id="reads-and-writes"), + pytest.param({"cache": {"no-store": True}}, id="no-store"), + pytest.param({"cache": {"ttl": 60, "namespace": "team"}}, id="ttl-and-namespace"), + pytest.param({"stream": True}, id="stream"), + ), +) +async def test_a_configured_response_cache_stays_native( monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object] ) -> None: rules: Final[Rules] = ( - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.STRICT_TOOLS_RETRY})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=MESSAGES_PYTHON_FEATURES), RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), ) native_result: Final = response() - monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "cache", Cache()) async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub - pytest.fail("Python cannot drop strict here") + pytest.fail("Rust reads and writes the configured response cache") async def native(request: NativeCall) -> AnthropicMessagesResponse: return native_result @@ -617,45 +579,61 @@ async def test_strict_tools_stay_native_unless_python_could_drop_strict( assert served is native_result -ASYNC_ONLY_FEATURES: Final = ( - pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"), - pytest.param( - Feature.FOUNDRY_BETA_ROW, - {"api_base": "https://team.services.ai.azure.com/anthropic"}, +@pytest.mark.asyncio +async def test_modify_params_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=catalog.MESSAGES_PYTHON_FEATURES), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "modify_params", True) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("modify_params must not fall back to Python") + + async def apython(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + pytest.fail("modify_params must not fall back to Python") + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + call_args: Final = (200_000, MESSAGES, "anthropic/claude-sonnet-4-5") + served: Final = _DISPATCH.run( + call_args, {}, - id="foundry-beta-row", - ), - pytest.param( - Feature.WEBSEARCH_INTERCEPTION, {}, {"callbacks": [WebSearchInterceptionLogger()]}, id="websearch-interception" - ), - pytest.param( - Feature.MESSAGES_INTERCEPTOR, - {"model": "deepseek/deepseek-v4-flash", "tools": [ADVISOR_TOOL]}, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, args, kwargs: hook(request), + rules=rules, + ) + aserved: Final = await _ADISPATCH.arun( + call_args, {}, - id="messages-interceptor", - ), - pytest.param( - Feature.STRICT_TOOLS_RETRY, {"drop_params": True, "tools": [STRICT_TOOL]}, {}, id="strict-tools-retry" - ), - pytest.param( - Feature.STRICT_TOOLS_RETRY, {"tools": [STRICT_TOOL]}, {"drop_params": True}, id="strict-tools-retry-global" - ), -) + python=apython, + binding=amessages_binding(native), + native=lambda hook, request, args, kwargs: hook(request), + rules=rules, + ) + + assert (served, aserved) == (native_result, native_result) @pytest.mark.asyncio -@pytest.mark.parametrize(("feature", "kwargs", "settings"), ASYNC_ONLY_FEATURES) -async def test_handler_hooks_the_sync_entry_never_runs_keep_only_async_calls_on_python( - monkeypatch: pytest.MonkeyPatch, feature: Feature, kwargs: Mapping[str, object], settings: Mapping[str, object] -) -> None: +@pytest.mark.parametrize( + "kwargs", + ( + pytest.param({"mock_response": "hi"}, id="unary"), + pytest.param({"mock_response": "hi", "stream": True}, id="streamed"), + ), +) +async def test_mock_response_is_served_natively(monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object]) -> None: rules: Final[Rules] = ( - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({feature})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.MOCK_RESPONSE})), RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), ) python_result: Final = response() native_result: Final = response() - model: Final = str(kwargs.get("model", "anthropic/claude-sonnet-4-5")) - call_kwargs: Final = {name: value for name, value in kwargs.items() if name != "model"} async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub return python_result @@ -664,48 +642,331 @@ async def test_handler_hooks_the_sync_entry_never_runs_keep_only_async_calls_on_ return native_result monkeypatch.setattr(litellm, "callbacks", []) - for name, value in settings.items(): - monkeypatch.setattr(litellm, name, value) served_async: Final = await _ADISPATCH.arun( - (16, MESSAGES, model), - call_kwargs, + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, python=python, binding=amessages_binding(native), native=lambda hook, request, args, native_kwargs: hook(request), rules=rules, ) served_sync: Final = _DISPATCH.run( - (16, MESSAGES, model), - call_kwargs, + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, python=lambda *call_args, **python_kwargs: python_result, binding=messages_binding(lambda request: native_result), native=lambda hook, request, args, native_kwargs: hook(request), rules=rules, ) - assert served_async is python_result + assert served_async is native_result assert served_sync is native_result +class _Claims(MessagesInterceptor): + def can_handle(self, tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "anthropic" + + async def handle(self, **kwargs: object) -> AnthropicMessagesResponse: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: test double only needs the override + return response() + + @pytest.mark.asyncio -async def test_interceptor_feature_follows_the_resolved_provider() -> None: +async def test_a_call_a_registered_interceptor_owns_runs_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str] +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(messages_interceptors, "_interceptors", [_Claims()]) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("Rust hands the call to the interceptor itself") + + with rebound(dispatch, "_PYTHON_AMESSAGES", python): + result: Final = await _public(True, "anthropic/claude-sonnet-4-5", {"tools": [ADVISOR_TOOL]}) + + assert isinstance(result, dict) + assert shipped_native == ["async"] + + +def test_a_federated_request_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) rules: Final[Rules] = ( - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.MESSAGES_INTERCEPTOR})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.WORKLOAD_IDENTITY})), RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), ) native_result: Final = response() - async def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub - pytest.fail("anthropic serves the advisor tool itself") + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("Rust mints the federation token itself") - async def native(request: NativeCall) -> AnthropicMessagesResponse: - return native_result - - result: Final = await _ADISPATCH.arun( - (16, MESSAGES, "claude-sonnet-4-5"), - {"tools": [ADVISOR_TOOL], "custom_llm_provider": "anthropic"}, + served: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, python=python, - binding=amessages_binding(native), - native=lambda hook, request, call_args, call_kwargs: hook(request), + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), rules=rules, ) - assert result is native_result + assert served is native_result + + +def _shipped_native(calls: list[str]) -> None: # mutable-ok: records which native entry ran + def native(request: NativeCall) -> AnthropicMessagesResponse: + calls.append("sync") + return response() + + async def anative(request: NativeCall) -> AnthropicMessagesResponse: + calls.append("async") + return response() + + NATIVE_MESSAGES.override(native) + NATIVE_AMESSAGES.override(anative) + + +@pytest.fixture +def shipped_native() -> Iterator[list[str]]: + calls: Final[list[str]] = [] + _shipped_native(calls) + with routed(*RULES): + yield calls + NATIVE_MESSAGES.reset() + NATIVE_AMESSAGES.reset() + + +async def _public(asynchronous: bool, model: str, kwargs: Mapping[str, object]) -> object: + call: Final = {"max_tokens": 16, "messages": MESSAGES, "model": model, "api_key": "sk-test", **kwargs} + if asynchronous: + return await dispatch.anthropic_messages(**call) + return dispatch.anthropic_messages_handler(**call) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_anthropic_runs_on_rust_without_the_opt_in( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool +) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setattr(litellm, "callbacks", []) + + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", {}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize( + "model", ("anthropic/parity-tool-model", "anthropic/claude-sonnet-4-5"), ids=("non-claude", "claude") +) +async def test_proxy_classified_safeguards_run_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool, model: str +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(proxy_server, "general_settings", CLASSIFIER_SETTINGS) + + result: Final = await _public(asynchronous, model, {"safeguards": DANGEROUS_TOOL_USE}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize( + "thread", + ( + pytest.param({"type": "create"}, id="create"), + pytest.param({"type": "continue", "previous_message_id": "msg_1"}, id="continue"), + ), +) +async def test_a_message_thread_runs_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool, thread: Mapping[str, str] +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", {"thread": dict(thread)}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +async def test_anthropic_runs_pre_request_hooks_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str] +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_Rewrite()]) + + result: Final = await _public(True, "anthropic/claude-sonnet-4-5", {}) + + assert isinstance(result, dict) + assert shipped_native == ["async"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_anthropic_runs_a_web_search_short_circuit_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_ShortCircuit()]) + + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", {}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize( + ("kwargs", "settings"), + ( + pytest.param({"client": object()}, {}, id="caller-client"), + pytest.param({"client": AsyncHTTPHandler()}, {}, id="caller-async-handler"), + pytest.param({}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}, id="tls-security-level"), + pytest.param({}, {"ssl_ecdh_curve": "secp521r1"}, id="tls-ecdh-curve"), + ), +) +async def test_anthropic_serves_the_http_transport_features_on_rust( + monkeypatch: pytest.MonkeyPatch, + shipped_native: list[str], + asynchronous: bool, + kwargs: Mapping[str, object], + settings: Mapping[str, object], +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + for name, value in settings.items(): + monkeypatch.setattr(litellm, name, value) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("anthropic has no Python Messages path") + + with rebound(dispatch, "_PYTHON_MESSAGES", python), rebound(dispatch, "_PYTHON_AMESSAGES", python): + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", kwargs) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize( + "callbacks", + ( + pytest.param([_Agentic()], id="agentic-loop"), + pytest.param([WebSearchInterceptionLogger(enabled_providers=["anthropic"])], id="websearch-interception"), + ), +) +async def test_anthropic_serves_agentic_loop_callbacks_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool, callbacks: list[object] +) -> None: + monkeypatch.setattr(litellm, "callbacks", callbacks) + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("anthropic has no Python Messages path") + + with rebound(dispatch, "_PYTHON_MESSAGES", python), rebound(dispatch, "_PYTHON_AMESSAGES", python): + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", {}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize( + "server_url", + ( + pytest.param("litellm_proxy", id="whole-gateway"), + pytest.param("litellm_proxy/mcp/deepwiki", id="one-server"), + pytest.param("https://gateway.example/mcp/deepwiki", id="proxy-path"), + ), +) +async def test_anthropic_serves_gateway_mcp_tools_on_rust( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool, server_url: str +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + tools: Final = [{"type": "mcp", "server_label": "gateway", "server_url": server_url, "require_approval": "never"}] + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("anthropic has no Python Messages path") + + with rebound(dispatch, "_PYTHON_MESSAGES", python), rebound(dispatch, "_PYTHON_AMESSAGES", python): + result: Final = await _public(asynchronous, "anthropic/claude-sonnet-4-5", {"tools": tools}) + + assert isinstance(result, dict) + assert shipped_native == ["async" if asynchronous else "sync"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_other_providers_keep_python_for_the_same_features( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + python_result: Final = response("deepseek/deepseek-v4-flash") + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result + + async def apython(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result + + with rebound(dispatch, "_PYTHON_MESSAGES", python), rebound(dispatch, "_PYTHON_AMESSAGES", apython): + result: Final = await _public(asynchronous, "deepseek/deepseek-v4-flash", {"client": object()}) + + assert result is python_result + assert shipped_native == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "refused"), + ( + pytest.param("anthropic/claude-sonnet-4-5", True, id="anthropic"), + pytest.param("deepseek/deepseek-v4-flash", False, id="deepseek"), + ), +) +async def test_a_failing_detector_refuses_only_the_rust_only_provider( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], model: str, refused: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + python_result: Final = response(model) + + def failing(request: NativeCall) -> bool: + raise RuntimeError("detector failed") + + async def apython(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result + + with ( + rebound(dispatch, "_FEATURES", {Feature.MOCK_RESPONSE: failing}), + rebound(dispatch, "_PYTHON_AMESSAGES", apython), + ): + if not refused: + assert await _public(True, model, {}) is python_result + return + with pytest.raises(UnsupportedParamsError, match="does not support this call"): + await _public(True, model, {}) + assert shipped_native == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_an_anthropic_call_its_signature_rejects_keeps_python_s_type_error( + monkeypatch: pytest.MonkeyPatch, shipped_native: list[str], asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + call: Final = {"messages": MESSAGES, "model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"} + + async def attempt() -> object: + if asynchronous: + return await dispatch.anthropic_messages(**call) + return dispatch.anthropic_messages_handler(**call) + + with pytest.raises(TypeError, match="max_tokens"): + await attempt() + + assert shipped_native == [] diff --git a/tests/test_litellm_rust/messages/test_workload_identity.py b/tests/test_litellm_rust/messages/test_workload_identity.py new file mode 100644 index 00000000000..d2c1a85de91 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_workload_identity.py @@ -0,0 +1,219 @@ +import re +from collections.abc import Iterator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.llms.base_llm.auth.token_exchange import default_token_exchange_engine +from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + +IDENTITY_TOKEN: Final = "eyJ.native-identity.jwt" +MINTED_TOKEN: Final = "sk-ant-oat01-native-minted" +FEDERATION_IDS: Final = {"anthropic_federation_rule_id": "fdrl_native", "anthropic_organization_id": "org-native"} +MINTED: Final = ResponseSpec(body={"access_token": MINTED_TOKEN, "token_type": "Bearer", "expires_in": 3600}) +EXCHANGE: Final = { + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "federation_rule_id": "fdrl_native", + "organization_id": "org-native", + "assertion": IDENTITY_TOKEN, +} +FEDERATION_ENV_NAMES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_FEDERATION_RULE_ID", + "ANTHROPIC_ORGANIZATION_ID", + "ANTHROPIC_IDENTITY_TOKEN", + "ANTHROPIC_IDENTITY_TOKEN_FILE", + "ANTHROPIC_IDENTITY_SOURCE", + "LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", +) + + +@pytest.fixture +def server(server: RecordingServer) -> Iterator[RecordingServer]: + yield from relaxed(server) + + +@pytest.fixture +def federation(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + for name in FEDERATION_ENV_NAMES: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", IDENTITY_TOKEN) + monkeypatch.setenv("LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS", "127.0.0.1") + with ( + rebound(default_token_exchange_engine, "_entries", {}), + rebound(default_token_exchange_engine, "_shared_store", None), + ): + yield + + +def wire(server: RecordingServer) -> tuple[tuple[str, object, Mapping[str, str]], ...]: + return tuple((sent.path, sent.body, sent.headers) for sent in server.requests) + + +def assert_minted_bearer(headers: Mapping[str, str]) -> None: + assert headers.get("authorization") == f"Bearer {MINTED_TOKEN}" + assert "x-api-key" not in headers + assert "oauth-2025-04-20" in headers.get("anthropic-beta", "").split(",") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("federation") +@pytest.mark.parametrize( + ("kwargs", "env", "exchange"), + ( + pytest.param(FEDERATION_IDS, {}, EXCHANGE, id="ids-from-params"), + pytest.param( + {}, + {"ANTHROPIC_FEDERATION_RULE_ID": "fdrl_env", "ANTHROPIC_ORGANIZATION_ID": "org-env"}, + {**EXCHANGE, "federation_rule_id": "fdrl_env", "organization_id": "org-env"}, + id="ids-from-environment", + ), + pytest.param( + { + **FEDERATION_IDS, + "anthropic_federation_workspace_id": "wrkspc_native", + "anthropic_service_account_id": "svac_native", + }, + {}, + {**EXCHANGE, "workspace_id": "wrkspc_native", "service_account_id": "svac_native"}, + id="workspace-and-service-account", + ), + ), +) +async def test_workload_identity_mints_a_bearer_for_a_call_without_a_key( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + kwargs: Mapping[str, object], + env: Mapping[str, str], + exchange: Mapping[str, str], +) -> None: + for name, value in env.items(): + monkeypatch.setenv(name, value) + server.enqueue(MINTED) + + response: Final = await call(server, api_key=None, **kwargs) + + assert_served_natively(server, count=2) + (token_path, token_body, _), (messages_path, _, headers) = wire(server) + assert (token_path, token_body) == ("/v1/oauth/token", exchange) + assert messages_path == "/v1/messages" + assert_minted_bearer(headers) + assert isinstance(response, dict) + assert response["id"] == MESSAGES_RESPONSE["id"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("federation") +async def test_workload_identity_keeps_the_caller_betas_and_headers(server: RecordingServer) -> None: + server.enqueue(MINTED) + + await call( + server, + api_key=None, + extra_headers={"anthropic-beta": "files-api-2025-04-14", "x-trace-id": "trace-1"}, + **FEDERATION_IDS, + ) + + assert_served_natively(server, count=2) + _, (_, _, headers) = wire(server) + assert_minted_bearer(headers) + assert "files-api-2025-04-14" in headers["anthropic-beta"].split(",") + assert headers["x-trace-id"] == "trace-1" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("federation") +async def test_workload_identity_streams_with_the_minted_bearer(server: RecordingServer) -> None: + server.enqueue(MINTED) + server.enqueue(ResponseSpec(body={}, events=MESSAGES_EVENTS)) + + stream: Final = await call(server, api_key=None, stream=True, **FEDERATION_IDS) + chunks: Final = [chunk async for chunk in stream] # pyright: ignore[reportGeneralTypeIssues] # the public call returns an async stream + + assert_served_natively(server, count=2) + _, (_, _, headers) = wire(server) + assert_minted_bearer(headers) + assert b"message_stop" in b"".join(chunk if isinstance(chunk, bytes) else str(chunk).encode() for chunk in chunks) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("federation") +async def test_a_static_key_shadows_workload_identity(server: RecordingServer) -> None: + await call(server, api_key="sk-ant-static", **FEDERATION_IDS) + + assert_served_natively(server) + ((path, _, headers),) = wire(server) + assert path == "/v1/messages" + assert headers.get("x-api-key") == "sk-ant-static" + assert "authorization" not in headers + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("federation") +@pytest.mark.parametrize( + ("kwargs", "env", "upstream", "message", "sent"), + ( + pytest.param( + FEDERATION_IDS, + {}, + (ResponseSpec(body={"error": "authentication_error"}, status=401),), + "Anthropic workload identity federation failed. The token endpoint returned HTTP 401", + ("/v1/oauth/token",), + id="exchange-denied", + ), + pytest.param( + FEDERATION_IDS, + {"LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS": None}, + (), + "Anthropic workload identity federation refused to use host", + (), + id="untrusted-exchange-host", + ), + pytest.param( + {"anthropic_identity_token": "oidc/env/ANTHROPIC_IDENTITY_TOKEN"}, + {}, + (), + "anthropic_identity_token is set, but anthropic_federation_rule_id and anthropic_organization_id", + (), + id="identity-token-without-ids", + ), + pytest.param( + {**FEDERATION_IDS, "anthropic_disable_workload_identity_federation": True}, + {}, + (), + "Missing Anthropic API Key", + (), + id="federation-disabled", + ), + ), +) +async def test_workload_identity_failures_raise_before_the_messages_call( + server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + kwargs: Mapping[str, object], + env: Mapping[str, str | None], + upstream: tuple[ResponseSpec, ...], + message: str, + sent: tuple[str, ...], +) -> None: + for name, value in env.items(): + if value is None: + monkeypatch.delenv(name) + else: + monkeypatch.setenv(name, value) + for spec in upstream: + server.enqueue(spec) + + with pytest.raises(litellm.AuthenticationError, match=re.escape(message)) as raised: + await call(server, api_key=None, **kwargs) + + assert raised.value.llm_provider == "anthropic" + assert_served_natively(server, count=len(sent)) + assert tuple(path for path, _, _ in wire(server)) == sent diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 5743658abdf..2b5e6fe93a6 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -1,5 +1,6 @@ import asyncio import copy +import gzip import json import threading import time @@ -29,13 +30,31 @@ class ResponseSpec: events: tuple[tuple[str, object], ...] = () # Splits a non-stream body into this many writes, `chunk_delay` apart, so it trickles in. body_chunks: int = 1 + cut_after: int | None = None + cut_delay: float = 0 + gzip: bool = False + # Writes every stream event in one write, so they reach the client in one read. + burst: bool = False + # Sends each stream event as its own HTTP/1.1 chunk instead of under a Content-Length. + chunked: bool = False def payloads(self) -> tuple[bytes, ...]: if not self.events: - body: Final = self.body if isinstance(self.body, bytes) else json.dumps(self.body).encode() + raw: Final = self.body if isinstance(self.body, bytes) else json.dumps(self.body).encode() + body: Final = gzip.compress(raw, mtime=0) if self.gzip else raw size: Final = max(-(-len(body) // max(self.body_chunks, 1)), 1) return tuple(body[start : start + size] for start in range(0, len(body), size)) or (body,) - return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) + frames: Final = tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) + if self.gzip: + return (gzip.compress(b"".join(frames), mtime=0),) + return (b"".join(frames),) if self.burst and not self.chunked else frames + + def writes(self) -> tuple[bytes, ...]: + payloads: Final = self.payloads() + if not self.chunked: + return payloads + chunks: Final = (*(b"%x\r\n%s\r\n" % (len(payload), payload) for payload in payloads), b"0\r\n\r\n") + return (b"".join(chunks),) if self.burst else chunks @dataclass @@ -84,15 +103,28 @@ def recording_service() -> Iterator[RecordingServer]: response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) if response.delay: time.sleep(response.delay) - payloads: Final = response.payloads() + payloads: Final = response.writes() + if response.chunked: + self.protocol_version = "HTTP/1.1" + self.close_connection = True self.send_response(response.status) self.send_header("Content-Type", "text/event-stream" if response.events else "application/json") - self.send_header("Content-Length", str(sum(len(payload) for payload in payloads))) + if response.chunked: + self.send_header("Transfer-Encoding", "chunked") + self.send_header("Connection", "close") + else: + self.send_header("Content-Length", str(sum(len(payload) for payload in payloads))) + if response.gzip: + self.send_header("Content-Encoding", "gzip") for name, value in response.headers.items(): self.send_header(name, value) self.end_headers() try: for index, payload in enumerate(payloads): + if index == response.cut_after: + time.sleep(response.cut_delay) + self.close_connection = True + break if index and response.chunk_delay: time.sleep(response.chunk_delay) self.wfile.write(payload) diff --git a/tests/test_litellm_rust/support/served.py b/tests/test_litellm_rust/support/served.py index 0890077f745..5dfaa2cfaeb 100644 --- a/tests/test_litellm_rust/support/served.py +++ b/tests/test_litellm_rust/support/served.py @@ -6,11 +6,13 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None: assert len(server.requests) == (len(responses) or count) - if not responses: - assert not any( - request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests - ) + marked: Final = tuple(response for response in responses if not isinstance(response, dict)) + if not marked: + assert not any(request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests) for response in responses: + if all(response is not other for other in marked): + assert "x-litellm-rust" not in str(get_hidden_params_dict(response)) + for response in marked: additional_headers: Final = get_hidden_params_dict(response).get("additional_headers") assert isinstance(additional_headers, dict) assert additional_headers.get("x-litellm-rust") == "true" diff --git a/tests/test_litellm_rust/test_dispatch.py b/tests/test_litellm_rust/test_dispatch.py index 193866409a2..10cc05e0801 100644 --- a/tests/test_litellm_rust/test_dispatch.py +++ b/tests/test_litellm_rust/test_dispatch.py @@ -189,7 +189,7 @@ def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"other"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, models=frozenset({"other"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS})), - RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.STREAMING})), ), ids=("providers", "models", "targets", "features"), ) diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 02476545cae..ef2048998b1 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, budget_reservation_from_metadata, drop_params_env_flag, + env_flag, drop_params_flag, get_litellm_metadata_from_kwargs, get_or_create_metadata_bucket, @@ -497,6 +498,32 @@ def test_drop_params_env_flag_keeps_a_non_flag_value_on_with_a_warning(configure ) +@pytest.mark.parametrize( + ("configured", "expected"), + [("", False), ("false", False), ("False", False), ("0", False), ("no", False), ("true", True), (" 1 ", True)], +) +def test_env_flag_parses_modify_params_as_a_bool(configured: str, expected: bool, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="env-flag-test"): + assert ( + env_flag({"LITELLM_MODIFY_PARAMS": configured}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) + is expected + ) + assert caplog.text == "" + + +def test_env_flag_treats_a_non_flag_value_as_on_with_a_warning(caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="env-flag-test"): + assert ( + env_flag({"LITELLM_MODIFY_PARAMS": "enabled"}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) + is True + ) + assert "LITELLM_MODIFY_PARAMS='enabled' is not a flag value, treating it as on" in caplog.text + + +def test_env_flag_is_off_when_unset(): + assert env_flag({}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) is False + + class TestIsExpectedClientError: def test_status_ranges(self): from litellm.litellm_core_utils.core_helpers import is_expected_client_error diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index c1d878d5b09..51fbcaa54bf 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2376,6 +2376,7 @@ async def test_anthropic_messages_marks_litellm_params_async(): import litellm from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.anthropic.messages.native import handler as python_messages captured = {} logged = asyncio.Event() @@ -2390,7 +2391,7 @@ async def test_anthropic_messages_marks_litellm_params_async(): original_callbacks = getattr(litellm, "callbacks", []) try: litellm.callbacks = [logger] - await litellm.anthropic_messages( + await python_messages.anthropic_messages( max_tokens=100, messages=[{"role": "user", "content": "Hey"}], model="anthropic/claude-sonnet-4-5", diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 54f55aba46a..f434b1f4d9f 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -15,6 +15,7 @@ from httpx import Headers from pydantic import BaseModel, ConfigDict import litellm +from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.constants import ( ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, @@ -4509,7 +4510,7 @@ async def test_native_compaction_wire_roundtrip( return httpx.Response(200, json=body) route.mock(side_effect=respond) - call: Final = litellm.anthropic.messages.acreate if messages_api else litellm.acompletion + call: Final = python_messages.anthropic_messages if messages_api else litellm.acompletion params: Final = dict( model=f"{'openai/' if gateway else ''}anthropic/claude-sonnet-5", api_key="test", diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index ff150af324d..54b76756975 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -20,6 +20,7 @@ from pydantic import TypeAdapter, ValidationError from unittest.mock import AsyncMock, MagicMock, patch import litellm +from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.anthropic_interface import messages from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -1470,7 +1471,7 @@ async def test_native_messages_strips_replayed_provider_specific_fields_from_wir patch.object(httpx.AsyncClient, "send", fake_send), pytest.raises(litellm.exceptions.InternalServerError), ): - await litellm.anthropic.messages.acreate( + await python_messages.anthropic_messages( max_tokens=100, messages=messages, model="anthropic/claude-haiku-4-5-20251001", @@ -2451,7 +2452,7 @@ async def test_anthropic_messages_with_extra_headers(): mock_client.post = AsyncMock(return_value=mock_response) # Call the handler with extra_headers and our mocked client - response = await litellm.anthropic.messages.acreate( + response = await python_messages.anthropic_messages( messages=messages, api_key=api_key, model="claude-haiku-4-5-20251001", @@ -2515,7 +2516,7 @@ async def test_anthropic_messages_with_thinking(): mock_client.post = AsyncMock(return_value=mock_response) # Call the handler with extra_headers and our mocked client - response = await litellm.anthropic.messages.acreate( + response = await python_messages.anthropic_messages( messages=messages, api_key=api_key, model="claude-haiku-4-5-20251001", @@ -2785,7 +2786,7 @@ async def test_anthropic_messages_strict_tools_rejection_is_sent_again_without_s anthropic: Final = _RejectsStrictToolsOnce() tool: Final = {"name": "lookup_record", "input_schema": {"type": "object"}, "strict": True} - response: Final = await litellm.anthropic_messages( + response: Final = await python_messages.anthropic_messages( max_tokens=16, messages=[{"role": "user", "content": "Use lookup_record to find invoice 42"}], model="anthropic/claude-sonnet-4-6", diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py index 1d0d70d8e60..a6aade2e3cf 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_router.py @@ -18,6 +18,7 @@ import respx import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.router import Router from litellm.types.llms.anthropic import ( AnthropicMessagesTextParam, @@ -220,6 +221,11 @@ def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) +@pytest.fixture(autouse=True) +def _python_messages(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) + + @pytest.mark.asyncio async def test_router_aanthropic_messages_non_streaming_posts_anthropic_body(respx_mock): route: Final = respx_mock.post(url__regex=_ANTHROPIC_URL).mock( @@ -362,7 +368,7 @@ async def test_acreate_non_streaming_returns_dict_content_blocks( ): _set_fake_aws_env(monkeypatch) route: Final = respx_mock.post(url__regex=url).mock(return_value=httpx.Response(200, json=body)) - response: Final = await litellm.anthropic.messages.acreate( + response: Final = await python_messages.anthropic_messages( messages=[{"role": "user", "content": _PROMPT}], model=model, max_tokens=100, @@ -602,7 +608,7 @@ async def test_bedrock_converse_system_prompt_caching_returns_cache_tokens(respx }, ) ) - response: Final = await litellm.anthropic.messages.acreate( + response: Final = await python_messages.anthropic_messages( model=f"bedrock/converse/{_BEDROCK_SONNET}", messages=[{"role": "user", "content": "What are the key terms?"}], system=_cached_system(), @@ -632,7 +638,7 @@ async def test_bedrock_invoke_system_prompt_caching_returns_cache_tokens(respx_m }, ) ) - response: Final = await litellm.anthropic.messages.acreate( + response: Final = await python_messages.anthropic_messages( model=f"bedrock/invoke/{_BEDROCK_SONNET}", messages=[{"role": "user", "content": "What are the key terms?"}], system=_cached_system(), @@ -664,7 +670,7 @@ async def test_bedrock_streaming_message_start_carries_cache_usage_fields( route: Final = respx_mock.post(url__regex=url).mock( return_value=httpx.Response(200, content=body, headers={"content-type": "application/vnd.amazon.eventstream"}) ) - stream: Final = await litellm.anthropic.messages.acreate( + stream: Final = await python_messages.anthropic_messages( model=model, messages=[ { @@ -734,7 +740,7 @@ async def test_tool_search_forwards_deferred_tools_and_beta_header( }, ) ) - response: Final = await litellm.anthropic.messages.acreate( + response: Final = await python_messages.anthropic_messages( model="anthropic/claude-sonnet-4-5-20250929", messages=[{"role": "user", "content": prompt}], tools=[dict(tool) for tool in _tool_search_tools()], diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index 97d5b0550d4..d86c25c2558 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -16,7 +16,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) -from litellm.llms.anthropic.wif import get_anthropic_wif_token from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig @@ -27,12 +26,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass ) from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, AnthropicFinishReason from litellm.types.router import GenericLiteLLMParams -from tests.unit.llms.anthropic.test_anthropic_wif import ( - ScriptedPoster, - make_engine, - token_response, - write_token_file, -) +from tests.unit.llms.anthropic.test_anthropic_wif import write_token_file _WIF_PARAMS: Final[dict] = { "anthropic_federation_rule_id": "fdrl_abc123", @@ -165,18 +159,6 @@ def test_messages_request_never_defers_tools_off_claude( assert request["tools"] == [{"name": "A", "input_schema": _EMPTY_SCHEMA} if drop_params else wrapped_tool] -def test_workload_identity_allowed_for_anthropic() -> None: - assert AnthropicMessagesConfig()._allows_workload_identity is True - - -def test_workload_identity_blocked_for_minimax() -> None: - assert MinimaxMessagesConfig()._allows_workload_identity is False - - -def test_workload_identity_blocked_for_tencent() -> None: - assert TencentAnthropicMessagesConfig()._allows_workload_identity is False - - def test_minimax_validate_environment_never_attaches_anthropic_wif_credential( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -228,37 +210,6 @@ def test_tencent_validate_environment_never_attaches_anthropic_wif_credential( ) -def test_wif_token_exchange_reaches_only_anthropic_not_minimax_or_tencent( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """get_anthropic_wif_token's engine parameter is the only DI seam in the WIF minting chain; - validate_anthropic_messages_environment always uses the module's default engine, so this - drives that seam directly with the exact litellm_params AnthropicModelInfo.get_auth_header - would receive from each config, proving MiniMax/Tencent never reach the token endpoint even - when a mint would otherwise succeed.""" - monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_prod") - monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-prod-uuid") - monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) - token_file = write_token_file(tmp_path, "jwt-assertion-value") - litellm_params = {"anthropic_identity_token_file": str(token_file)} - poster = ScriptedPoster([token_response("sk-ant-oat01-canary")]) - engine = make_engine(poster) - - minted: Final = get_anthropic_wif_token( - litellm_params, - "https://api.anthropic.com", - "claude-sonnet-4-5", - engine, - ) - assert minted == "sk-ant-oat01-canary" - assert len(poster.requests) == 1 - - for config in (MinimaxMessagesConfig(), TencentAnthropicMessagesConfig()): - 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"}} diff --git a/tests/unit/llms/anthropic/messages/native/test_response_cache.py b/tests/unit/llms/anthropic/messages/native/test_response_cache.py index ac104b55bde..bc26d5612df 100644 --- a/tests/unit/llms/anthropic/messages/native/test_response_cache.py +++ b/tests/unit/llms/anthropic/messages/native/test_response_cache.py @@ -99,9 +99,9 @@ async def test_non_streaming_request_is_served_from_cache(local_cache, request_k fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - first = await litellm.anthropic_messages(**request_kwargs) + first = await handler.anthropic_messages(**request_kwargs) await asyncio.sleep(0) - second = await litellm.anthropic_messages(**request_kwargs) + second = await handler.anthropic_messages(**request_kwargs) assert len(fake_handler.calls) == 1 assert first == second @@ -115,9 +115,9 @@ async def test_cache_key_separates_different_system_prompts(local_cache, request fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - first = await litellm.anthropic_messages(**request_kwargs, system="Always answer ALPHA") + first = await handler.anthropic_messages(**request_kwargs, system="Always answer ALPHA") await asyncio.sleep(0) - second = await litellm.anthropic_messages(**request_kwargs, system="Always answer BETA") + second = await handler.anthropic_messages(**request_kwargs, system="Always answer BETA") assert len(fake_handler.calls) == 2 assert first["content"][0]["text"] == "ALPHA" @@ -130,9 +130,9 @@ async def test_cache_key_separates_anthropic_native_params(local_cache, request_ fake_handler = _CountingHandler([_anthropic_response("msg_1", "ALPHA"), _anthropic_response("msg_2", "BETA")]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - await litellm.anthropic_messages(**request_kwargs) + await handler.anthropic_messages(**request_kwargs) await asyncio.sleep(0) - await litellm.anthropic_messages(**request_kwargs, **anthropic_param) + await handler.anthropic_messages(**request_kwargs, **anthropic_param) assert len(fake_handler.calls) == 2 @@ -181,9 +181,9 @@ async def test_streaming_request_is_replayed_from_cache(local_cache, request_kwa fake_handler = _CountingHandler([_byte_stream(events), _byte_stream([b"event: never_used\n\n"])]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + first = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) await asyncio.sleep(0) - second_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + second_stream = await handler.anthropic_messages(**request_kwargs, stream=True) second = await _collect(second_stream) assert len(fake_handler.calls) == 1 @@ -197,8 +197,8 @@ async def test_streaming_cache_is_not_shared_with_non_streaming(local_cache, req fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _anthropic_response("msg_2", "ALPHA")]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) - non_streaming = await litellm.anthropic_messages(**request_kwargs) + await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) + non_streaming = await handler.anthropic_messages(**request_kwargs) assert len(fake_handler.calls) == 2 assert non_streaming["content"][0]["text"] == "ALPHA" @@ -219,11 +219,11 @@ async def test_stream_without_output_is_not_cached(local_cache, request_kwargs, ) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - empty = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + empty = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) await asyncio.sleep(0) - refilled = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + refilled = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) await asyncio.sleep(0) - replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert empty == empty_events assert refilled == STREAM_EVENTS @@ -239,8 +239,8 @@ async def test_failed_stream_is_not_cached(local_cache, request_kwargs, monkeypa fake_handler = _CountingHandler([_byte_stream(error_events), _byte_stream(STREAM_EVENTS)]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - failed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) - replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + failed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert failed == error_events assert len(fake_handler.calls) == 2 @@ -260,9 +260,9 @@ async def test_multibyte_utf8_split_across_chunks_streams_and_caches(local_cache fake_handler = _CountingHandler([_byte_stream(chunks), _byte_stream([b"event: never_used\n\n"])]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + first = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) await asyncio.sleep(0) - second = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + second = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert len(fake_handler.calls) == 1 assert first == chunks @@ -278,9 +278,9 @@ async def test_message_stop_split_across_chunks_still_caches(local_cache, reques fake_handler = _CountingHandler([_byte_stream(chunks), _byte_stream([b"event: never_used\n\n"])]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - first = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + first = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) await asyncio.sleep(0) - second = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + second = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert len(fake_handler.calls) == 1 assert first == chunks @@ -296,8 +296,8 @@ async def test_error_event_split_across_chunks_is_not_cached(local_cache, reques fake_handler = _CountingHandler([_byte_stream(chunks), _byte_stream(STREAM_EVENTS)]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - failed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) - replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + failed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert failed == chunks assert len(fake_handler.calls) == 2 @@ -309,11 +309,11 @@ async def test_abandoned_stream_is_not_cached(local_cache, request_kwargs, monke fake_handler = _CountingHandler([_byte_stream(STREAM_EVENTS), _byte_stream(STREAM_EVENTS)]) monkeypatch.setattr(handler, "anthropic_messages_handler", fake_handler) - partial_stream = await litellm.anthropic_messages(**request_kwargs, stream=True) + partial_stream = await handler.anthropic_messages(**request_kwargs, stream=True) await partial_stream.__anext__() await partial_stream.aclose() - replayed = await _collect(await litellm.anthropic_messages(**request_kwargs, stream=True)) + replayed = await _collect(await handler.anthropic_messages(**request_kwargs, stream=True)) assert len(fake_handler.calls) == 2 assert replayed == STREAM_EVENTS @@ -488,7 +488,5 @@ def test_cache_writer_forwards_has_buffered_provider_output(request_kwargs): ) held_back = AnthropicMessagesStreamCacheWriter(stream=_HeldBackStream(), caching_handler=caching_handler) assert held_back.has_buffered_provider_output is True - replayable = AnthropicMessagesStreamCacheWriter( - stream=_byte_stream(STREAM_EVENTS), caching_handler=caching_handler - ) + replayable = AnthropicMessagesStreamCacheWriter(stream=_byte_stream(STREAM_EVENTS), caching_handler=caching_handler) assert replayable.has_buffered_provider_output is False diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 4e8dc69fd01..d105a7958b4 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -3266,23 +3266,6 @@ class TestWifLitellmParamsPlumbing: assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" assert len(poster.requests) == 1 - def test_messages_mints_from_litellm_params(self, wif_engine): - from litellm.llms.anthropic.messages.native.transformation import ( - AnthropicMessagesConfig, - ) - - poster, _ = wif_engine - headers, _ = AnthropicMessagesConfig().validate_anthropic_messages_environment( - headers={}, - model="claude-sonnet-4-5", - messages=[{"role": "user", "content": "Hello"}], - optional_params={}, - litellm_params=dict(WIF_PARAMS_ONLY), - ) - - assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert len(poster.requests) == 1 - class TestWifTokenUrlParity: """Both credential tiers must derive the SAME clean token URL from any form of @@ -3344,27 +3327,6 @@ class TestWifAsyncSeam: assert poster.thread_ids == [poster.thread_ids[0]] assert poster.thread_ids[0] != threading.get_ident() - @pytest.mark.asyncio - async def test_avalidate_messages_environment_mints_off_loop(self, wif_async_engine, monkeypatch): - from litellm.llms.anthropic.messages.native.transformation import ( - AnthropicMessagesConfig, - ) - - poster, sync_calls = wif_async_engine - monkeypatch.setenv("WIF_PARAMS_TEST_TOKEN", "params-jwt") - - headers, _ = await AnthropicMessagesConfig().avalidate_anthropic_messages_environment( - headers={}, - model="claude-sonnet-4-5", - messages=[{"role": "user", "content": "Hello"}], - optional_params={}, - litellm_params=dict(WIF_PARAMS_ONLY), - ) - - assert headers["authorization"] == f"Bearer {FAKE_MINTED_TOKEN}" - assert sync_calls == [] - assert poster.thread_ids[0] != threading.get_ident() - @pytest.mark.asyncio async def test_avalidate_delegates_to_subclass_sync_override(self): """A provider subclass that only overrides the sync method must keep its diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 6a946c2e75c..e07695fcd03 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -11,6 +11,7 @@ import respx from pydantic import JsonValue import litellm +from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.caching.dual_cache import DualCache from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.anthropic.count_tokens import handler as count_handler @@ -165,6 +166,7 @@ async def test_environment_credential_matches_native_count_and_observed_scope( monkeypatch.setenv("LIT7658_PROVIDER_KEY", _KEY) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) params: Final = { "model": f"anthropic/{_MODEL}", "api_key": "os.environ/LIT7658_PROVIDER_KEY", diff --git a/tests/unit/llms/anthropic/test_messages_http_error_retry_cases.py b/tests/unit/llms/anthropic/test_messages_http_error_retry_cases.py new file mode 100644 index 00000000000..a9edb2160f7 --- /dev/null +++ b/tests/unit/llms/anthropic/test_messages_http_error_retry_cases.py @@ -0,0 +1,188 @@ +import copy +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.llms.anthropic.common_utils import ( + ToolsWithoutStrict, + is_anthropic_invalid_thinking_block_error, + is_anthropic_thinking_binding_mismatch_error, + is_claude_strict_tools_rejection, + tools_without_strict, +) +from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig +from litellm.llms.deepseek.messages.transformation import DeepSeekAnthropicMessagesConfig +from litellm.llms.base_llm.anthropic_messages.transformation import BaseAnthropicMessagesConfig + +FIXTURE_PATH: Final = ( + Path(__file__).parents[4] / "litellm-rust/crates/llms/tests/fixtures/anthropic_messages_http_error_retry.json" +) + + +class _TextCase(TypedDict): + name: ReadOnly[str] + error_text: ReadOnly[str] + expected: ReadOnly[bool] + + +class _RetryCase(TypedDict): + name: ReadOnly[str] + status: ReadOnly[int] + error_text: ReadOnly[str] + thinking: ReadOnly[Mapping[str, object]] + retried: ReadOnly[bool] + stripped: ReadOnly[bool] + + +class _RetryTable(TypedDict): + request: ReadOnly[Mapping[str, object]] + stripped_request: ReadOnly[Mapping[str, object]] + cases: ReadOnly[tuple[_RetryCase, ...]] + + +class _StrictRejectionCase(TypedDict): + name: ReadOnly[str] + status: ReadOnly[int] + error_text: ReadOnly[str] + expected: ReadOnly[bool] + + +class _ToolsWithoutStrictExpected(TypedDict): + tools: ReadOnly[tuple[Mapping[str, object], ...]] + dropped_count: ReadOnly[int] + + +class _ToolsWithoutStrictCase(TypedDict): + name: ReadOnly[str] + tools: ReadOnly[object] + expected: ReadOnly[_ToolsWithoutStrictExpected | None] + + +class _StrictRetryCase(TypedDict): + name: ReadOnly[str] + status: ReadOnly[int] + error_text: ReadOnly[str] + thinking: ReadOnly[Mapping[str, object]] + drops_unsupported_params: ReadOnly[bool] + retried: ReadOnly[bool] + stripped: ReadOnly[bool] + strict_dropped: ReadOnly[bool] + + +class _StrictRetryTable(TypedDict): + tools: ReadOnly[tuple[Mapping[str, object], ...]] + tools_without_strict: ReadOnly[tuple[Mapping[str, object], ...]] + dropped: ReadOnly[int] + cases: ReadOnly[tuple[_StrictRetryCase, ...]] + + +class _Fixture(TypedDict): + invalid_thinking_block_error: ReadOnly[tuple[_TextCase, ...]] + thinking_binding_mismatch_error: ReadOnly[tuple[_TextCase, ...]] + http_error_retry: ReadOnly[_RetryTable] + claude_strict_tools_rejection: ReadOnly[tuple[_StrictRejectionCase, ...]] + tools_without_strict: ReadOnly[tuple[_ToolsWithoutStrictCase, ...]] + strict_tools_retry: ReadOnly[_StrictRetryTable] + + +FIXTURE: Final = TypeAdapter(_Fixture).validate_json(FIXTURE_PATH.read_text()) +RETRY: Final = FIXTURE["http_error_retry"] +STRICT_RETRY: Final = FIXTURE["strict_tools_retry"] + + +@pytest.mark.parametrize( + "case", [pytest.param(case, id=case["name"]) for case in FIXTURE["invalid_thinking_block_error"]] +) +def test_invalid_thinking_block_errors_match_the_shared_cases(case: _TextCase) -> None: + assert is_anthropic_invalid_thinking_block_error(case["error_text"]) is case["expected"] + + +@pytest.mark.parametrize( + "case", [pytest.param(case, id=case["name"]) for case in FIXTURE["thinking_binding_mismatch_error"]] +) +def test_binding_mismatch_errors_match_the_shared_cases(case: _TextCase) -> None: + assert is_anthropic_thinking_binding_mismatch_error(case["error_text"]) is case["expected"] + + +@pytest.mark.parametrize( + "case", [pytest.param(case, id=case["name"]) for case in FIXTURE["claude_strict_tools_rejection"]] +) +def test_strict_tools_rejections_match_the_shared_cases(case: _StrictRejectionCase) -> None: + assert is_claude_strict_tools_rejection(case["status"], case["error_text"]) is case["expected"] + + +@pytest.mark.parametrize("case", [pytest.param(case, id=case["name"]) for case in FIXTURE["tools_without_strict"]]) +def test_tools_without_strict_matches_the_shared_cases(case: _ToolsWithoutStrictCase) -> None: + expected: Final = case["expected"] + + assert tools_without_strict(copy.deepcopy(case["tools"])) == ( + None + if expected is None + else ToolsWithoutStrict(tools=tuple(expected["tools"]), dropped_count=expected["dropped_count"]) + ) + + +def _rejection(case: _RetryCase | _StrictRetryCase) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + return httpx.HTTPStatusError( + "rejected", + request=request, + response=httpx.Response(case["status"], text=case["error_text"], request=request), + ) + + +@pytest.mark.parametrize( + "config", + [AnthropicMessagesConfig(), DeepSeekAnthropicMessagesConfig()], + ids=["anthropic", "deepseek"], +) +@pytest.mark.parametrize("case", [pytest.param(case, id=case["name"]) for case in RETRY["cases"]]) +def test_the_http_error_retry_matches_the_shared_cases(config: BaseAnthropicMessagesConfig, case: _RetryCase) -> None: + body: Final = {**copy.deepcopy(RETRY["request"]), "thinking": case["thinking"]} + + retried: Final = config.should_retry_anthropic_messages_on_http_error( + e=_rejection(case), litellm_params={}, request_body=body + ) + transformed: Final = config.transform_anthropic_messages_request_on_http_error( + e=_rejection(case), request_data=copy.deepcopy(body) + ) + + assert retried is case["retried"] + assert transformed == (RETRY["stripped_request"] if case["stripped"] else body) + + +@pytest.mark.parametrize( + "config", + [AnthropicMessagesConfig(), DeepSeekAnthropicMessagesConfig()], + ids=["anthropic", "deepseek"], +) +@pytest.mark.parametrize("case", [pytest.param(case, id=case["name"]) for case in STRICT_RETRY["cases"]]) +def test_the_strict_tools_retry_matches_the_shared_cases( + monkeypatch: pytest.MonkeyPatch, config: BaseAnthropicMessagesConfig, case: _StrictRetryCase +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + body: Final = { + **copy.deepcopy(RETRY["request"]), + "thinking": case["thinking"], + "tools": copy.deepcopy(list(STRICT_RETRY["tools"])), + } + expected: Final = { + **copy.deepcopy(RETRY["stripped_request"] if case["stripped"] else body), + "tools": list(STRICT_RETRY["tools_without_strict"] if case["strict_dropped"] else STRICT_RETRY["tools"]), + } + + retried: Final = config.should_retry_anthropic_messages_on_http_error( + e=_rejection(case), litellm_params={"drop_params": case["drops_unsupported_params"]}, request_body=body + ) + transformed: Final = config.transform_anthropic_messages_request_on_http_error( + e=_rejection(case), request_data=copy.deepcopy(body) + ) + + assert retried is case["retried"] + assert transformed == expected diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py index 991a0c36f03..4d4a92f632e 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py @@ -4736,7 +4736,10 @@ async def test_bedrock_messages_fallback_to_anthropic_sends_the_tool_names_bedro unnamed_tool_extras: dict[str, object], sent_extras: dict[str, object], ) -> None: + from litellm.llms.anthropic.messages.native import handler as python_messages + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) upstream: Final = BedrockDownAnthropicUp() respx_mock.post(url__regex=r"https://.*").mock(side_effect=upstream) router: Final = litellm.Router( diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 4d8ba26d8d8..08f38352735 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -19,6 +19,7 @@ from botocore.credentials import RefreshableCredentials from h2.events import ConnectionTerminated import litellm +from litellm.llms.anthropic.messages.native import handler as python_messages from litellm._logging import verbose_logger from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, @@ -6493,7 +6494,7 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(anthropic_answers)) try: - await litellm.anthropic.messages.acreate( + await python_messages.anthropic_messages( max_tokens=16, messages=[{"role": "user", "content": "hi"}], model="anthropic/claude-sonnet-4-6", @@ -6546,7 +6547,6 @@ async def _search_through(transport: httpx.MockTransport, is_async: bool) -> obj @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) @pytest.mark.parametrize( ("api_base", "explicit_header", "expected_session"), [ @@ -6557,7 +6557,6 @@ async def _search_through(transport: httpx.MockTransport, is_async: bool) -> obj ) async def test_messages_opencode_session_header( stream: bool, - environment_base: str | None, api_base: str, explicit_header: str | None, expected_session: str | None, @@ -6565,8 +6564,6 @@ async def test_messages_opencode_session_header( ) -> None: monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) - if environment_base: - monkeypatch.setenv(environment_base, "https://opencode.ai/inference/go/anthropic") received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() payload: Final = { "id": "msg_test", @@ -6594,11 +6591,11 @@ async def test_messages_opencode_session_header( upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) try: - result: Final = await litellm.anthropic.messages.acreate( + result: Final = await python_messages.anthropic_messages( model="anthropic/claude-haiku-5-5", messages=[{"role": "user", "content": "hello"}], max_tokens=128000, - api_base=None if environment_base and api_base == "https://opencode.ai/inference/go/anthropic" else api_base, + api_base=api_base, api_key="test-key", stream=stream, litellm_session_id="conversation-123", @@ -6737,20 +6734,6 @@ def _anthropic_messages_logging_obj(): @pytest.mark.parametrize( ("custom_llm_provider", "messages_config", "foundry_api_base", "forwarded_betas"), [ - pytest.param( - "anthropic", - "anthropic", - "https://my-resource.services.ai.azure.com/anthropic", - f"thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}", - id="anthropic-drops-what-the-azure-ai-row-rejects", - ), - pytest.param( - "anthropic", - "anthropic", - "https://my-resource.cognitiveservices.azure.com/anthropic", - f"thinking-binding-controls-2026-08-01,{_BETA_THE_TABLE_DOES_NOT_LIST}", - id="anthropic-on-a-cognitiveservices-host", - ), pytest.param( "azure_ai", "azure_ai", @@ -6775,7 +6758,6 @@ async def test_async_anthropic_messages_handler_picks_the_beta_row_by_deployment forwarded_betas: str, ) -> None: from litellm import anthropic_beta_headers_manager - from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig @@ -6783,7 +6765,6 @@ async def test_async_anthropic_messages_handler_picks_the_beta_row_by_deployment monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) assert _BETA_THE_TABLE_DOES_NOT_LIST not in anthropic_beta_headers_manager._load_beta_headers_config()["azure_ai"] configs: Final = { - "anthropic": AnthropicMessagesConfig, "azure_ai": AzureAnthropicMessagesConfig, "pass_through": OpenAILikeAnthropicMessagesConfig, } @@ -6817,7 +6798,7 @@ async def test_first_foundry_request_with_betas_fetches_the_beta_table_off_the_e monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm import anthropic_beta_headers_manager, anthropic_beta_headers_url - from litellm.llms.anthropic.messages.native.transformation import AnthropicMessagesConfig + from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) monkeypatch.delenv("LITELLM_OFFLINE_MODE", raising=False) @@ -6835,9 +6816,9 @@ async def test_first_foundry_request_with_betas_fetches_the_beta_table_off_the_e await BaseLLMHTTPHandler().async_anthropic_messages_handler( model="claude-sonnet-4-6", messages=[{"role": "user", "content": "Hello"}], - anthropic_messages_provider_config=AnthropicMessagesConfig(), + anthropic_messages_provider_config=AzureAnthropicMessagesConfig(), anthropic_messages_optional_request_params={"max_tokens": 8}, - custom_llm_provider="anthropic", + custom_llm_provider="azure_ai", litellm_params=GenericLiteLLMParams(api_base=foundry_api_base), logging_obj=_anthropic_messages_logging_obj(), api_key="sk-test", diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index 12aa850afb6..86a7e03db84 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -1150,3 +1150,10 @@ async def test_a_turn_routed_to_a_non_anthropic_tier_says_its_provider_is_unsupp captured: Final = _observation(await rig.capture.payload()) assert captured.observation.outcome == "uncertain" assert captured.observation.reason == "unsupported_selected_provider" + + +@pytest.fixture(autouse=True) +def _python_messages(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.messages.native import handler as python_messages + + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) diff --git a/tests/unit/responses/mcp/test_tool_precheck.py b/tests/unit/responses/mcp/test_tool_precheck.py index bd37114ccea..04670f6b4ed 100644 --- a/tests/unit/responses/mcp/test_tool_precheck.py +++ b/tests/unit/responses/mcp/test_tool_precheck.py @@ -36,16 +36,6 @@ ANTHROPIC_FUNCTION_TOOL: Final = { "description": "No operation", "input_schema": {"type": "object", "properties": {}}, } -ANTHROPIC_MESSAGE_PAYLOAD: Final = { - "id": "msg_1", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-5-5", - "content": [{"type": "text", "text": "ok"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 1}, -} GATEWAY_MCP_TOOL: Final = {"type": "mcp", "server_label": "litellm", "server_url": "litellm_proxy"} GATEWAY_SERVER_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "litellm_proxy/mcp/github"} GATEWAY_PATH_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "https://gw.example.com/mcp/github"} @@ -193,23 +183,16 @@ def test_responses_with_gateway_mcp_tool_still_reaches_the_handler(): @pytest.mark.asyncio @pytest.mark.usefixtures("gateway_handler_unavailable") -async def test_anthropic_messages_with_function_tool_does_not_import_the_gateway_handler( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter -): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - upstream: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( - return_value=httpx.Response(200, json=ANTHROPIC_MESSAGE_PAYLOAD) - ) - +async def test_anthropic_messages_with_function_tool_does_not_import_the_gateway_handler(): response: Final = await litellm.anthropic_interface.messages.acreate( model="anthropic/claude-sonnet-5-5", messages=[{"role": "user", "content": "hi"}], max_tokens=16, tools=[ANTHROPIC_FUNCTION_TOOL], api_key="sk-test", + mock_response="ok", ) - assert upstream.called assert response["content"][0]["text"] == "ok" diff --git a/tests/unit/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py index 850accf574b..3fed862ae50 100644 --- a/tests/unit/router_strategy/complexity_router/test_context_compaction.py +++ b/tests/unit/router_strategy/complexity_router/test_context_compaction.py @@ -503,3 +503,19 @@ async def test_retry_reuses_summary_or_terminal_cancellation(outcome: Literal["s assert calls.qsize() == 1 and stopped.is_set() finally: compaction_executor.reset(token) + + +@pytest.fixture(autouse=True) +def _python_messages(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import catalog + from litellm.rust_bridge.configuration import Rollout + + monkeypatch.setattr( + catalog, + "RULES", + tuple( + rule + for rule in catalog.RULES + if not (rule.route is catalog.Route.MESSAGES and rule.rollout is Rollout.RUST_REQUIRED) + ), + ) diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 054c48bbb5e..0190a7e929d 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -15783,6 +15783,21 @@ class TestModalityRouting: @pytest.mark.usefixtures("local_model_cost_map") class TestHealthFallbackDispatch: + @pytest.fixture(autouse=True) + def _python_messages(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import catalog + from litellm.rust_bridge.configuration import Rollout + + monkeypatch.setattr( + catalog, + "RULES", + tuple( + rule + for rule in catalog.RULES + if not (rule.route is catalog.Route.MESSAGES and rule.rollout is Rollout.RUST_REQUIRED) + ), + ) + @pytest.mark.asyncio @pytest.mark.parametrize("classifier", ("capability", "llm_v2")) @pytest.mark.parametrize("calibrated", (False, True), ids=("raw", "calibrated")) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 47d58a3b869..1bfe88275c2 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -2630,3 +2630,19 @@ async def test_run_async_fallback_never_re_reserves_the_requested_model_group() with pytest.raises(Exception, match="model group 'primary' failed"): await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) assert checked == ["fallback"] + + +@pytest.fixture(autouse=True) +def _python_messages(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import catalog + from litellm.rust_bridge.configuration import Rollout + + monkeypatch.setattr( + catalog, + "RULES", + tuple( + rule + for rule in catalog.RULES + if not (rule.route is catalog.Route.MESSAGES and rule.rollout is Rollout.RUST_REQUIRED) + ), + ) diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index aa7dcc74841..ed60951002d 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -1,10 +1,29 @@ +import datetime +import functools +import json +from collections.abc import Iterator, Mapping, Sequence from types import MappingProxyType from typing import Final +import httpx import pytest +from pydantic import ValidationError import litellm +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases +from litellm.llms.anthropic.messages.native import interceptors, mcp_handler +from litellm.llms.anthropic.messages.native.interceptors.base import MessagesInterceptor +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + exclude_rsa_key_exchange_and_chacha20, + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params dict is untyped +) +from litellm.rust_bridge.host import callbacks_legacy_python from litellm.rust_bridge.host.routes import messages as messages_host from litellm.rust_bridge.host.routes.messages import response @@ -30,11 +49,26 @@ def test_response_is_a_detached_public_messages_dict() -> None: assert "_hidden_params" not in native +@pytest.mark.parametrize( + ("disabled", "joined"), + [pytest.param("true", False, id="httpx_transport"), pytest.param("false", True, id="aiohttp_transport")], +) +def test_settings_join_buffered_stream_reads_only_on_the_aiohttp_transport( + monkeypatch: pytest.MonkeyPatch, disabled: str, joined: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", disabled) + + assert messages_host.settings({})["joined_reads"] is joined + + def test_settings_project_caller_configuration_without_resolving_a_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", False) monkeypatch.setattr(litellm, "reasoning_auto_summary", True) monkeypatch.setattr(litellm, "request_timeout", 1800) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) projected: Final = messages_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) @@ -43,9 +77,44 @@ def test_settings_project_caller_configuration_without_resolving_a_model(monkeyp "reasoning_auto_summary": True, "additional_drop_params": ("metadata.user_id",), "request_timeout": 1800.0, + "drops_unsupported_params": True, + "model": None, + "mock_response": None, + "stream_cost_model": None, + "extra_headers": None, + "mcp_gateway": False, + "ascii_response_headers": True, + "joined_reads": True, + "ascii_request_headers": True, } +@pytest.mark.parametrize( + ("global_drop_params", "kwargs", "drop_params", "drops_unsupported_params"), + [ + pytest.param(True, {}, True, True, id="global-true"), + pytest.param(1, {}, True, False, id="global-truthy-non-bool"), + pytest.param(False, {"drop_params": "True"}, True, True, id="kwarg-flag-string"), + pytest.param(False, {"drop_params": 1}, True, True, id="kwarg-one"), + pytest.param(False, {"drop_params": "maybe"}, False, False, id="kwarg-not-a-flag"), + pytest.param(False, {}, False, False, id="unset"), + ], +) +def test_strict_tools_retry_reads_drop_params_as_the_python_retry_check_does( + monkeypatch: pytest.MonkeyPatch, + global_drop_params: object, + kwargs: Mapping[str, object], + drop_params: bool, + drops_unsupported_params: bool, +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + + projected: Final = messages_host.settings(kwargs) + + assert projected["drop_params"] is drop_params + assert projected["drops_unsupported_params"] is drops_unsupported_params + + @pytest.mark.parametrize( ("timeout", "explicitly_set", "expected"), [ @@ -170,7 +239,7 @@ def test_body_params_round_effort_to_the_levels_the_deployment_declares( "model_info": {"id": "deployment-1", "reasoning_effort_levels": levels}, } - assert messages_host.body_params(request)["output_config"] == {"effort": expected} + assert messages_host.body_params(request, asynchronous=True)["output_config"] == {"effort": expected} def test_body_params_keep_the_requested_effort_without_deployment_levels() -> None: @@ -181,7 +250,7 @@ def test_body_params_keep_the_requested_effort_without_deployment_levels() -> No "output_config": {"effort": "max"}, } - assert messages_host.body_params(request)["output_config"] == {"effort": "max"} + assert messages_host.body_params(request, asynchronous=True)["output_config"] == {"effort": "max"} @pytest.mark.parametrize( @@ -203,7 +272,28 @@ def test_body_params_derive_no_effort_from_a_param_the_deployment_lists( "model_info": {"id": "deployment-1", "reasoning_effort_levels": ["low", "medium", "high"]}, } - assert set(messages_host.body_params(request)).difference({"messages", "model"}) == expected_keys + assert set(messages_host.body_params(request, asynchronous=True)).difference({"messages", "model"}) == expected_keys + + +@pytest.mark.parametrize( + ("thread", "expected"), + [ + ({"type": "create"}, {"type": "create"}), + ({"type": "continue", "previous_message_id": "msg_1"}, {"type": "continue", "previous_message_id": "msg_1"}), + (None, None), + ], +) +def test_body_params_forward_the_thread_like_the_python_handler( + thread: dict[str, str] | None, expected: dict[str, str] | None +) -> None: + request: Final[dict[str, object]] = { + "model": "claude-sonnet-4-5", + "messages": [], + "max_tokens": 8, + "thread": thread, + } + + assert messages_host.body_params(request, asynchronous=True).get("thread") == expected def _marked(phase: str | None = None, *args: object) -> ValueError: @@ -213,13 +303,40 @@ def _marked(phase: str | None = None, *args: object) -> ValueError: return error -def test_a_stalled_stream_body_reaches_the_caller_as_the_transport_timeout() -> None: +@pytest.mark.parametrize( + ("disabled", "text"), ((None, "Timeout on reading data from socket"), ("true", "")), ids=("aiohttp", "httpx") +) +def test_a_stalled_stream_body_reaches_the_caller_as_the_transport_timeout( + monkeypatch: pytest.MonkeyPatch, disabled: str | None, text: str +) -> None: import httpx + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + if disabled is None: + monkeypatch.delenv("DISABLE_AIOHTTP_TRANSPORT", raising=False) + else: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", disabled) + mapped: Final = messages_host.map_failure(_marked("Body"), {"model": _ADAPTIVE_OPUS}, "anthropic") assert isinstance(mapped, httpx.ReadTimeout) - assert str(mapped) == "Timeout on reading data from socket" + assert str(mapped) == text + + +@pytest.mark.parametrize( + ("value", "received", "text"), + ( + ("h\ud800".encode("utf-8", "surrogatepass"), False, "can't encode character '\\ud800' in position 1"), + ("caf\u00e9".encode(), True, "can't encode character '\\xe9' in position 3"), + (b"caf\xe9", True, "can't encode character '\\udce9' in position 3"), + ), + ids=("sent-surrogate", "received-utf-8", "received-latin-1"), +) +def test_an_unencodable_header_is_the_error_pythons_ascii_codec_raises(value: bytes, received: bool, text: str) -> None: + error: Final = messages_host.unencodable_header(value, received) + + assert isinstance(error, UnicodeEncodeError) + assert text in str(error) def test_a_native_read_timeout_maps_to_the_public_timeout_with_python_text() -> None: @@ -365,12 +482,47 @@ def test_a_stream_body_cut_reaches_the_caller_as_the_transport_raised_it( native: Final = RuntimeError("upstream network error") native.messages_body_cut = cut # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets - mapped: Final = messages_host.map_failure(native, {"model": _ADAPTIVE_OPUS}, "anthropic") + mapped: Final = messages_host.map_failure(native, {"model": _ADAPTIVE_OPUS, "stream": True}, "anthropic") assert type(mapped) is getattr(httpx, error_class) assert str(mapped) == text +@pytest.mark.parametrize( + ("aiohttp", "cut", "text"), + ( + ( + True, + (19, 100, False), + "Response payload is not completed: ", + ), + ( + False, + (19, 100, False), + "peer closed connection without sending complete message body (received 19 bytes, expected 100)", + ), + ), +) +def test_a_complete_body_cut_is_a_synthesized_server_error( + monkeypatch: pytest.MonkeyPatch, + aiohttp: bool, + cut: tuple[int, int | None, bool], + text: str, +) -> None: + from litellm.rust_bridge.host import http_transport + + monkeypatch.setattr(http_transport, "uses_aiohttp_transport", lambda: aiohttp) + native: Final = RuntimeError("upstream network error") + native.messages_body_cut = cut # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure(native, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert type(mapped) is litellm.InternalServerError + assert mapped.status_code == 500 + assert text in str(mapped) + + _TRANSPORT: Final = "error sending request: client error" @@ -445,3 +597,1153 @@ def test_a_failure_before_any_response_maps_like_the_python_transport_error( f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." ) assert "This looks like a bug in LiteLLM" not in str(mapped) + + +_CALL: Final = MappingProxyType({"messages": [{"role": "user", "content": "hi"}], "max_tokens": 8}) +_OPENCODE: Final = "https://opencode.ai/inference/go/anthropic/v1" +_COST_HEADER: Final = MappingProxyType({"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}) + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ( + "anthropic/claude-opus-latest", + None, + f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}", + ), + ("claude-opus-latest", "anthropic", get_model_latest_aliases()["anthropic"]["claude-opus-latest"]), + ("anthropic/claude-sonnet-4-5", None, "anthropic/claude-sonnet-4-5"), + ], +) +def test_settings_resolve_a_latest_alias_keeping_the_provider_prefix( + model: str, custom_llm_provider: str | None, expected: str +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": model, "custom_llm_provider": custom_llm_provider}) + + assert settings["model"] == expected + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + ({"api_base": _OPENCODE, "session_id": "s-1"}, {"x-opencode-session": "s-1"}), + ( + {"api_base": _OPENCODE, "session_id": "s-1", "headers": {"x-trace": "t"}}, + {"x-trace": "t", "x-opencode-session": "s-1"}, + ), + ({"api_base": _OPENCODE, "session_id": "s-1", "extra_headers": {"x-opencode-session": "mine"}}, None), + ({"api_base": "https://other.example/v1", "session_id": "s-1"}, None), + ({"api_base": _OPENCODE, "session_id": "s-1", "mock_response": "hi"}, None), + ], +) +def test_settings_add_only_the_affinity_header_the_caller_did_not_send( + extra: dict[str, object], expected: dict[str, str] | None +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + assert settings["extra_headers"] == expected + + +@pytest.mark.parametrize("variable", ["ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +def test_settings_add_the_affinity_header_for_an_anthropic_base_from_the_environment( + monkeypatch: pytest.MonkeyPatch, variable: str +) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + monkeypatch.setenv(variable, _OPENCODE) + + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", "session_id": "s-1"}) + + assert settings["extra_headers"] == {"x-opencode-session": "s-1"} + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + ({"stream": True, "proxy_server_request": dict(_COST_HEADER)}, "claude-sonnet-4-5"), + ({"proxy_server_request": dict(_COST_HEADER)}, None), + ({"stream": True}, None), + ({"stream": True, "proxy_server_request": dict(_COST_HEADER), "mock_response": "hi"}, None), + ], +) +def test_settings_name_the_cost_model_only_for_an_opted_in_stream( + extra: dict[str, object], expected: str | None +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + assert settings["stream_cost_model"] == expected + + +def test_settings_carry_the_mock_the_python_handler_returns() -> None: + settings: Final = messages_host.settings( + {}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", "mock_response": "hi"} + ) + + assert isinstance(settings["mock_response"], dict) + assert settings["mock_response"]["content"] == [{"type": "text", "text": "hi"}] + + +def test_settings_infer_no_provider_off_an_affinity_host(monkeypatch: pytest.MonkeyPatch) -> None: + def unexpected(model: str, api_base: str | None) -> str | None: + pytest.fail("the provider is only inferred for an affinity api_base") + + monkeypatch.setattr(messages_host, "inferred_provider", unexpected) + + settings: Final = messages_host.settings({}, {**_CALL, "model": "claude-sonnet-4-5"}) + + assert settings["model"] == "claude-sonnet-4-5" + assert settings["extra_headers"] is None + + +def test_body_params_leave_the_messages_to_the_route_for_a_call_without_prompt_caching( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "system": "Be terse.", + "enable_prompt_caching": False, + "cache_control_injection_points": [], + }, + asynchronous=True, + ) + + assert fields == {"system": "Be terse."} + + +@pytest.mark.parametrize("asynchronous", [True, False]) +def test_body_params_mark_the_turns_tools_and_deployment_the_handler_marks( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + gateway: Final[dict[str, object]] = {} + caller_metadata: Final = {"user_id": "u"} + tool: Final = {"name": "get_weather", "input_schema": {"type": "object"}} + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "again"}, + ], + "system": "Be terse.", + "tools": [tool], + "metadata": caller_metadata, + "litellm_metadata": gateway, + "model_info": {"id": "deployment-1"}, + "cache_control_injection_points": [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1, "control": {"type": "ephemeral", "ttl": "1h"}}, + {"location": "tool_config"}, + ], + }, + asynchronous=asynchronous, + ) + + assert fields == { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "user", + "content": [{"type": "text", "text": "again", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + }, + ], + "system": [{"type": "text", "text": "Be terse.", "cache_control": {"type": "ephemeral"}}], + "tools": [{**tool, "cache_control": {"type": "ephemeral"}}], + "metadata": caller_metadata, + } + assert gateway == {"litellm_gateway_injected_cache": "deployment-1"} + assert caller_metadata == {"user_id": "u"} + + +@pytest.mark.parametrize( + ("asynchronous", "expected_system"), + [ + (True, [{"type": "text", "text": "Be terse.", "cache_control": {"type": "ephemeral"}}]), + (False, "Be terse."), + ], +) +def test_body_params_prepare_prompt_caching_as_the_entry_that_ran( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, expected_system: object +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "system": "Be terse.", + "cache_control_injection_points": [{"location": "message", "role": "assistant"}], + }, + asynchronous=asynchronous, + ) + + assert fields["system"] == expected_system + + +@pytest.mark.parametrize( + ("modify_params", "args", "kwargs", "expected"), + [ + pytest.param(True, (), {"max_tokens": 200_000}, 64_000, id="clamped-to-model-output-limit"), + pytest.param(True, (), {"max_tokens": 1}, 1, id="under-the-limit"), + pytest.param(False, (), {"max_tokens": 200_000}, 200_000, id="modify-params-off"), + pytest.param("true", (), {"max_tokens": 200_000}, 200_000, id="truthy-but-not-true"), + pytest.param(True, (200_000,), {"max_tokens": 200_000}, 200_000, id="positional-model-is-not-a-model"), + ], +) +def test_modify_max_tokens_clamps_the_keyword_view_as_the_client_wrapper_does( + monkeypatch: pytest.MonkeyPatch, + modify_params: object, + args: tuple[object, ...], + kwargs: dict[str, object], + expected: int, +) -> None: + monkeypatch.setattr(litellm, "modify_params", modify_params) + model: Final = args[0] if args else "anthropic/claude-sonnet-4-5" + view: Final = {**kwargs, "messages": [{"role": "user", "content": "hi"}]} + + messages_host.modify_max_tokens(model, args, view) + + assert view["max_tokens"] == expected + + +def test_workload_identity_token_mints_with_the_params_python_validates(monkeypatch: pytest.MonkeyPatch) -> None: + minted: Final[list[tuple[dict[str, object], str | None, str]]] = [] + + def mint(litellm_params: dict[str, object], api_base: str | None, model: str) -> str: + federation = {name: value for name, value in litellm_params.items() if name.startswith("anthropic_")} + minted.append(({name: value for name, value in federation.items() if value is not None}, api_base, model)) + return "sk-ant-oat01-minted" + + monkeypatch.setattr(messages_host, "get_anthropic_wif_token", mint) + + token: Final = messages_host.workload_identity_token( + {"anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": "60"}, + "https://gateway.test", + ) + + assert (token, minted) == ( + "sk-ant-oat01-minted", + [({"anthropic_federation_rule_id": "fdrl_1", "anthropic_issuer_ttl_seconds": 60}, "https://gateway.test", "")], + ) + + +def test_workload_identity_token_is_none_without_a_federation_rule(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("ANTHROPIC_FEDERATION_RULE_ID", "ANTHROPIC_ORGANIZATION_ID", "ANTHROPIC_IDENTITY_TOKEN"): + monkeypatch.delenv(name, raising=False) + + assert messages_host.workload_identity_token({}, None) is None + + +def test_workload_identity_fields_are_the_federation_litellm_params() -> None: + assert set(messages_host.WORKLOAD_IDENTITY_FIELDS) >= { + "anthropic_federation_rule_id", + "anthropic_organization_id", + "anthropic_disable_workload_identity_federation", + } + + +class _Handler(AsyncHTTPHandler): + def __init__(self, answer: httpx.Response | Exception) -> None: + super().__init__() # pyright: ignore[reportUnknownMemberType] # the handler's untyped concurrent_limit + self.answer: Final = answer + self.posted: list[Mapping[str, object]] = [] + + async def post(self, url: str, **kwargs: object) -> httpx.Response: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: records the post's keywords + self.posted.append({"url": url, **kwargs}) + if isinstance(self.answer, Exception): + raise self.answer + return self.answer + + +_URL: Final = "https://api.anthropic.com/v1/messages" +_REQUEST: Final = httpx.Request("POST", _URL) + + +@pytest.fixture +def no_tls_settings(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + + +@pytest.mark.usefixtures("no_tls_settings") +@pytest.mark.parametrize( + "client", (None, object(), httpx.AsyncClient(), HTTPHandler()), ids=("none", "object", "httpx", "sync-handler") +) +def test_the_native_client_stands_in_for_any_client_python_ignores(client: object) -> None: + assert messages_host.transport_handler({"client": client}) is None + + +@pytest.mark.usefixtures("no_tls_settings") +def test_a_caller_async_handler_is_the_handler_python_posts_through() -> None: + client: Final = AsyncHTTPHandler() + + assert messages_host.transport_handler({"client": client}) is client + + +@pytest.mark.usefixtures("no_tls_settings") +@pytest.mark.parametrize( + "setting", ({"ssl_security_level": "DEFAULT@SECLEVEL=1"}, {"ssl_ecdh_curve": "secp521r1"}), ids=("level", "curve") +) +def test_tls_the_native_client_cannot_apply_posts_through_pythons_anthropic_handler( + monkeypatch: pytest.MonkeyPatch, setting: Mapping[str, str] +) -> None: + for name, value in setting.items(): + monkeypatch.setattr(litellm, name, value) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + + handler: Final = messages_host.transport_handler({"client": object()}) + + assert handler is get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) + assert callbacks_legacy_python.times_provider_exchange(object(), None) is False + + +@pytest.fixture +def fips_cipher_exclusion() -> Iterator[None]: + exclude_rsa_key_exchange_and_chacha20(True) + yield + exclude_rsa_key_exchange_and_chacha20(False) + + +@pytest.mark.usefixtures("no_tls_settings", "fips_cipher_exclusion") +def test_fips_cipher_exclusion_posts_through_pythons_anthropic_handler(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + + handler: Final = messages_host.transport_handler({"client": None}) + + assert handler is get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) + assert callbacks_legacy_python.times_provider_exchange(None, None) is False + + +@pytest.mark.usefixtures("no_tls_settings") +def test_only_a_native_post_is_timed_by_the_native_call() -> None: + assert callbacks_legacy_python.times_provider_exchange(None, None) is True + assert callbacks_legacy_python.times_provider_exchange(None, 30.0) is True + assert callbacks_legacy_python.times_provider_exchange(AsyncHTTPHandler(), None) is False + + +@pytest.mark.usefixtures("no_tls_settings") +@pytest.mark.parametrize( + "timeout", + (-1.0, 0, -0.5, httpx.Timeout(10.0, connect=0)), + ids=("negative", "zero", "fraction", "httpx-zero-connect"), +) +def test_a_non_positive_timeout_posts_through_pythons_anthropic_handler( + monkeypatch: pytest.MonkeyPatch, timeout: object +) -> None: + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + + handler: Final = messages_host.transport_handler({"client": None, "timeout": timeout}) + + assert handler is get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC) + assert callbacks_legacy_python.times_provider_exchange(None, timeout) is False + + +@pytest.mark.usefixtures("no_tls_settings") +@pytest.mark.parametrize("timeout", (None, 0.5, httpx.Timeout(10.0)), ids=("unset", "positive", "httpx")) +def test_a_positive_timeout_keeps_the_native_client(timeout: object) -> None: + assert messages_host.transport_handler({"client": None, "timeout": timeout}) is None + + +@pytest.mark.asyncio +async def test_transport_send_posts_what_the_python_handler_posts() -> None: + logger: Final = object() + timeout: Final = httpx.Timeout(30.0, connect=2.0) + handler: Final = _Handler( + httpx.Response(200, json={"id": "msg_1"}, headers={"request-id": "req_1"}, request=_REQUEST) + ) + + sent: Final = await messages_host.transport_send( + handler, + {"litellm_logging_obj": logger, "timeout": timeout}, + _URL, + [("x-api-key", "sk-ant"), ("anthropic-version", "2023-06-01")], + {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "héllo"}]}, + False, + ) + + assert handler.posted == [ + { + "url": _URL, + "headers": {"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, + "data": json.dumps({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "héllo"}]}), + "stream": False, + "logging_obj": logger, + "timeout": timeout, + } + ] + assert (sent.status, sent.text, sent.chunks, sent.rejection) == (200, '{"id":"msg_1"}', None, None) + assert ("request-id", "req_1") in sent.headers + + +@pytest.mark.asyncio +async def test_transport_send_hands_signed_bytes_on_unchanged() -> None: + handler: Final = _Handler(httpx.Response(200, json={}, request=_REQUEST)) + + await messages_host.transport_send(handler, {}, _URL, [], b'{"signed":true}', False) + + assert handler.posted[0]["data"] == b'{"signed":true}' + assert handler.posted[0]["timeout"] is None + + +@pytest.mark.asyncio +async def test_transport_send_leaves_a_stream_for_the_route_to_read() -> None: + handler: Final = _Handler( + httpx.Response(200, stream=httpx.ByteStream(b"event: message_stop\n\n"), request=_REQUEST) + ) + + sent: Final = await messages_host.transport_send(handler, {}, _URL, [], {}, True) + + assert sent.chunks is not None + assert sent.text is None + assert [chunk async for chunk in sent.chunks] == [b"event: message_stop\n\n"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("checked", (True, False), ids=("handler-raised", "status-checked")) +async def test_transport_send_reports_a_rejection_with_the_status_error(checked: bool) -> None: + rejected: Final = httpx.Response(429, text="slow down", headers={"retry-after": "7"}, request=_REQUEST) + error: Final = httpx.HTTPStatusError("rate limited", request=_REQUEST, response=rejected) + handler: Final = _Handler(error if checked else rejected) + + sent: Final = await messages_host.transport_send(handler, {}, _URL, [], {}, False) + + assert (sent.status, sent.text) == (429, "slow down") + assert ("retry-after", "7") in sent.headers + assert isinstance(sent.rejection, httpx.HTTPStatusError) + assert sent.rejection.response is rejected + + +@pytest.mark.asyncio +async def test_a_rejection_python_cannot_read_reaches_the_caller_as_raised() -> None: + unread: Final = httpx.Response(500, stream=httpx.ByteStream(b"unread"), request=_REQUEST) + + with pytest.raises(messages_host.RawTransportFailure) as raised: + await messages_host.transport_send(_Handler(unread), {}, _URL, [], {}, True) + + assert isinstance(raised.value.error, httpx.ResponseNotRead) + assert isinstance(raised.value.error.__context__, httpx.HTTPStatusError) + + +@pytest.mark.asyncio +async def test_a_handler_failure_is_raised_for_the_route_to_map() -> None: + failure: Final = httpx.ConnectError("refused") + + with pytest.raises(httpx.ConnectError) as raised: + await messages_host.transport_send(_Handler(failure), {}, _URL, [], {}, False) + + assert raised.value is failure + + +@pytest.mark.parametrize( + ("error", "expected", "status"), + ( + pytest.param( + litellm.Timeout( + message="Read timed out.", model="default-model-name", llm_provider="litellm-httpx-handler" + ), + litellm.Timeout, + 408, + id="timeout", + ), + pytest.param( + httpx.ConnectError("All connection attempts failed"), litellm.InternalServerError, 500, id="connect" + ), + pytest.param( + httpx.HTTPStatusError( + "unavailable", + request=_REQUEST, + response=httpx.Response(503, text="Unavailable", request=_REQUEST), + ), + litellm.ServiceUnavailableError, + 503, + id="status", + ), + ), +) +def test_a_handler_failure_maps_as_the_python_handler_maps_it( + error: Exception, expected: type[Exception], status: int +) -> None: + mapped: Final = messages_host.map_transport_failure( + error, {"model": "anthropic/claude-sonnet-4-5", "api_base": None}, "anthropic" + ) + + assert isinstance(mapped, expected) + assert getattr(mapped, "status_code", None) == status + assert mapped.__context__ is not None + assert mapped.__context__.__context__ is error + + +class _Hook(CustomLogger): + rewrite: Mapping[str, object] = MappingProxyType({}) + seen: tuple[tuple[str, list[dict[str, object]], dict[str, object]], ...] = () + + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + self.seen = (*self.seen, (model, list(messages), dict(kwargs))) + return {**kwargs, **self.rewrite} + + +def _hook(rewrite: Mapping[str, object]) -> _Hook: + hook: Final = _Hook() + hook.rewrite = rewrite + return hook + + +def _hook_request(**overrides: object) -> dict[str, object]: + return { + "max_tokens": 64, + "messages": [{"role": "user", "content": [{"type": "text", "text": ""}, {"type": "text", "text": "hi"}]}], + "model": "anthropic/claude-sonnet-4-5", + "metadata": None, + "stop_sequences": None, + "stream": False, + "system": None, + "temperature": None, + "thinking": None, + "tool_choice": None, + "tools": None, + "top_k": None, + "top_p": None, + "api_key": "sk-test", + "api_base": None, + "client": None, + "custom_llm_provider": None, + "litellm_call_id": "call-1", + **overrides, + } + + +def _instance_hook() -> CustomLogger: + logger: Final = CustomLogger() + logger.async_pre_request_hook = _hook({}).async_pre_request_hook # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: an instance-level override + return logger + + +@pytest.mark.parametrize( + ("callbacks", "expected"), + [ + pytest.param([_hook({})], True, id="override"), + pytest.param([CustomLogger()], False, id="base-logger"), + pytest.param([_instance_hook()], True, id="instance-override"), + pytest.param(["not-a-logger"], False, id="not-a-logger"), + ], +) +def test_runs_pre_request_hooks_needs_a_global_override( + monkeypatch: pytest.MonkeyPatch, callbacks: list[object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", callbacks) + + assert messages_host.runs_pre_request_hooks() is expected + + +@pytest.mark.asyncio +async def test_pre_request_hooks_hand_the_handler_what_the_python_entry_hands_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hook: Final = _hook({"temperature": 0.2, "tools": [{"name": "t"}], "extra_headers": {"x-hooked": "1"}}) + monkeypatch.setattr(litellm, "callbacks", [hook]) + + hooked: Final = await messages_host.pre_request_hooks(_hook_request(temperature=0.9)) + + assert hook.seen == ( + ( + "anthropic/claude-sonnet-4-5", + [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + { + "tools": None, + "stream": False, + "litellm_params": {"custom_llm_provider": "anthropic"}, + "tool_choice": None, + "litellm_call_id": "call-1", + }, + ), + ) + assert hooked["temperature"] == 0.2 + assert hooked["tools"] == [{"name": "t"}] + assert hooked["extra_headers"] == {"x-hooked": "1"} + assert hooked["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + assert hooked["custom_llm_provider"] is None + assert "litellm_params" not in hooked + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", ["model", "max_tokens", "messages", "api_key", "api_base", "custom_llm_provider"]) +async def test_pre_request_hooks_naming_a_handler_argument_raise_the_python_type_error( + monkeypatch: pytest.MonkeyPatch, name: str +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_hook({name: "hooked"})]) + + with pytest.raises(TypeError, match=f"functools.partial\\(\\) got multiple values for keyword argument '{name}'"): + await messages_host.pre_request_hooks(_hook_request()) + + +@pytest.mark.asyncio +async def test_pre_request_hooks_refuse_a_route_to_another_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", [_hook({"litellm_params": {"custom_llm_provider": "openai"}})]) + + with pytest.raises(litellm.UnsupportedParamsError, match="routing the call to openai") as refused: + await messages_host.pre_request_hooks(_hook_request()) + + assert refused.value.llm_provider == "anthropic" + + +@pytest.mark.asyncio +async def test_pre_request_hooks_keep_the_callers_provider_over_the_hook(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", [_hook({"litellm_params": {"custom_llm_provider": "openai"}})]) + + hooked: Final = await messages_host.pre_request_hooks(_hook_request(custom_llm_provider="anthropic")) + + assert hooked["custom_llm_provider"] == "anthropic" + + +def test_body_params_hand_hooked_messages_on_as_already_sanitized() -> None: + hooked: Final = [{"role": "user", "content": [{"type": "text", "text": ""}]}] + + fields: Final = messages_host.body_params( + {"model": "anthropic/claude-sonnet-4-5", "messages": hooked}, asynchronous=True, hooked=True + ) + + assert fields == {"messages": hooked} + + +class _Agentic(CustomLogger): + async def async_should_run_agentic_loop(self, *args: object, **kwargs: object) -> tuple[bool, dict[str, object]]: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return False, {} + + +def _agentic_logger(callbacks: Sequence[CustomLogger]) -> Logging: + logger: Final = Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="agentic", + function_id="agentic", + dynamic_success_callbacks=list(callbacks), + ) + return logger + + +@pytest.mark.parametrize( + ("callbacks", "dynamic", "request_extra", "expected"), + [ + pytest.param([], [], {}, False, id="no-hook"), + pytest.param([_Agentic()], [], {}, True, id="global-hook"), + pytest.param([], [_Agentic()], {}, True, id="call-hook"), + pytest.param([], [], {"_websearch_interception_converted_stream": True}, True, id="converted-stream"), + ], +) +def test_runs_agentic_hooks_gates_as_the_python_handler_does( + monkeypatch: pytest.MonkeyPatch, + callbacks: list[object], + dynamic: Sequence[CustomLogger], + request_extra: Mapping[str, object], + expected: bool, +) -> None: + monkeypatch.setattr(litellm, "callbacks", callbacks) + request: Final = {"model": "anthropic/claude-sonnet-4-5", "litellm_logging_obj": _agentic_logger(dynamic)} + + assert messages_host.runs_agentic_hooks({**request, **request_extra}) is expected + assert messages_host.runs_agentic_hooks({"model": "anthropic/claude-sonnet-4-5"}) is False + + +def test_record_agentic_loop_names_the_requested_model_for_follow_up_calls() -> None: + logger: Final = _agentic_logger([]) + messages_host.record_agentic_loop( + { + "model": "anthropic/claude-sonnet-4-5", + "litellm_logging_obj": logger, + "_websearch_interception_converted_stream": True, + }, + "anthropic", + ) + + details: Final = logger.model_call_details + assert details["agentic_loop_params"] == { + "model": "anthropic/claude-sonnet-4-5", + "custom_llm_provider": "anthropic", + } + assert details["websearch_interception_converted_stream"] is True + + +def test_agentic_stream_tags_the_head_only_when_a_hook_runs(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge.messages.entrypoints import MessagesStreamHead + + request: Final = { + "model": "anthropic/claude-sonnet-4-5", + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi", "provider_specific_fields": {"x": 1}}]} + ], + "tools": [{"name": "get_weather", "input_schema": {"type": "object"}}], + "api_key": "sk-test", + "litellm_call_id": "agentic", + "litellm_logging_obj": _agentic_logger([]), + } + sent: Final = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "stream": True, + "tools": [{"name": "get_weather", "input_schema": {"type": "object"}}], + } + requested: Final = {"max_tokens": 16, "stream": True, "tools": request["tools"]} + head: Final = MessagesStreamHead() + + monkeypatch.setattr(litellm, "callbacks", []) + messages_host.agentic_stream(head, request, sent, "claude-sonnet-4-5", "anthropic", requested, hooked=False) + assert head.agentic is None + + monkeypatch.setattr(litellm, "callbacks", [_Agentic()]) + messages_host.agentic_stream(head, request, sent, "claude-sonnet-4-5", "anthropic", requested, hooked=False) + agentic: Final = head.agentic + assert isinstance(agentic, functools.partial) + bound: Final[Mapping[str, object]] = agentic.keywords + call: Final = bound["call"] + assert isinstance(call, messages_host._AgenticCall) # pyright: ignore[reportPrivateUsage] # the partial binds the handler call + assert call.messages == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + assert call.params == {"stream": True, "tools": request["tools"]} + assert call.params["tools"] is request["tools"] + assert call.kwargs["api_key"] == "sk-test" + assert "tools" not in call.kwargs + + +_DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use"}] +_CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "openai/classifier"} +_NO_TOOL_USES: Final[dict[str, object]] = {} +_AVAILABLE: Final = ({"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": _NO_TOOL_USES}},) + + +@pytest.mark.parametrize( + ("request_fields", "general_settings", "classifies"), + ( + pytest.param( + {"model": "anthropic/parity-tool-model", "safeguards": _DANGEROUS_TOOL_USE}, + _CLASSIFIER_SETTINGS, + True, + id="non-claude", + ), + pytest.param( + {"model": "anthropic/claude-sonnet-4-5", "safeguards": _DANGEROUS_TOOL_USE}, + _CLASSIFIER_SETTINGS, + False, + id="claude", + ), + pytest.param( + {"model": "anthropic/parity-tool-model", "safeguards": _DANGEROUS_TOOL_USE}, {}, False, id="unconfigured" + ), + pytest.param({"model": "anthropic/parity-tool-model"}, _CLASSIFIER_SETTINGS, False, id="not-requested"), + pytest.param( + {"model": "anthropic/parity-tool-model", "safeguards": [{"type": "other"}]}, + _CLASSIFIER_SETTINGS, + False, + id="other-safeguard", + ), + ), +) +def test_safeguards_classify_only_what_the_proxy_answers( + monkeypatch: pytest.MonkeyPatch, + request_fields: Mapping[str, object], + general_settings: Mapping[str, object], + classifies: bool, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", general_settings) + + assert (messages_host.safeguards(request_fields) is not None) is classifies + + +@pytest.mark.asyncio +async def test_safeguards_frames_split_and_stamp_like_the_python_stream(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.messages.safeguards import SafeguardsEvaluator + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", _CLASSIFIER_SETTINGS) + evaluated: Final[list[object]] = [] + + async def evaluate(_evaluator: SafeguardsEvaluator, tool_uses: object) -> object: + evaluated.append(tool_uses) + return _AVAILABLE + + monkeypatch.setattr(SafeguardsEvaluator, "evaluate", evaluate) + classifier: Final = messages_host.safeguards( + { + "model": "anthropic/parity-tool-model", + "safeguards": _DANGEROUS_TOOL_USE, + "messages": [{"role": "user", "content": "hi"}], + } + ) + assert classifier is not None + assert await classifier.open_stream() + + opened: Final = await classifier.frames( + b'event: ping\ndata: {"type": "ping"}\n\nevent: message_delta\ndata: {"type": ' + ) + closed: Final = await classifier.frames(b'"message_delta", "delta": {"stop_reason": "end_turn"}}\n\ntail') + flushed: Final = await classifier.flush() + + assert opened == (b'event: ping\ndata: {"type": "ping"}\n\n',) + assert closed == ( + b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn", ' + b'"safeguard_results": [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}]}}\n\n', + ) + assert flushed == (b"tail",) + assert evaluated == [()] + + +@pytest.mark.asyncio +async def test_safeguards_publish_the_results_python_stamped(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.messages.safeguards import SafeguardsEvaluator + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", _CLASSIFIER_SETTINGS) + stamped: Final = ({"type": "dangerous_tool_use", "status": {"type": "unsupported"}},) + + async def evaluate(_evaluator: SafeguardsEvaluator, tool_uses: object) -> object: + return stamped + + monkeypatch.setattr(SafeguardsEvaluator, "evaluate", evaluate) + request: Final = { + "model": "anthropic/parity-tool-model", + "safeguards": _DANGEROUS_TOOL_USE, + "messages": [{"role": "user", "content": "hi"}], + } + unanswered: Final = messages_host.safeguards(request) + answered: Final = messages_host.safeguards(request) + assert unanswered is not None and answered is not None + + assert await unanswered.results({"content": [], "safeguard_results": []}) is None + assert unanswered.published({"id": "msg"}) == {"id": "msg"} + assert await answered.results({"content": []}) is stamped + assert answered.published({"id": "msg", "safeguard_results": [stamped[0]]}) == { + "id": "msg", + "safeguard_results": stamped, + } + + +class _Claims(MessagesInterceptor): + def __init__(self, provider: str) -> None: + self.provider = provider + self.asked: list[tuple[object, str | None]] = [] + + def can_handle(self, tools: Sequence[Mapping[str, object]] | None, custom_llm_provider: str | None) -> bool: + self.asked.append((tools, custom_llm_provider)) + return custom_llm_provider == self.provider + + async def handle(self, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: the interceptor contract's keyword call + return kwargs + + +_TOOLS: Final = [{"type": "custom_tool", "name": "t"}] + + +@pytest.mark.parametrize( + ("request_fields", "provider"), + ( + pytest.param({"model": "anthropic/claude-sonnet-4-5"}, "anthropic", id="prefixed"), + pytest.param({"model": "claude-sonnet-4-5", "custom_llm_provider": "anthropic"}, "anthropic", id="explicit"), + pytest.param({"model": "not-a-known-model"}, None, id="unresolved"), + ), +) +def test_interceptor_asks_with_the_tools_and_the_provider_the_handler_resolves( + monkeypatch: pytest.MonkeyPatch, request_fields: Mapping[str, object], provider: str | None +) -> None: + declined: Final = _Claims("other") + claims: Final = _Claims("anthropic") + monkeypatch.setattr(interceptors, "_interceptors", [declined, claims]) + + chosen: Final = messages_host.interceptor({**request_fields, "tools": _TOOLS}) + + assert chosen is (claims if provider == "anthropic" else None) + assert declined.asked == [(_TOOLS, provider)] + + +@pytest.mark.asyncio +async def test_answer_hands_the_interceptor_the_call_as_the_handler_does() -> None: + owner: Final = _Claims("anthropic") + request: Final = { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}, {"type": "text", "text": ""}]}], + "max_tokens": 16, + "tools": _TOOLS, + "stream": None, + "system": None, + "client": object(), + "custom_llm_provider": None, + "litellm_call_id": "call-1", + } + + handed: Final = await messages_host.answer(owner, request, stream=False, hooked=False) + + assert isinstance(handed, dict) + assert handed["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + assert handed["custom_llm_provider"] == "anthropic" + assert handed["stream"] is False + assert handed["litellm_call_id"] == "call-1" + assert "client" not in handed + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "converted", "expected"), + ( + pytest.param(True, False, True, id="caller-stream"), + pytest.param(False, True, True, id="converted-stream"), + pytest.param(False, False, False, id="unary"), + ), +) +async def test_answer_hands_a_hooked_call_on_with_the_callers_stream( + stream: bool, converted: bool, expected: bool +) -> None: + hooked_messages: Final = [{"role": "user", "content": [{"type": "text", "text": ""}]}] + request: Final = { + "model": "anthropic/claude-sonnet-4-5", + "messages": hooked_messages, + "max_tokens": 16, + "tools": _TOOLS, + "stream": False, + "system": "hooked", + "custom_llm_provider": None, + **({"_websearch_interception_converted_stream": True} if converted else {}), + } + + handed: Final = await messages_host.answer(_Claims("anthropic"), request, stream=stream, hooked=True) + + assert isinstance(handed, dict) + assert handed["messages"] is hooked_messages + assert handed["system"] == "hooked" + assert handed["stream"] is expected + + +MCP_REFERENCE: Final = MappingProxyType( + {"type": "mcp", "server_label": "litellm", "server_url": "litellm_proxy/mcp/deepwiki", "require_approval": "never"} +) + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + pytest.param({"tools": [dict(MCP_REFERENCE)]}, True, id="gateway-reference"), + pytest.param({"tools": [{**MCP_REFERENCE, "server_url": "litellm_proxy"}]}, True, id="whole-gateway"), + pytest.param({"tools": [dict(MCP_REFERENCE)], "_skip_mcp_handler": True}, False, id="gateway-follow-up"), + pytest.param({"tools": [{"name": "a", "input_schema": {"type": "object"}}]}, False, id="native-tool"), + pytest.param({"tools": []}, False, id="no-tools"), + pytest.param({"tools": [dict(MCP_REFERENCE)], "mock_response": "hi"}, False, id="mock-answers-first"), + ], +) +def test_settings_hand_the_call_to_the_mcp_gateway_where_the_python_handler_does( + extra: dict[str, object], expected: bool +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + assert settings["mcp_gateway"] is expected + + +def test_gateway_is_none_for_a_call_the_handler_serves_itself() -> None: + assert messages_host.gateway({**_CALL, "model": "anthropic/claude-sonnet-4-5"}, True) is None + + +def _gateway_request() -> dict[str, object]: + return { + **_CALL, + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": [{"type": "text", "text": ""}, {"type": "text", "text": "hi"}]}], + "tools": [dict(MCP_REFERENCE)], + "metadata": {"user_id": "u-1", "session": "dropped"}, + "stream": True, + "temperature": 0.2, + "litellm_metadata": {"tags": ["t1"]}, + } + + +def test_sync_gateway_hands_the_mcp_run_what_the_python_handler_hands_it(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.messages.native import handler + + sentinel: Final = object() + handed: Final[list[dict[str, object]]] = [] + + def run(**kwargs: object) -> object: + handed.append(kwargs) + return sentinel + + monkeypatch.setattr(mcp_handler, "anthropic_messages_with_mcp", run) + request: Final = _gateway_request() + + assert handler.anthropic_messages_handler(**_gateway_request()) is sentinel # noqa: TID251 # the Python handler under comparison + assert messages_host.gateway(request, False) is sentinel + assert handed[1] == handed[0] + assert handed[1]["model"] == "claude-sonnet-4-5" + assert handed[1]["custom_llm_provider"] == "anthropic" + assert handed[1]["metadata"] == {"user_id": "u-1"} + assert handed[1]["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + assert "_skip_mcp_handler" not in handed[1] + + +def test_async_gateway_tells_the_logger_the_agentic_loop_params_as_the_handler_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + sentinel: Final = object() + handed: Final[list[dict[str, object]]] = [] + + def run(**kwargs: object) -> object: + handed.append(kwargs) + return sentinel + + monkeypatch.setattr(mcp_handler, "anthropic_messages_with_mcp", run) + logger: Final = Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=False, + call_type="aanthropic_messages", + start_time=None, + litellm_call_id="call-1", + function_id="fn", + ) + + assert messages_host.gateway({**_gateway_request(), "litellm_logging_obj": logger}, True) is sentinel + assert logger.model_call_details["agentic_loop_params"] == { + "model": "anthropic/claude-sonnet-4-5", + "custom_llm_provider": "anthropic", + } + assert logger.model_call_details["litellm_params"]["aanthropic_messages"] is True + assert "is_async" not in handed[0] + assert handed[0]["litellm_logging_obj"] is logger + + +def test_async_gateway_hands_hooked_messages_on_as_already_sanitized(monkeypatch: pytest.MonkeyPatch) -> None: + handed: Final[list[dict[str, object]]] = [] + + def run(**kwargs: object) -> object: + handed.append(kwargs) + return None + + monkeypatch.setattr(mcp_handler, "anthropic_messages_with_mcp", run) + hooked: Final = [ + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]} + ] + + messages_host.gateway({**_gateway_request(), "messages": hooked}, True, hooked=True) + + assert handed[0]["messages"] == hooked + + +def test_sync_answer_serves_the_gateway_after_a_mock(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel: Final = object() + + def run(**_: object) -> object: + return sentinel + + monkeypatch.setattr(mcp_handler, "anthropic_messages_with_mcp", run) + + served: Final = messages_host.sync_answer(_gateway_request()) + mocked: Final = messages_host.sync_answer({**_gateway_request(), "mock_response": "hi"}) + + assert served is not None and served.value is sentinel + assert mocked is not None and isinstance(mocked.value, dict) + assert messages_host.sync_answer({**_CALL, "model": "anthropic/claude-sonnet-4-5"}) is None + + +@pytest.mark.parametrize( + ("extra", "raised", "text"), + [ + pytest.param({"messages": None}, TypeError, "'NoneType' object is not iterable", id="messages-none"), + pytest.param({"messages": 5}, TypeError, "'int' object is not iterable", id="messages-int"), + pytest.param({"tools": {"name": "t"}}, AttributeError, "'str' object has no attribute 'get'", id="tools-map"), + pytest.param({"model_info": "not-a-mapping"}, ValidationError, "model_info", id="model-info-str"), + pytest.param({"custom_llm_provider": 5}, ValidationError, "custom_llm_provider", id="provider-int"), + ], +) +def test_settings_raise_what_the_python_handler_raises_for_arguments_of_the_wrong_type( + extra: dict[str, object], raised: type[Exception], text: str +) -> None: + with pytest.raises(raised, match=text): + messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + +def test_settings_leave_well_typed_tools_to_the_interceptors_after_the_cache(monkeypatch: pytest.MonkeyPatch) -> None: + asked: Final[list[object]] = [] + + def can_handle(tools: object, custom_llm_provider: object) -> bool: + asked.append(tools) + return False + + advisor: Final = interceptors.get_messages_interceptors()[0] + monkeypatch.setattr(advisor, "can_handle", can_handle) + + messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", "tools": [{"name": "t"}]}) + + assert asked == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("hook_params", "raised", "text"), + [ + pytest.param("not-a-mapping", AttributeError, "'str' object has no attribute 'get'", id="params-str"), + pytest.param({"custom_llm_provider": 5}, ValidationError, "custom_llm_provider", id="provider-int"), + ], +) +async def test_pre_request_hooks_raise_what_the_python_handler_raises_for_the_hooks_params( + monkeypatch: pytest.MonkeyPatch, hook_params: object, raised: type[Exception], text: str +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_hook({"litellm_params": hook_params})]) + + with pytest.raises(raised, match=text): + await messages_host.pre_request_hooks(_hook_request(model="claude-sonnet-4-5")) + + +def test_a_provider_config_rejection_maps_as_the_python_http_handler_maps_it() -> None: + logger: Final = Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="rejected", + function_id="rejected", + ) + rejected: Final = ValueError("missing required field: max_tokens") + rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + rejected.messages_provider_rejection = "max_tokens is required for Anthropic /v1/messages API" # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure( + rejected, + {"model": "anthropic/claude-sonnet-4-5", "messages": [], "litellm_logging_obj": logger}, + "anthropic", + ) + + assert isinstance(mapped, litellm.BadRequestError) + assert str(mapped) == ( + "litellm.BadRequestError: AnthropicException - max_tokens is required for Anthropic /v1/messages API" + ) + assert mapped.model == "anthropic/claude-sonnet-4-5" # pyright: ignore[reportUnknownMemberType] # the public error leaves model untyped + assert mapped.llm_provider == "anthropic" # pyright: ignore[reportUnknownMemberType] # the public error leaves llm_provider untyped + assert logger.model_call_details["custom_llm_provider"] == "anthropic" + + +def test_short_circuits_needs_a_replaced_web_search_short_circuit(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger + + class _ShortCircuit(WebSearchInterceptionLogger): + async def try_short_circuit_search(self, *args: object, **kwargs: object) -> dict[str, object] | None: # pyright: ignore[reportIncompatibleMethodOverride] # test double only needs the override + return None + + monkeypatch.setattr(litellm, "callbacks", [CustomLogger()]) + assert not messages_host.short_circuits() + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger()]) + assert not messages_host.short_circuits() + monkeypatch.setattr(litellm, "callbacks", [_ShortCircuit()]) + assert messages_host.short_circuits() diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py index c2b87cca1dc..9b58524a444 100644 --- a/tests/unit/rust_bridge/host/test_cache.py +++ b/tests/unit/rust_bridge/host/test_cache.py @@ -1,12 +1,22 @@ import asyncio +import datetime +from collections.abc import AsyncIterator, Callable from typing import Final import pytest import litellm from litellm.caching.caching import Cache -from litellm.caching.caching_handler import _PENDING_CACHE_WRITES -from litellm.rust_bridge.host.cache import configured_cache, store_response +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES, LLMCachingHandler +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.messages.native.response_cache import AnthropicMessagesStreamCacheWriter +from litellm.rust_bridge.host.cache import ( + caching_handler, + configured_cache, + store_response, + store_verbatim, + wrap_stream, +) from litellm.types.caching import LiteLLMCacheType @@ -26,9 +36,9 @@ def test_configured_cache_preserves_identity_and_observes_runtime_replacement(mo _REQUEST: Final = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]} -async def _stored(response: dict[str, object]) -> object: +async def _stored(response: dict[str, object], store: Callable[..., None] = store_response) -> object: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) - store_response(cache, response, **_REQUEST) + store(cache, response, **_REQUEST) await asyncio.gather(*_PENDING_CACHE_WRITES) return await cache.async_get_cache(**_REQUEST) @@ -51,3 +61,104 @@ async def test_store_response_drops_only_the_rust_marker_header_from_the_cached_ async def test_store_response_keeps_an_entry_without_the_rust_marker_unchanged() -> None: response: Final = {"id": "msg_2", "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}}} assert await _stored(response) == response + + +async def _events() -> AsyncIterator[bytes]: + yield b"event: message_stop\ndata: {}\n\n" + + +async def anthropic_messages() -> None: + return None + + +def _logging() -> Logging: + return Logging( + model="claude-sonnet-4-5", + messages=[], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id="messages", + function_id="messages", + ) + + +def test_wrap_stream_hands_the_stream_to_the_cache_writer_the_client_wrapper_uses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL, supported_call_types=["anthropic_messages"]) + monkeypatch.setattr(litellm, "cache", cache) + stream: Final = _events() + + wrapped: Final = wrap_stream(cache, anthropic_messages, stream, None, **_REQUEST, litellm_logging_obj=_logging()) + + assert isinstance(wrapped, AnthropicMessagesStreamCacheWriter) + assert wrapped.stream is stream + + +def test_wrap_stream_leaves_a_stream_the_cache_would_not_store(monkeypatch: pytest.MonkeyPatch) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL, supported_call_types=["anthropic_messages"]) + monkeypatch.setattr(litellm, "cache", cache) + stream: Final = _events() + + assert ( + wrap_stream( + cache, + anthropic_messages, + stream, + None, + **_REQUEST, + cache={"no-store": True}, + litellm_logging_obj=_logging(), + ) + is stream + ) + + +@pytest.mark.asyncio +async def test_store_verbatim_keeps_the_rust_marker_header_in_the_cached_entry() -> None: + response: Final[dict[str, object]] = { + "id": "msg_3", + "_hidden_params": {"additional_headers": {"x-litellm-rust": "true", "llm_provider-x": "1"}}, + } + assert await _stored(response, store_verbatim) == response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store", [store_response, store_verbatim]) +@pytest.mark.parametrize("content", [[], [{"type": "text", "text": ""}]]) +async def test_a_message_without_output_is_never_stored( + store: Callable[..., None], content: list[dict[str, object]] +) -> None: + assert await _stored({"id": "msg_4", "type": "message", "content": content}, store) is None + + +def test_wrap_stream_writes_through_the_handler_that_read_the_cache_over_a_later_call_s( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL, supported_call_types=["anthropic_messages"]) + monkeypatch.setattr(litellm, "cache", cache) + logging_obj: Final = _logging() + read: Final = LLMCachingHandler( + original_function=anthropic_messages, request_kwargs=dict(_REQUEST), start_time=datetime.datetime.now() + ) + logging_obj.llm_caching_handler = LLMCachingHandler( + original_function=anthropic_messages, request_kwargs={}, start_time=datetime.datetime.now() + ) + + wrapped: Final = wrap_stream( + cache, anthropic_messages, _events(), read, **_REQUEST, litellm_logging_obj=logging_obj + ) + + assert isinstance(wrapped, AnthropicMessagesStreamCacheWriter) + assert wrapped.caching_handler is read + + +def test_caching_handler_is_the_one_the_logging_object_holds() -> None: + logging_obj: Final = _logging() + assert caching_handler(litellm_logging_obj=logging_obj) is None + handler: Final = LLMCachingHandler( + original_function=anthropic_messages, request_kwargs={}, start_time=datetime.datetime.now() + ) + logging_obj.llm_caching_handler = handler + assert caching_handler(litellm_logging_obj=logging_obj) is handler diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index cb594767908..2a448df6e5a 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -4,7 +4,7 @@ import inspect from collections.abc import Mapping from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, cast # noqa: TID251 # a duck-typed logger double stands in for Logging import httpx import pytest @@ -45,6 +45,20 @@ def test_native_stream_headers_reach_spend_callbacks() -> None: assert logger.model_call_details["response_headers"] == {"llm_provider-request-id": "req_native"} +def test_recorded_messages_stream_keeps_the_provider_header_names() -> None: + logger: Final = _supplied_logger() + sent: Final[tuple[str, Mapping[str, str], Mapping[str, object]]] = ( + "http://upstream/v1/messages", + {}, + {"stream": True}, + ) + legacy.record_response(logger, 200, [("request-id", "req_native")], "", sent) + legacy.stream_opened(logger, {"additional_headers": {"llm_provider-request-id": "req_native"}}) + + assert logger.stream is True + assert logger.model_call_details["response_headers"] == {"request-id": "req_native"} + + def test_setup_reuses_a_supplied_logger() -> None: supplied: Final = _supplied_logger() result: Final = setup( @@ -73,7 +87,7 @@ def test_cache_hit_finalization_preserves_execution_provider_attribution(explici ) prepared.logger.model_call_details.update({"cache_hit": True, "cache_key": "cached-response"}) response: Final = ModelResponse(model="cache-test-model") - legacy.finalize(response, prepared.logger, prepared.kwargs, now, now) + legacy.finalize(response, prepared.logger, prepared.kwargs, now, now, False) assert prepared.logger.model_call_details["custom_llm_provider"] == "azure" assert prepared.logger.model_call_details["model"] == "resolved-cache-model" assert prepared.logger.litellm_params["metadata"]["user_api_key"] == "key-hash" @@ -303,7 +317,7 @@ def _messages_logger() -> Logging: def test_native_provider_response_is_recorded_for_logging_replay() -> None: logger: Final = _messages_logger() - legacy.record_response(logger, 200, (("request-id", "req_native"),), '{"id": "msg_native"}') + legacy.record_response(logger, 200, (("request-id", "req_native"),), '{"id": "msg_native"}', None) recorded: Final[object] = logger.model_call_details["httpx_response"] assert isinstance(recorded, httpx.Response) @@ -312,6 +326,51 @@ def test_native_provider_response_is_recorded_for_logging_replay() -> None: assert logger.model_call_details["response_headers"]["request-id"] == "req_native" +def test_a_non_ascii_provider_header_is_recorded_as_httpx_decodes_it() -> None: + logger: Final = _messages_logger() + + legacy.record_response(logger, 200, (("x-fz-unicode", "caf\u00e9"),), '{"id": "msg_native"}', None) + + assert logger.model_call_details["response_headers"]["x-fz-unicode"] == "caf\u00e9" + + +_HTTPX_RESPONSE: Final = TypeAdapter(httpx.Response, config={"arbitrary_types_allowed": True}) +_SENT_HEADERS: Final = MappingProxyType({"x-api-key": "sk-ant", "anthropic-beta": "message-threads-2026-08-12"}) + + +def test_native_provider_response_carries_the_request_that_was_sent() -> None: + logger: Final = _messages_logger() + body: Final = {"model": "claude-sonnet-4-5", "thread": {"type": "create"}} + + legacy.record_response( + logger, 200, (), '{"id": "msg_native"}', ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, body) + ) + + recorded: Final = _HTTPX_RESPONSE.validate_python(logger.model_call_details["httpx_response"]) + assert recorded.request.method == "POST" + assert str(recorded.request.url) == "https://api.anthropic.com/v1/messages" + assert recorded.request.headers["anthropic-beta"] == "message-threads-2026-08-12" + assert recorded.request.content == b'{"model": "claude-sonnet-4-5", "thread": {"type": "create"}}' + assert recorded.content == b'{"id": "msg_native"}' + + +def test_native_stream_response_is_recorded_unread_like_the_python_stream() -> None: + logger: Final = _messages_logger() + + legacy.record_response( + logger, + 200, + (("request-id", "req_native"),), + "", + ("https://api.anthropic.com/v1/messages", _SENT_HEADERS, {"stream": True}), + ) + + recorded: Final = _HTTPX_RESPONSE.validate_python(logger.model_call_details["httpx_response"]) + assert recorded.request.headers["anthropic-beta"] == "message-threads-2026-08-12" + with pytest.raises(httpx.ResponseNotRead): + _ = recorded.content + + def test_native_provider_exchange_window_is_recorded_like_the_timed_post() -> None: logger: Final = _messages_logger() logger.model_call_details["litellm_params"] = {"litellm_metadata": {"user_api_key": "hashed"}} @@ -323,6 +382,76 @@ def test_native_provider_exchange_window_is_recorded_like_the_timed_post() -> No assert metadata["llm_api_timing_windows"] == ((1_700_000_000.0, 1_700_000_000.25),) +@pytest.mark.asyncio +async def test_a_retried_request_rewrites_the_logged_body_in_place_like_the_python_retry_loop() -> None: + logger: Final = _messages_logger() + metadata: Final = {"user_id": "u-1"} + replayed: Final = [ + {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]}, + {"role": "user", "content": "again"}, + ] + logged: Final[dict[str, object]] = { + "model": "claude-sonnet-4-5", + "messages": replayed, + "metadata": metadata, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + } + logger.model_call_details.update(logged) + retried: Final = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "again"}], + "metadata": {"user_id": "u-1"}, + } + + await legacy.retried_request(logger, logged, retried, ({"kind": "thinking_signature"},)) + + assert logged == retried + assert logged["metadata"] is metadata + assert logger.model_call_details["messages"] == [{"role": "user", "content": "again"}] + assert logger.model_call_details["thinking"] == {"type": "enabled", "budget_tokens": 1024} + + +class _BaselineLogger: + def __init__(self) -> None: + self.baseline_cache_context: object = object() + self.model_call_details: dict[str, object] = {} + self.invalidated: list[str] = [] + + async def invalidate_baseline_cache_estimate(self, reason: str) -> None: + self.invalidated.append(reason) + + +@pytest.mark.asyncio +async def test_a_retried_request_invalidates_the_baseline_cache_estimate() -> None: + logger: Final = _BaselineLogger() + + await legacy.retried_request(cast(Logging, logger), {"messages": []}, {"messages": []}, ()) # cast-ok: duck-typed + + assert logger.invalidated == ["retried_request"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reasons", "warned"), + [ + pytest.param(({"kind": "thinking_signature"},), [], id="thinking_signature"), + pytest.param(({"kind": "strict_tools", "dropped": 2},), [2], id="strict_tools"), + pytest.param(({"kind": "thinking_signature"}, {"kind": "strict_tools", "dropped": 3}), [3], id="both_reasons"), + ], +) +async def test_a_retry_that_dropped_strict_tools_warns_with_the_dropped_count( + monkeypatch: pytest.MonkeyPatch, reasons: tuple[Mapping[str, object], ...], warned: list[int] +) -> None: + from litellm.llms.anthropic import common_utils + + warnings: Final[list[int]] = [] + monkeypatch.setattr(common_utils, "warn_claude_strict_tools_dropped", warnings.append) + + await legacy.retried_request(_messages_logger(), {"messages": []}, {"messages": []}, reasons) + + assert warnings == warned + + def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: logger: Final = _messages_logger() head: Final[dict[str, object]] = {"additional_headers": {"request-id": "req_native"}} @@ -334,6 +463,7 @@ def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: {"model": "anthropic/claude-sonnet-4-5", "litellm_call_id": "messages"}, started, started + datetime.timedelta(milliseconds=5), + False, ) assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"} @@ -403,3 +533,22 @@ async def test_a_proxy_stream_abandoned_mid_cancel_bills_after_the_disconnect_cl await asyncio.sleep(0) assert events == ["disconnect recorded", "billed"] + + +@pytest.mark.parametrize( + ("requested", "keeps_requested", "expected"), + [ + pytest.param("false", True, "false", id="messages-keeps-a-truthy-string"), + pytest.param(1, True, 1, id="messages-keeps-a-truthy-int"), + pytest.param(False, True, True, id="messages-hook-streamed"), + pytest.param("false", False, True, id="other-routes-stream-true"), + ], +) +def test_a_stream_logs_the_flag_python_logs(requested: object, keeps_requested: bool, expected: object) -> None: + logger: Final = _supplied_logger() + logger.model_call_details["stream"] = requested + + legacy.stream_opened(logger, {}, keeps_requested) + + assert logger.model_call_details["stream"] == expected + assert logger.stream == expected diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 1994dd00709..56fcad14851 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -20,16 +20,10 @@ from litellm.rust_bridge.public_call import NativeCall @pytest.mark.parametrize( - ("model", "expected"), - ( - ("anthropic/claude-opus-latest", Rollout.PYTHON_ONLY), - (f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}", Rollout.RUST_OPT_IN), - ), + "model", + ("anthropic/claude-opus-latest", f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}"), ) -def test_latest_alias_messages_falls_back_without_changing_the_pinned_model_rollout( - model: str, - expected: Rollout, -) -> None: +def test_latest_alias_messages_stays_native_like_the_pinned_model(model: str) -> None: request: Final = NativeCall(args=(), kwargs={"model": model}, base={}) context: Final = RouteContext( Route.MESSAGES, @@ -41,7 +35,7 @@ def test_latest_alias_messages_falls_back_without_changing_the_pinned_model_roll RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=catalog.MESSAGES_PYTHON_FEATURES), RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), ) - assert catalog.rollout(context, rules) is expected + assert catalog.rollout(context, rules) is Rollout.RUST_OPT_IN @pytest.fixture(autouse=True) @@ -74,7 +68,7 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt (RouteContext(Route.MESSAGES, target=Target.MESSAGES), Decision.RUST_REQUIRED), (RouteContext(Route.MESSAGES, target=Target.CHAT_COMPLETIONS), Decision.PYTHON), (RouteContext(Route.MESSAGES), Decision.PYTHON), - (RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), Decision.PYTHON), + (RouteContext(Route.OCR, features=frozenset({Feature.STREAMING})), Decision.PYTHON), (RouteContext(Route.OCR), Decision.RUST_REQUIRED), ), ) @@ -88,7 +82,7 @@ def test_first_matching_rule_respects_every_constraint(context: RouteContext, ex ), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})), - RouteRule(Route.OCR, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY, features=frozenset({Feature.STREAMING})), RouteRule(Route.OCR, Rollout.RUST_REQUIRED), ) @@ -152,6 +146,7 @@ def test_empty_constraints_match_nothing() -> None: ) assert catalog.decision(RouteContext(Route.OCR), rules) is Decision.PYTHON - assert catalog.decision( - RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules - ) is Decision.PYTHON + assert ( + catalog.decision(RouteContext(Route.OCR, features=frozenset({Feature.STREAMING})), rules) + is Decision.PYTHON + ) diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index fccc01dd8e6..7ad8069758e 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -1,4 +1,5 @@ from collections.abc import Callable, Mapping +import ssl from copy import deepcopy from datetime import datetime from functools import reduce @@ -14,10 +15,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge import detectors from litellm.rust_bridge.chat_completions import detectors as chat_detectors -from litellm.rust_bridge.messages import detectors as messages_detectors from litellm.rust_bridge.public_call import NativeCall from litellm.rust_bridge.responses import detectors as responses_detectors -from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES +from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_TOOL_TYPES GATEWAY_TOOL: Final = {"type": "mcp", "server_label": "gateway", "server_url": "litellm_proxy"} REMOTE_TOOL: Final = {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"} @@ -31,9 +31,6 @@ MAPPED_FILE: Final = { } MAPPED_FILE_INPUT: Final = [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-file"}]}] CACHED_TEXT: Final = {"type": "input_text", "text": "hi", "cache_control": {"type": "ephemeral"}} -THINKING_TURN: Final = {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]} -REDACTED: Final = {"type": "redacted_thinking", "data": "x"} -STRICT_TOOL: Final = {"name": "b", "input_schema": {"type": "object"}, "strict": True} TOOL_SEARCH_TYPE: Final = sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0] @@ -63,25 +60,6 @@ def test_floating_alias_requires_python(model: str, provider: str | None, expect assert detectors.latest_alias(call(model=model, custom_llm_provider=provider)) is expected -@pytest.mark.parametrize( - ("model", "provider", "tools", "expected"), - ( - ("deepseek/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True), - ("anthropic/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False), - ("test-model", "anthropic", [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False), - ("unresolved-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True), - ("deepseek/test-model", None, [{"type": "function", "name": "f"}], False), - ("deepseek/test-model", None, ["not-a-tool"], False), - ("deepseek/test-model", None, [], False), - ("deepseek/test-model", None, None, False), - ), -) -def test_messages_interception_matches_tools_and_resolved_provider( - model: str, provider: str | None, tools: object, expected: bool -) -> None: - assert messages_detectors.intercepted(call(model=model, custom_llm_provider=provider, tools=tools)) is expected - - def _wrap_nested_level(value: object, depth: int) -> object: return {"nested": value} if depth % 2 else [value] @@ -136,8 +114,6 @@ def _wrap_nested_level(value: object, depth: int) -> object: (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), (responses_detectors.provider_affinity, {}, False), (responses_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/openai/v1"}, True), - (messages_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/anthropic/v1"}, True), - (messages_detectors.provider_affinity, {"api_base": "https://other.example/v1"}, False), (detectors.streaming, {"stream": True}, True), (detectors.streaming, {"stream": False}, False), (detectors.timeout_variants, {"request_timeout": 5}, True), @@ -196,36 +172,6 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.cache_control_injection, {"enable_prompt_caching": True}, True), (detectors.cache_control_injection, {"cache_control_injection_points": [{"location": "message"}]}, True), (detectors.cache_control_injection, {"cache_control_injection_points": []}, False), - (messages_detectors.thinking_signature_retry, {"messages": [THINKING_TURN]}, True), - ( - messages_detectors.thinking_signature_retry, - {"messages": [{"role": "assistant", "content": [REDACTED]}]}, - True, - ), - (messages_detectors.thinking_signature_retry, {"messages": [{"role": "assistant", "content": "done"}]}, False), - ( - messages_detectors.thinking_signature_retry, - {"messages": [{"role": "user", "content": [{"type": "text"}]}]}, - False, - ), - (messages_detectors.tool_custom_field, {"tools": [{"name": "a"}, {"name": "b", "custom": {}}]}, True), - (messages_detectors.tool_custom_field, {"tools": [{"name": "a", "defer_loading": True}]}, False), - (messages_detectors.tool_custom_field, {}, False), - (messages_detectors.strict_tools_retry, {"drop_params": True, "tools": [{"name": "a"}, STRICT_TOOL]}, True), - ( - messages_detectors.strict_tools_retry, - {"drop_params": True, "tools": [{**STRICT_TOOL, "strict": False}]}, - False, - ), - ( - messages_detectors.strict_tools_retry, - {"additional_drop_params": ["metadata.user_id"], "tools": [STRICT_TOOL]}, - False, - ), - (messages_detectors.strict_tools_retry, {"tools": [STRICT_TOOL]}, False), - (messages_detectors.strict_tools_retry, {"drop_params": False, "tools": [STRICT_TOOL]}, False), - (messages_detectors.strict_tools_retry, {"drop_params": True, "tools": [{"name": "a"}]}, False), - (messages_detectors.strict_tools_retry, {"drop_params": True}, False), (detectors.drop_params, {"drop_params": True}, True), (detectors.drop_params, {"drop_params": False}, False), (detectors.drop_params, {"drop_params": " True "}, True), @@ -323,8 +269,6 @@ def test_input_rewrite_preserves_caller_input(input_value: object, expected: boo (detectors.modify_params, "modify_params", True, {}), (detectors.retries_fallbacks, "num_retries", 3, {}), (detectors.cache_control_injection, "enable_anthropic_prompt_caching", True, {}), - (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", True, {"stream": True}), - (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", "true", {"stream": True}), (chat_detectors.chat_retries_fallbacks, "model_fallbacks", [{"gpt-5": ["gpt-5-mini"]}], {}), (chat_detectors.chat_retries_fallbacks, "num_retries", 3, {}), (chat_detectors.model_alias_map, "model_alias_map", {"gpt-5": "openai/gpt-5"}, {}), @@ -347,30 +291,6 @@ def test_global_settings_are_detected( assert detector(call(**kwargs)) -def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) - assert not messages_detectors.stream_cost_injection(call()) - - -def test_stream_cost_injection_reads_the_request_header_the_proxy_reads(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) - opted_in: Final = {"proxy_server_request": {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}} - assert messages_detectors.stream_cost_injection(call(stream=True, **opted_in)) - assert not messages_detectors.stream_cost_injection(call(**opted_in)) - assert not messages_detectors.stream_cost_injection(call(stream=True, proxy_server_request={"headers": {}})) - - -def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None: - for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): - monkeypatch.delenv(name, raising=False) - assert messages_detectors.foundry_beta_row(call(api_base="https://team.services.ai.azure.com/anthropic")) - assert messages_detectors.foundry_beta_row(call(base_url="https://team.cognitiveservices.azure.com/anthropic")) - assert not messages_detectors.foundry_beta_row(call(api_base="https://api.anthropic.com")) - assert not messages_detectors.foundry_beta_row(call()) - monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://team.services.ai.azure.com/anthropic") - assert messages_detectors.foundry_beta_row(call()) - - def test_model_alias_map_applies_only_to_the_requested_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_alias_map", {"fast": "openai/gpt-5-mini"}) assert chat_detectors.model_alias_map(call(model="fast")) @@ -413,20 +333,6 @@ def test_openai_workload_identity_applies_only_without_a_static_key(monkeypatch: assert not responses_detectors.openai_workload_identity(call()) -def test_anthropic_workload_identity_applies_only_without_a_static_key(monkeypatch: pytest.MonkeyPatch) -> None: - for name in ( - "ANTHROPIC_API_KEY", - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_FEDERATION_RULE_ID", - "ANTHROPIC_ORGANIZATION_ID", - ): - monkeypatch.delenv(name, raising=False) - federation: Final = {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"} - assert messages_detectors.anthropic_workload_identity(call(**federation)) - assert not messages_detectors.anthropic_workload_identity(call(**federation, api_key="sk-ant-static")) - assert not messages_detectors.anthropic_workload_identity(call(anthropic_federation_rule_id="rule-1")) - - class _Logger(CustomLogger): pass @@ -436,10 +342,6 @@ class _Rewrite(CustomLogger): return kwargs -class _InheritedRewrite(_Rewrite): - pass - - class _Agentic(CustomLogger): async def async_should_run_agentic_loop(self, *args: object, **kwargs: object) -> tuple[bool, dict]: # pyright: ignore[reportIncompatibleMethodOverride, reportMissingTypeArgument] # test double only needs the override return False, {} @@ -460,10 +362,6 @@ class _Deployment(CustomLogger): @pytest.mark.parametrize( ("detector", "callbacks", "expected"), ( - (messages_detectors.pre_request_hook, [_Rewrite()], True), - (messages_detectors.pre_request_hook, [_InheritedRewrite()], True), - (messages_detectors.pre_request_hook, [_Logger(), _Agentic()], False), - (messages_detectors.pre_request_hook, ["not-a-logger"], False), (detectors.agentic_loop_hook, [_Agentic()], True), (detectors.agentic_loop_hook, [_Logger(), _Rewrite()], False), (detectors.agentic_loop_hook, [_ChatAgentic()], False), @@ -501,7 +399,6 @@ def _logging_with(callback: CustomLogger) -> Logging: (detectors.agentic_loop_hook, {"success_callback": [_Agentic()]}, True), (detectors.agentic_loop_hook, {"litellm_logging_obj": _logging_with(_Agentic())}, True), (detectors.agentic_loop_hook, {"success_callback": [_Logger()]}, False), - (messages_detectors.pre_request_hook, {"success_callback": [_Rewrite()]}, False), (chat_detectors.chat_agentic_loop_hook, {"litellm_logging_obj": _logging_with(_ChatAgentic())}, True), (responses_detectors.deployment_hook, {"success_callback": [_Deployment()]}, True), ), @@ -509,7 +406,6 @@ def _logging_with(callback: CustomLogger) -> Logging: "success-callback", "logging-object", "plain-logger", - "pre-request-reads-globals-only", "chat-agentic-logging-object", "deployment-success-callback", ), @@ -524,15 +420,6 @@ def test_callback_detectors_read_the_request_callbacks_the_handler_reads( assert detector(call(**kwargs)) is expected -def test_websearch_interception_needs_the_interception_logger(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger - - monkeypatch.setattr(litellm, "callbacks", [_Rewrite()]) - assert not messages_detectors.websearch_interception(call()) - monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger()]) - assert messages_detectors.websearch_interception(call()) - - @pytest.mark.parametrize( ("cache", "kwargs", "expected"), ( @@ -551,7 +438,14 @@ def test_response_cache_is_skipped_only_by_no_cache( @pytest.mark.parametrize( ("environment", "setting", "expected"), - ((None, None, False), ("DEFAULT@SECLEVEL=1", None, True), (None, "DEFAULT@SECLEVEL=1", True)), + ( + (None, None, False), + ("DEFAULT@SECLEVEL=1", None, True), + (None, "DEFAULT@SECLEVEL=1", True), + (None, "ECDHE-RSA-AES128-GCM-SHA256", False), + (None, "ECDHE-RSA-AES128-GCM-SHA256:RC4-SHA", True), + ("HIGH", None, False), + ), ) def test_tls_policy_reads_the_security_level_environment_then_the_setting( monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool @@ -567,7 +461,19 @@ def test_tls_policy_reads_the_security_level_environment_then_the_setting( @pytest.mark.parametrize( ("environment", "setting", "expected"), - ((None, "X25519", False), (None, " P-256 ", False), ("secp521r1", None, True), (None, "secp521r1", True)), + ( + (None, "X25519", False), + (None, "prime256v1", False), + (None, "secp384r1", False), + (None, " P-256 ", True), + (None, "P-384", True), + (None, "secp256r1", True), + (None, "x25519", True), + ("secp521r1", None, True), + (None, "secp521r1", True), + ("", "secp521r1", False), + (None, "", False), + ), ) def test_tls_policy_flags_only_curves_native_tls_lacks( monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool @@ -579,3 +485,24 @@ def test_tls_policy_flags_only_curves_native_tls_lacks( monkeypatch.setenv("SSL_ECDH_CURVE", environment) monkeypatch.setattr(litellm, "ssl_ecdh_curve", setting) assert detectors.tls_policy(call()) is expected + + +@pytest.mark.parametrize(("setting", "value"), (("ssl_security_level", 1), ("ssl_ecdh_curve", object()))) +def test_tls_policy_hands_an_undecodable_setting_to_python_without_raising( + monkeypatch: pytest.MonkeyPatch, setting: str, value: object +) -> None: + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + monkeypatch.setattr(litellm, setting, value) + assert detectors.tls_policy(call()) is True + + +def test_tls_policy_ignores_a_live_ssl_context(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + monkeypatch.setattr(litellm, "ssl_verify", ssl.create_default_context()) + assert detectors.tls_policy(call()) is False diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index 75cdfde087b..548591a2903 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/tests/unit/rust_bridge/test_runtime.py @@ -216,9 +216,12 @@ async def test_python_fallback_does_not_claim_rust_execution(missing: bool) -> N @pytest.mark.parametrize("shape", ("model", "dict")) @pytest.mark.parametrize("asynchronous", (False, True)) @pytest.mark.parametrize("cache_key", (None, "test-cache-key")) +@pytest.mark.parametrize("route", (Route.MESSAGES, Route.OCR)) async def test_native_response_marker_reaches_caller_with_existing_metadata( - shape: str, asynchronous: bool, cache_key: str | None + shape: str, asynchronous: bool, cache_key: str | None, route: Route ) -> None: + context: Final = RouteContext(route, provider="anthropic", model="model") + required: Final = (RouteRule(route, Rollout.RUST_REQUIRED, providers=frozenset({"anthropic"})),) hidden: Final = { "additional_headers": {"x-request-id": "upstream"}, "response_cost": 0.01, @@ -232,6 +235,47 @@ async def test_native_response_marker_reaches_caller_with_existing_metadata( bound: Final[bindings.NativeBinding[Callable[[], object]]] = bindings.NativeBinding("ocr", validate=lambda _: None) bound.override(lambda: response) + def python() -> object: + pytest.fail("native success must not fall back") + + async def anative(fn: Callable[[], object]) -> object: + return fn() + + async def apython() -> object: + return python() + + result: Final = ( + await runtime.arun(context, binding=bound, native=anative, python=apython, rules=required) + if asynchronous + else runtime.run(context, binding=bound, native=lambda fn: fn(), python=python, rules=required) + ) + assert result is response + assert get_hidden_params_dict(result) == ( + { + "response_cost": 0.01, + "additional_headers": {"x-request-id": "upstream"}, + **({"cache_key": cache_key} if cache_key is not None else {}), + } + if shape == "dict" + else { + "response_cost": 0.01, + "additional_headers": { + "x-request-id": "upstream", + "x-litellm-rust": "true", + **({"x-litellm-cache-key": cache_key} if cache_key is not None and route is not Route.MESSAGES else {}), + }, + **({"cache_key": cache_key} if cache_key is not None else {}), + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_native_dict_response_reaches_caller_without_hidden_params_key(asynchronous: bool) -> None: + response: Final[dict[str, object]] = {"content": "native"} + bound: Final[bindings.NativeBinding[Callable[[], object]]] = bindings.NativeBinding("ocr", validate=lambda _: None) + bound.override(lambda: response) + def python() -> object: pytest.fail("native success must not fall back") @@ -249,15 +293,7 @@ async def test_native_response_marker_reaches_caller_with_existing_metadata( ) ) assert result is response - assert get_hidden_params_dict(result) == { - "response_cost": 0.01, - "additional_headers": { - "x-request-id": "upstream", - "x-litellm-rust": "true", - **({"x-litellm-cache-key": cache_key} if cache_key is not None else {}), - }, - **({"cache_key": cache_key} if cache_key is not None else {}), - } + assert result == {"content": "native"} class ScriptedStreamExecution: diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 6ceea94d843..44b5876a078 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1223,6 +1223,13 @@ ALIAS_CASES: Final = ( False, ), ("litellm.caching.caching_handler", "LLMCachingHandler", "_async_get_cache", "async_get_cache", False), + ( + "litellm.caching.caching_handler", + "LLMCachingHandler", + "_async_log_cache_hit_on_callbacks", + "async_log_cache_hit_on_callbacks", + False, + ), ( "litellm.caching.caching_handler", "LLMCachingHandler", @@ -1957,7 +1964,13 @@ ALIAS_CASES: Final = ( "get_tags_from_request_kwargs", False, ), - ("litellm.litellm_core_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), + ( + "litellm.litellm_core_utils.add_retry_fallback_headers", + "HiddenParamsAsyncIteratorWrapper", + "_inner", + "inner", + True, + ), ("litellm.litellm_core_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), ( "litellm.litellm_core_utils.batch_utils", @@ -2301,1837 +2314,1837 @@ LLMS_ALIAS_CASES: Final = ( ) PROXY_ALIAS_CASES: Final = ( ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - '', - '_is_mcp_admitted_user_subject', - 'is_mcp_admitted_user_subject', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "", + "_is_mcp_admitted_user_subject", + "is_mcp_admitted_user_subject", False, ), ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - 'MCPRequestHandler', - '_get_mcp_auth_header_from_headers', - 'get_mcp_auth_header_from_headers', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "MCPRequestHandler", + "_get_mcp_auth_header_from_headers", + "get_mcp_auth_header_from_headers", False, ), ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - 'MCPRequestHandler', - '_get_mcp_server_auth_headers_from_headers', - 'get_mcp_server_auth_headers_from_headers', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "MCPRequestHandler", + "_get_mcp_server_auth_headers_from_headers", + "get_mcp_server_auth_headers_from_headers", False, ), ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - 'MCPRequestHandler', - '_get_mcp_servers_from_access_groups', - 'get_mcp_servers_from_access_groups', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "MCPRequestHandler", + "_get_mcp_servers_from_access_groups", + "get_mcp_servers_from_access_groups", False, ), ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - 'MCPRequestHandler', - '_get_oauth2_headers_from_headers', - 'get_oauth2_headers_from_headers', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "MCPRequestHandler", + "_get_oauth2_headers_from_headers", + "get_oauth2_headers_from_headers", False, ), ( - 'litellm.proxy.mcp_server.auth.user_api_key_auth_mcp', - 'MCPRequestHandler', - '_safe_get_headers_from_scope', - 'safe_get_headers_from_scope', + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp", + "MCPRequestHandler", + "_safe_get_headers_from_scope", + "safe_get_headers_from_scope", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_bridge_mint_error_response', - 'bridge_mint_error_response', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_bridge_mint_error_response", + "bridge_mint_error_response", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_extract_user_id_from_request', - 'extract_user_id_from_request', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_extract_user_id_from_request", + "extract_user_id_from_request", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_finish_bridge_mint', - 'finish_bridge_mint', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_finish_bridge_mint", + "finish_bridge_mint", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_prepare_bridge_mint', - 'prepare_bridge_mint', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_prepare_bridge_mint", + "prepare_bridge_mint", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_prepare_bridge_refresh', - 'prepare_bridge_refresh', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_prepare_bridge_refresh", + "prepare_bridge_refresh", False, ), ( - 'litellm.proxy.mcp_server.bridge_token_flow', - '', - '_reload_active_user_by_id', - 'reload_active_user_by_id', + "litellm.proxy.mcp_server.bridge_token_flow", + "", + "_reload_active_user_by_id", + "reload_active_user_by_id", False, ), ( - 'litellm.proxy.mcp_server.byok_oauth_endpoints', - '', - '_user_id_from_session_cookie', - 'user_id_from_session_cookie', + "litellm.proxy.mcp_server.byok_oauth_endpoints", + "", + "_user_id_from_session_cookie", + "user_id_from_session_cookie", False, ), ( - 'litellm.proxy.mcp_server.db', - '', - '_decode_oauth_payload', - 'decode_oauth_payload', + "litellm.proxy.mcp_server.db", + "", + "_decode_oauth_payload", + "decode_oauth_payload", False, ), ( - 'litellm.proxy.mcp_server.discoverable_endpoints', - '', - '_raise_if_not_oauth2', - 'raise_if_not_oauth2', + "litellm.proxy.mcp_server.discoverable_endpoints", + "", + "_raise_if_not_oauth2", + "raise_if_not_oauth2", False, ), ( - 'litellm.proxy.mcp_server.mcp_context', - '', - '_mcp_active_toolset_id', - 'mcp_active_toolset_id', + "litellm.proxy.mcp_server.mcp_context", + "", + "_mcp_active_toolset_id", + "mcp_active_toolset_id", False, ), ( - 'litellm.proxy.mcp_server.mcp_context', - '', - '_mcp_gateway_initialize_instructions', - 'mcp_gateway_initialize_instructions', + "litellm.proxy.mcp_server.mcp_context", + "", + "_mcp_gateway_initialize_instructions", + "mcp_gateway_initialize_instructions", False, ), ( - 'litellm.proxy.mcp_server.mcp_context', - '', - '_mcp_gateway_server_name', - 'mcp_gateway_server_name', + "litellm.proxy.mcp_server.mcp_context", + "", + "_mcp_gateway_server_name", + "mcp_gateway_server_name", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - '', - '_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES', - 'UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES', + "litellm.proxy.mcp_server.mcp_server_manager", + "", + "_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES", + "UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - '', - '_caller_authorization_fans_out', - 'caller_authorization_fans_out', + "litellm.proxy.mcp_server.mcp_server_manager", + "", + "_caller_authorization_fans_out", + "caller_authorization_fans_out", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - '', - '_client_forwarded_authorization_headers', - 'client_forwarded_authorization_headers', + "litellm.proxy.mcp_server.mcp_server_manager", + "", + "_client_forwarded_authorization_headers", + "client_forwarded_authorization_headers", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - '', - '_resolve_openapi_tool_auth', - 'resolve_openapi_tool_auth', + "litellm.proxy.mcp_server.mcp_server_manager", + "", + "_resolve_openapi_tool_auth", + "resolve_openapi_tool_auth", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - '', - '_should_strip_caller_authorization', - 'should_strip_caller_authorization', + "litellm.proxy.mcp_server.mcp_server_manager", + "", + "_should_strip_caller_authorization", + "should_strip_caller_authorization", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_build_mcp_server_table', - 'build_mcp_server_table', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_build_mcp_server_table", + "build_mcp_server_table", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_build_stdio_env', - 'build_stdio_env', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_build_stdio_env", + "build_stdio_env", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_create_mcp_client', - 'create_mcp_client', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_create_mcp_client", + "create_mcp_client", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_ensure_upstream_initialize_instructions_cached', - 'ensure_upstream_initialize_instructions_cached', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_ensure_upstream_initialize_instructions_cached", + "ensure_upstream_initialize_instructions_cached", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_extract_subject_token', - 'extract_subject_token', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_extract_subject_token", + "extract_subject_token", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_get_mcp_server_from_tool_name', - 'get_mcp_server_from_tool_name', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_get_mcp_server_from_tool_name", + "get_mcp_server_from_tool_name", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_get_tools_from_server', - 'get_tools_from_server', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_get_tools_from_server", + "get_tools_from_server", False, ), ( - 'litellm.proxy.mcp_server.mcp_server_manager', - 'MCPServerManager', - '_is_server_accessible_from_ip', - 'is_server_accessible_from_ip', + "litellm.proxy.mcp_server.mcp_server_manager", + "MCPServerManager", + "_is_server_accessible_from_ip", + "is_server_accessible_from_ip", False, ), ( - 'litellm.proxy.mcp_server.oauth2_token_cache', - '', - '_compute_per_user_token_ttl', - 'compute_per_user_token_ttl', + "litellm.proxy.mcp_server.oauth2_token_cache", + "", + "_compute_per_user_token_ttl", + "compute_per_user_token_ttl", False, ), ( - 'litellm.proxy.mcp_server.oauth_utils', - '', - '_redact_mcp_resource_url', - 'redact_mcp_resource_url', + "litellm.proxy.mcp_server.oauth_utils", + "", + "_redact_mcp_resource_url", + "redact_mcp_resource_url", False, ), ( - 'litellm.proxy.mcp_server.server', - '', - '_redact_mcp_resource_url', - 'redact_mcp_resource_url', + "litellm.proxy.mcp_server.server", + "", + "_redact_mcp_resource_url", + "redact_mcp_resource_url", False, ), ( - 'litellm.proxy.mcp_server.openapi_to_mcp_generator', - '', - '_OPENAPI_TOOL_NAME_MAX_LEN', - 'OPENAPI_TOOL_NAME_MAX_LEN', + "litellm.proxy.mcp_server.openapi_to_mcp_generator", + "", + "_OPENAPI_TOOL_NAME_MAX_LEN", + "OPENAPI_TOOL_NAME_MAX_LEN", False, ), ( - 'litellm.proxy.mcp_server.openapi_to_mcp_generator', - '', - '_request_auth_header', - 'request_auth_header', + "litellm.proxy.mcp_server.openapi_to_mcp_generator", + "", + "_request_auth_header", + "request_auth_header", False, ), ( - 'litellm.proxy.mcp_server.openapi_to_mcp_generator', - '', - '_request_extra_headers', - 'request_extra_headers', + "litellm.proxy.mcp_server.openapi_to_mcp_generator", + "", + "_request_extra_headers", + "request_extra_headers", False, ), ( - 'litellm.proxy.mcp_server.openapi_to_mcp_generator', - '', - '_request_resolved_auth_headers', - 'request_resolved_auth_headers', + "litellm.proxy.mcp_server.openapi_to_mcp_generator", + "", + "_request_resolved_auth_headers", + "request_resolved_auth_headers", False, ), ( - 'litellm.proxy.mcp_server.semantic_tool_filter', - 'SemanticMCPToolFilter', - '_extract_tool_info', - 'extract_tool_info', + "litellm.proxy.mcp_server.semantic_tool_filter", + "SemanticMCPToolFilter", + "_extract_tool_info", + "extract_tool_info", False, ), ( - 'litellm.proxy.mcp_server.server', - '', - '_apply_toolset_scope', - 'apply_toolset_scope', + "litellm.proxy.mcp_server.server", + "", + "_apply_toolset_scope", + "apply_toolset_scope", False, ), ( - 'litellm.proxy.mcp_server.server_resolution', - 'MCPServerRegistry', - '_build_mcp_server_table', - 'build_mcp_server_table', + "litellm.proxy.mcp_server.server_resolution", + "MCPServerRegistry", + "_build_mcp_server_table", + "build_mcp_server_table", False, ), ( - 'litellm.proxy.mcp_server.server_resolution', - 'MCPServerRegistry', - '_is_server_accessible_from_ip', - 'is_server_accessible_from_ip', + "litellm.proxy.mcp_server.server_resolution", + "MCPServerRegistry", + "_is_server_accessible_from_ip", + "is_server_accessible_from_ip", False, ), ( - 'litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace', - '', - '_get_prisma_client', - 'get_prisma_client', + "litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace", + "", + "_get_prisma_client", + "get_prisma_client", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_cache_access_object', - 'cache_access_object', + "litellm.proxy.auth.auth_checks", + "", + "_cache_access_object", + "cache_access_object", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_cache_key_object', - 'cache_key_object', + "litellm.proxy.auth.auth_checks", + "", + "_cache_key_object", + "cache_key_object", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_cache_team_object', - 'cache_team_object', + "litellm.proxy.auth.auth_checks", + "", + "_cache_team_object", + "cache_team_object", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_can_object_call_model', - 'can_object_call_model', + "litellm.proxy.auth.auth_checks", + "", + "_can_object_call_model", + "can_object_call_model", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_check_end_user_budget', - 'check_end_user_budget', + "litellm.proxy.auth.auth_checks", + "", + "_check_end_user_budget", + "check_end_user_budget", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_check_model_access_helper', - 'check_model_access_helper', + "litellm.proxy.auth.auth_checks", + "", + "_check_model_access_helper", + "check_model_access_helper", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_check_team_member_model_access', - 'check_team_member_model_access', + "litellm.proxy.auth.auth_checks", + "", + "_check_team_member_model_access", + "check_team_member_model_access", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_copy_user_api_key_auth_for_cache', - 'copy_user_api_key_auth_for_cache', + "litellm.proxy.auth.auth_checks", + "", + "_copy_user_api_key_auth_for_cache", + "copy_user_api_key_auth_for_cache", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_delete_cache_access_object', - 'delete_cache_access_object', + "litellm.proxy.auth.auth_checks", + "", + "_delete_cache_access_object", + "delete_cache_access_object", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_delete_cache_key_object', - 'delete_cache_key_object', + "litellm.proxy.auth.auth_checks", + "", + "_delete_cache_key_object", + "delete_cache_key_object", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_fetch_key_object_from_db_with_reconnect', - 'fetch_key_object_from_db_with_reconnect', + "litellm.proxy.auth.auth_checks", + "", + "_fetch_key_object_from_db_with_reconnect", + "fetch_key_object_from_db_with_reconnect", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_get_agent_ids_from_access_groups', - 'get_agent_ids_from_access_groups', + "litellm.proxy.auth.auth_checks", + "", + "_get_agent_ids_from_access_groups", + "get_agent_ids_from_access_groups", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_get_mcp_server_ids_from_access_groups', - 'get_mcp_server_ids_from_access_groups', + "litellm.proxy.auth.auth_checks", + "", + "_get_mcp_server_ids_from_access_groups", + "get_mcp_server_ids_from_access_groups", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_get_models_from_access_groups', - 'get_models_from_access_groups', + "litellm.proxy.auth.auth_checks", + "", + "_get_models_from_access_groups", + "get_models_from_access_groups", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_get_team_object_from_cache', - 'get_team_object_from_cache', + "litellm.proxy.auth.auth_checks", + "", + "_get_team_object_from_cache", + "get_team_object_from_cache", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_get_user_role', - 'get_user_role', + "litellm.proxy.auth.auth_checks", + "", + "_get_user_role", + "get_user_role", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_is_model_cost_zero', - 'is_model_cost_zero', + "litellm.proxy.auth.auth_checks", + "", + "_is_model_cost_zero", + "is_model_cost_zero", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_is_user_proxy_admin', - 'is_user_proxy_admin', + "litellm.proxy.auth.auth_checks", + "", + "_is_user_proxy_admin", + "is_user_proxy_admin", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_key_access_group_grants_model', - 'key_access_group_grants_model', + "litellm.proxy.auth.auth_checks", + "", + "_key_access_group_grants_model", + "key_access_group_grants_model", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_organization_max_budget_check', - 'organization_max_budget_check', + "litellm.proxy.auth.auth_checks", + "", + "_organization_max_budget_check", + "organization_max_budget_check", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_team_max_budget_check', - 'team_max_budget_check', + "litellm.proxy.auth.auth_checks", + "", + "_team_max_budget_check", + "team_max_budget_check", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_team_member_max_budget_alert_check', - 'team_member_max_budget_alert_check', + "litellm.proxy.auth.auth_checks", + "", + "_team_member_max_budget_alert_check", + "team_member_max_budget_alert_check", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_virtual_key_max_budget_alert_check', - 'virtual_key_max_budget_alert_check', + "litellm.proxy.auth.auth_checks", + "", + "_virtual_key_max_budget_alert_check", + "virtual_key_max_budget_alert_check", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_virtual_key_max_budget_check', - 'virtual_key_max_budget_check', + "litellm.proxy.auth.auth_checks", + "", + "_virtual_key_max_budget_check", + "virtual_key_max_budget_check", False, ), ( - 'litellm.proxy.auth.auth_checks', - '', - '_virtual_key_soft_budget_check', - 'virtual_key_soft_budget_check', + "litellm.proxy.auth.auth_checks", + "", + "_virtual_key_soft_budget_check", + "virtual_key_soft_budget_check", False, ), ( - 'litellm.proxy.auth.auth_checks_organization', - '', - '_user_is_org_admin', - 'user_is_org_admin', + "litellm.proxy.auth.auth_checks_organization", + "", + "_user_is_org_admin", + "user_is_org_admin", False, ), ( - 'litellm.proxy.auth.auth_exception_handler', - 'UserAPIKeyAuthExceptionHandler', - '_handle_authentication_error', - 'handle_authentication_error', + "litellm.proxy.auth.auth_exception_handler", + "UserAPIKeyAuthExceptionHandler", + "_handle_authentication_error", + "handle_authentication_error", False, ), ( - 'litellm.proxy.auth.auth_utils', - '', - '_get_request_ip_address', - 'get_request_ip_address', + "litellm.proxy.auth.auth_utils", + "", + "_get_request_ip_address", + "get_request_ip_address", False, ), ( - 'litellm.proxy.auth.resolvers.store', - 'IdentityStore', - '_principal_from_key', - 'principal_from_key', + "litellm.proxy.auth.resolvers.store", + "IdentityStore", + "_principal_from_key", + "principal_from_key", False, ), ( - 'litellm.proxy.auth.route_checks', - 'RouteChecks', - '_get_request_method', - 'get_request_method', + "litellm.proxy.auth.route_checks", + "RouteChecks", + "_get_request_method", + "get_request_method", False, ), ( - 'litellm.proxy.auth.route_checks', - 'RouteChecks', - '_is_assistants_api_request', - 'is_assistants_api_request', + "litellm.proxy.auth.route_checks", + "RouteChecks", + "_is_assistants_api_request", + "is_assistants_api_request", False, ), ( - 'litellm.proxy.auth.route_checks', - 'RouteChecks', - '_is_wildcard_pattern', - 'is_wildcard_pattern', + "litellm.proxy.auth.route_checks", + "RouteChecks", + "_is_wildcard_pattern", + "is_wildcard_pattern", False, ), ( - 'litellm.proxy.auth.user_api_key_auth', - '', - '_enforce_key_and_fallback_model_access', - 'enforce_key_and_fallback_model_access', + "litellm.proxy.auth.user_api_key_auth", + "", + "_enforce_key_and_fallback_model_access", + "enforce_key_and_fallback_model_access", False, ), ( - 'litellm.proxy.auth.user_api_key_auth', - '', - '_fetch_global_spend_with_event_coordination', - 'fetch_global_spend_with_event_coordination', + "litellm.proxy.auth.user_api_key_auth", + "", + "_fetch_global_spend_with_event_coordination", + "fetch_global_spend_with_event_coordination", False, ), ( - 'litellm.proxy.auth.user_api_key_auth', - '', - '_get_bearer_token', - 'get_bearer_token', + "litellm.proxy.auth.user_api_key_auth", + "", + "_get_bearer_token", + "get_bearer_token", False, ), ( - 'litellm.proxy.auth.user_api_key_auth', - '', - '_run_centralized_common_checks', - 'run_centralized_common_checks', + "litellm.proxy.auth.user_api_key_auth", + "", + "_run_centralized_common_checks", + "run_centralized_common_checks", False, ), ( - 'litellm.proxy.auth.user_api_key_auth', - '', - '_user_api_key_auth_builder', - 'user_api_key_auth_builder', + "litellm.proxy.auth.user_api_key_auth", + "", + "_user_api_key_auth_builder", + "user_api_key_auth_builder", False, ), ( - 'litellm.proxy.common_request_processing', - '', - '_is_azure_model_router_request', - 'is_azure_model_router_request', + "litellm.proxy.common_request_processing", + "", + "_is_azure_model_router_request", + "is_azure_model_router_request", False, ), ( - 'litellm.proxy.common_request_processing', - '', - '_should_return_raw_model_name', - 'should_return_raw_model_name', + "litellm.proxy.common_request_processing", + "", + "_should_return_raw_model_name", + "should_return_raw_model_name", False, ), ( - 'litellm.proxy.common_request_processing', - 'ProxyBaseLLMRequestProcessing', - '_finalize_streaming_generator_cleanup', - 'finalize_streaming_generator_cleanup', + "litellm.proxy.common_request_processing", + "ProxyBaseLLMRequestProcessing", + "_finalize_streaming_generator_cleanup", + "finalize_streaming_generator_cleanup", False, ), ( - 'litellm.proxy.common_request_processing', - 'ProxyBaseLLMRequestProcessing', - '_handle_llm_api_exception', - 'handle_llm_api_exception', + "litellm.proxy.common_request_processing", + "ProxyBaseLLMRequestProcessing", + "_handle_llm_api_exception", + "handle_llm_api_exception", False, ), ( - 'litellm.proxy.common_request_processing', - 'ProxyBaseLLMRequestProcessing', - '_process_chunk_with_cost_injection', - 'process_chunk_with_cost_injection', + "litellm.proxy.common_request_processing", + "ProxyBaseLLMRequestProcessing", + "_process_chunk_with_cost_injection", + "process_chunk_with_cost_injection", False, ), ( - 'litellm.proxy.common_utils.callback_utils', - '', - '_CALLBACK_VAR_ENCRYPTED_PREFIX', - 'CALLBACK_VAR_ENCRYPTED_PREFIX', + "litellm.proxy.common_utils.callback_utils", + "", + "_CALLBACK_VAR_ENCRYPTED_PREFIX", + "CALLBACK_VAR_ENCRYPTED_PREFIX", False, ), ( - 'litellm.proxy.common_utils.config_sync_pubsub', - '', - '_ConfigSyncPubSub', - 'ConfigSyncPubSub', + "litellm.proxy.common_utils.config_sync_pubsub", + "", + "_ConfigSyncPubSub", + "ConfigSyncPubSub", False, ), ( - 'litellm.proxy.common_utils.config_sync_pubsub', - '', - '_pubsub_capable_client', - 'pubsub_capable_client', + "litellm.proxy.common_utils.config_sync_pubsub", + "", + "_pubsub_capable_client", + "pubsub_capable_client", False, ), ( - 'litellm.proxy.common_utils.encrypt_decrypt_utils', - '', - '_ALGO_AES_GCM', - 'ALGO_AES_GCM', + "litellm.proxy.common_utils.encrypt_decrypt_utils", + "", + "_ALGO_AES_GCM", + "ALGO_AES_GCM", False, ), ( - 'litellm.proxy.common_utils.encrypt_decrypt_utils', - '', - '_ENCRYPTION_ALGORITHM_SETTING', - 'ENCRYPTION_ALGORITHM_SETTING', + "litellm.proxy.common_utils.encrypt_decrypt_utils", + "", + "_ENCRYPTION_ALGORITHM_SETTING", + "ENCRYPTION_ALGORITHM_SETTING", False, ), ( - 'litellm.proxy.common_utils.encrypt_decrypt_utils', - '', - '_V2_GCM_PREFIX', - 'V2_GCM_PREFIX', + "litellm.proxy.common_utils.encrypt_decrypt_utils", + "", + "_V2_GCM_PREFIX", + "V2_GCM_PREFIX", False, ), ( - 'litellm.proxy.common_utils.encrypt_decrypt_utils', - '', - '_get_salt_key', - 'get_salt_key', + "litellm.proxy.common_utils.encrypt_decrypt_utils", + "", + "_get_salt_key", + "get_salt_key", False, ), ( - 'litellm.proxy.common_utils.http_parsing_utils', - '', - '_read_request_body', - 'read_request_body', + "litellm.proxy.common_utils.http_parsing_utils", + "", + "_read_request_body", + "read_request_body", False, ), ( - 'litellm.proxy.common_utils.http_parsing_utils', - '', - '_safe_get_request_headers', - 'safe_get_request_headers', + "litellm.proxy.common_utils.http_parsing_utils", + "", + "_safe_get_request_headers", + "safe_get_request_headers", False, ), ( - 'litellm.proxy.common_utils.http_parsing_utils', - '', - '_safe_get_request_query_params', - 'safe_get_request_query_params', + "litellm.proxy.common_utils.http_parsing_utils", + "", + "_safe_get_request_query_params", + "safe_get_request_query_params", False, ), ( - 'litellm.proxy.common_utils.http_parsing_utils', - '', - '_safe_set_request_parsed_body', - 'safe_set_request_parsed_body', + "litellm.proxy.common_utils.http_parsing_utils", + "", + "_safe_set_request_parsed_body", + "safe_set_request_parsed_body", False, ), ( - 'litellm.proxy.common_utils.realtime_utils', - '', - '_realtime_request_body', - 'realtime_request_body', + "litellm.proxy.common_utils.realtime_utils", + "", + "_realtime_request_body", + "realtime_request_body", False, ), ( - 'litellm.proxy.db.db_spend_update_writer', - 'DBSpendUpdateWriter', - '_commit_daily_tag_spend_to_db', - 'commit_daily_tag_spend_to_db', + "litellm.proxy.db.db_spend_update_writer", + "DBSpendUpdateWriter", + "_commit_daily_tag_spend_to_db", + "commit_daily_tag_spend_to_db", False, ), ( - 'litellm.proxy.db.db_spend_update_writer', - 'DBSpendUpdateWriter', - '_commit_daily_tag_spend_to_db_with_redis', - 'commit_daily_tag_spend_to_db_with_redis', + "litellm.proxy.db.db_spend_update_writer", + "DBSpendUpdateWriter", + "_commit_daily_tag_spend_to_db_with_redis", + "commit_daily_tag_spend_to_db_with_redis", False, ), ( - 'litellm.proxy.db.db_spend_update_writer', - 'DBSpendUpdateWriter', - '_handle_spend_update_failure', - 'handle_spend_update_failure', + "litellm.proxy.db.db_spend_update_writer", + "DBSpendUpdateWriter", + "_handle_spend_update_failure", + "handle_spend_update_failure", False, ), ( - 'litellm.proxy.db.db_transaction_queue.redis_update_buffer', - 'RedisUpdateBuffer', - '_should_commit_spend_updates_to_redis', - 'should_commit_spend_updates_to_redis', + "litellm.proxy.db.db_transaction_queue.redis_update_buffer", + "RedisUpdateBuffer", + "_should_commit_spend_updates_to_redis", + "should_commit_spend_updates_to_redis", False, ), ( - 'litellm.proxy.db.log_db_metrics', - '', - '_is_exception_related_to_db', - 'is_exception_related_to_db', + "litellm.proxy.db.log_db_metrics", + "", + "_is_exception_related_to_db", + "is_exception_related_to_db", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.azure.base', - '', - '_RESPONSES_API_CALL_TYPES', - 'RESPONSES_API_CALL_TYPES', + "litellm.proxy.guardrails.guardrail_hooks.azure.base", + "", + "_RESPONSES_API_CALL_TYPES", + "RESPONSES_API_CALL_TYPES", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.cisco_ai_defense.cisco_ai_defense_mcp', - '', - '_CiscoAIDefenseMcpMixin', - 'CiscoAIDefenseMcpMixin', + "litellm.proxy.guardrails.guardrail_hooks.cisco_ai_defense.cisco_ai_defense_mcp", + "", + "_CiscoAIDefenseMcpMixin", + "CiscoAIDefenseMcpMixin", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base', - '', - '_compile_marker', - 'compile_marker', + "litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base", + "", + "_compile_marker", + "compile_marker", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base', - '', - '_count_signals', - 'count_signals', + "litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base", + "", + "_count_signals", + "count_signals", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base', - '', - '_word_boundary_match', - 'word_boundary_match', + "litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base", + "", + "_word_boundary_match", + "word_boundary_match", False, ), ( - 'litellm.proxy.guardrails.guardrail_hooks.presidio', - '', - '_OPTIONAL_PresidioPIIMasking', - 'OPTIONAL_PresidioPIIMasking', + "litellm.proxy.guardrails.guardrail_hooks.presidio", + "", + "_OPTIONAL_PresidioPIIMasking", + "OPTIONAL_PresidioPIIMasking", False, ), ( - 'litellm.proxy.health_check', - '', - '_clean_endpoint_data', - 'clean_endpoint_data', + "litellm.proxy.health_check", + "", + "_clean_endpoint_data", + "clean_endpoint_data", False, ), ( - 'litellm.proxy.health_check', - '', - '_update_litellm_params_for_health_check', - 'update_litellm_params_for_health_check', + "litellm.proxy.health_check", + "", + "_update_litellm_params_for_health_check", + "update_litellm_params_for_health_check", False, ), ( - 'litellm.proxy.health_endpoints._health_endpoints', - '', - '_convert_health_check_to_dict', - 'convert_health_check_to_dict', + "litellm.proxy.health_endpoints._health_endpoints", + "", + "_convert_health_check_to_dict", + "convert_health_check_to_dict", False, ), ( - 'litellm.proxy.health_endpoints._health_endpoints', - '', - '_save_background_health_checks_to_db', - 'save_background_health_checks_to_db', + "litellm.proxy.health_endpoints._health_endpoints", + "", + "_save_background_health_checks_to_db", + "save_background_health_checks_to_db", False, ), ( - 'litellm.proxy.hooks.azure_content_safety', - '', - '_PROXY_AzureContentSafety', - 'PROXY_AzureContentSafety', + "litellm.proxy.hooks.azure_content_safety", + "", + "_PROXY_AzureContentSafety", + "PROXY_AzureContentSafety", False, ), ( - 'litellm.proxy.hooks.batch_rate_limiter', - '', - '_PROXY_BatchRateLimiter', - 'PROXY_BatchRateLimiter', + "litellm.proxy.hooks.batch_rate_limiter", + "", + "_PROXY_BatchRateLimiter", + "PROXY_BatchRateLimiter", False, ), ( - 'litellm.proxy.hooks.batch_redis_get', - '', - '_PROXY_BatchRedisRequests', - 'PROXY_BatchRedisRequests', + "litellm.proxy.hooks.batch_redis_get", + "", + "_PROXY_BatchRedisRequests", + "PROXY_BatchRedisRequests", False, ), ( - 'litellm.proxy.hooks.cache_control_check', - '', - '_PROXY_CacheControlCheck', - 'PROXY_CacheControlCheck', + "litellm.proxy.hooks.cache_control_check", + "", + "_PROXY_CacheControlCheck", + "PROXY_CacheControlCheck", False, ), ( - 'litellm.proxy.hooks.dynamic_rate_limiter', - '', - '_PROXY_DynamicRateLimitHandler', - 'PROXY_DynamicRateLimitHandler', + "litellm.proxy.hooks.dynamic_rate_limiter", + "", + "_PROXY_DynamicRateLimitHandler", + "PROXY_DynamicRateLimitHandler", False, ), ( - 'litellm.proxy.hooks.dynamic_rate_limiter_v3', - '', - '_PROXY_DynamicRateLimitHandlerV3', - 'PROXY_DynamicRateLimitHandlerV3', + "litellm.proxy.hooks.dynamic_rate_limiter_v3", + "", + "_PROXY_DynamicRateLimitHandlerV3", + "PROXY_DynamicRateLimitHandlerV3", False, ), ( - 'litellm.proxy.hooks.max_budget_per_session_limiter', - '', - '_PROXY_MaxBudgetPerSessionHandler', - 'PROXY_MaxBudgetPerSessionHandler', + "litellm.proxy.hooks.max_budget_per_session_limiter", + "", + "_PROXY_MaxBudgetPerSessionHandler", + "PROXY_MaxBudgetPerSessionHandler", False, ), ( - 'litellm.proxy.hooks.max_iterations_limiter', - '', - '_PROXY_MaxIterationsHandler', - 'PROXY_MaxIterationsHandler', + "litellm.proxy.hooks.max_iterations_limiter", + "", + "_PROXY_MaxIterationsHandler", + "PROXY_MaxIterationsHandler", False, ), ( - 'litellm.proxy.hooks.model_max_budget_limiter', - '', - '_PROXY_VirtualKeyModelMaxBudgetLimiter', - 'PROXY_VirtualKeyModelMaxBudgetLimiter', + "litellm.proxy.hooks.model_max_budget_limiter", + "", + "_PROXY_VirtualKeyModelMaxBudgetLimiter", + "PROXY_VirtualKeyModelMaxBudgetLimiter", False, ), ( - 'litellm.proxy.hooks.parallel_request_limiter', - '', - '_PROXY_MaxParallelRequestsHandler', - 'PROXY_MaxParallelRequestsHandler', + "litellm.proxy.hooks.parallel_request_limiter", + "", + "_PROXY_MaxParallelRequestsHandler", + "PROXY_MaxParallelRequestsHandler", False, ), ( - 'litellm.proxy.hooks.parallel_request_limiter_v3', - '', - '_PROXY_MaxParallelRequestsHandler_v3', - 'PROXY_MaxParallelRequestsHandler_v3', + "litellm.proxy.hooks.parallel_request_limiter_v3", + "", + "_PROXY_MaxParallelRequestsHandler_v3", + "PROXY_MaxParallelRequestsHandler_v3", False, ), ( - 'litellm.proxy.hooks.parallel_request_limiter_v3', - '_PROXY_MaxParallelRequestsHandler_v3', - '_create_rate_limit_descriptors', - 'create_rate_limit_descriptors', + "litellm.proxy.hooks.parallel_request_limiter_v3", + "_PROXY_MaxParallelRequestsHandler_v3", + "_create_rate_limit_descriptors", + "create_rate_limit_descriptors", False, ), ( - 'litellm.proxy.hooks.prompt_injection_detection', - '', - '_OPTIONAL_PromptInjectionDetection', - 'OPTIONAL_PromptInjectionDetection', + "litellm.proxy.hooks.prompt_injection_detection", + "", + "_OPTIONAL_PromptInjectionDetection", + "OPTIONAL_PromptInjectionDetection", False, ), ( - 'litellm.proxy.hooks.proxy_track_cost_callback', - '', - '_ProxyDBLogger', - 'ProxyDBLogger', + "litellm.proxy.hooks.proxy_track_cost_callback", + "", + "_ProxyDBLogger", + "ProxyDBLogger", False, ), ( - 'litellm.proxy.hooks.sensitive_data_routing', - '', - '_PROXY_SensitiveDataRoutingHandler', - 'PROXY_SensitiveDataRoutingHandler', + "litellm.proxy.hooks.sensitive_data_routing", + "", + "_PROXY_SensitiveDataRoutingHandler", + "PROXY_SensitiveDataRoutingHandler", False, ), ( - 'litellm.proxy.litellm_pre_call_utils', - '', - '_add_guardrails_from_key_or_team_metadata', - 'add_guardrails_from_key_or_team_metadata', + "litellm.proxy.litellm_pre_call_utils", + "", + "_add_guardrails_from_key_or_team_metadata", + "add_guardrails_from_key_or_team_metadata", False, ), ( - 'litellm.proxy.litellm_pre_call_utils', - '', - '_get_dynamic_logging_metadata', - 'get_dynamic_logging_metadata', + "litellm.proxy.litellm_pre_call_utils", + "", + "_get_dynamic_logging_metadata", + "get_dynamic_logging_metadata", False, ), ( - 'litellm.proxy.litellm_pre_call_utils', - '', - '_get_metadata_variable_name', - 'get_metadata_variable_name', + "litellm.proxy.litellm_pre_call_utils", + "", + "_get_metadata_variable_name", + "get_metadata_variable_name", False, ), ( - 'litellm.proxy.litellm_pre_call_utils', - '', - '_get_validated_callback_metadata', - 'get_validated_callback_metadata', + "litellm.proxy.litellm_pre_call_utils", + "", + "_get_validated_callback_metadata", + "get_validated_callback_metadata", False, ), ( - 'litellm.proxy.litellm_pre_call_utils', - 'LiteLLMProxyRequestSetup', - '_merge_tags', - 'merge_tags', + "litellm.proxy.litellm_pre_call_utils", + "LiteLLMProxyRequestSetup", + "_merge_tags", + "merge_tags", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_check_disable_global_guardrails_caller_permission', - 'check_disable_global_guardrails_caller_permission', + "litellm.proxy.management_endpoints.common_utils", + "", + "_check_disable_global_guardrails_caller_permission", + "check_disable_global_guardrails_caller_permission", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_check_passthrough_routes_caller_permission', - 'check_passthrough_routes_caller_permission', + "litellm.proxy.management_endpoints.common_utils", + "", + "_check_passthrough_routes_caller_permission", + "check_passthrough_routes_caller_permission", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_set_object_metadata_field', - 'set_object_metadata_field', + "litellm.proxy.management_endpoints.common_utils", + "", + "_set_object_metadata_field", + "set_object_metadata_field", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_team_member_has_permission', - 'team_member_has_permission', + "litellm.proxy.management_endpoints.common_utils", + "", + "_team_member_has_permission", + "team_member_has_permission", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_update_metadata_fields', - 'update_metadata_fields', + "litellm.proxy.management_endpoints.common_utils", + "", + "_update_metadata_fields", + "update_metadata_fields", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_upsert_budget_and_membership', - 'upsert_budget_and_membership', + "litellm.proxy.management_endpoints.common_utils", + "", + "_upsert_budget_and_membership", + "upsert_budget_and_membership", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_user_has_admin_privileges', - 'user_has_admin_privileges', + "litellm.proxy.management_endpoints.common_utils", + "", + "_user_has_admin_privileges", + "user_has_admin_privileges", False, ), ( - 'litellm.proxy.management_endpoints.common_utils', - '', - '_user_has_admin_view', - 'user_api_key_has_admin_view', + "litellm.proxy.management_endpoints.common_utils", + "", + "_user_has_admin_view", + "user_api_key_has_admin_view", False, ), ( - 'litellm.proxy.management_endpoints.config_override_endpoints', - '', - '_clear_hashicorp_vault_state', - 'clear_hashicorp_vault_state', + "litellm.proxy.management_endpoints.config_override_endpoints", + "", + "_clear_hashicorp_vault_state", + "clear_hashicorp_vault_state", False, ), ( - 'litellm.proxy.management_endpoints.config_override_endpoints', - '', - '_get_current_env_values', - 'get_current_env_values', + "litellm.proxy.management_endpoints.config_override_endpoints", + "", + "_get_current_env_values", + "get_current_env_values", False, ), ( - 'litellm.proxy.management_endpoints.config_override_endpoints', - '', - '_parse_config_value', - 'parse_config_value', + "litellm.proxy.management_endpoints.config_override_endpoints", + "", + "_parse_config_value", + "parse_config_value", False, ), ( - 'litellm.proxy.management_endpoints.config_override_endpoints', - '', - '_set_env_vars', - 'set_env_vars', + "litellm.proxy.management_endpoints.config_override_endpoints", + "", + "_set_env_vars", + "set_env_vars", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_calculate_key_rotation_time', - 'calculate_key_rotation_time', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_calculate_key_rotation_time", + "calculate_key_rotation_time", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_check_permissions_caller_permission', - 'check_permissions_caller_permission', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_check_permissions_caller_permission", + "check_permissions_caller_permission", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_get_caller_team_role', - 'get_caller_team_role', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_get_caller_team_role", + "get_caller_team_role", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_list_key_helper', - 'list_key_helper', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_list_key_helper", + "list_key_helper", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_persist_deleted_verification_tokens', - 'persist_deleted_verification_tokens', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_persist_deleted_verification_tokens", + "persist_deleted_verification_tokens", False, ), ( - 'litellm.proxy.management_endpoints.key_management_endpoints', - '', - '_rotate_master_key', - 'rotate_master_key', + "litellm.proxy.management_endpoints.key_management_endpoints", + "", + "_rotate_master_key", + "rotate_master_key", False, ), ( - 'litellm.proxy.management_endpoints.mcp_management_endpoints', - '', - '_inherit_credentials_from_existing_server', - 'inherit_credentials_from_existing_server', + "litellm.proxy.management_endpoints.mcp_management_endpoints", + "", + "_inherit_credentials_from_existing_server", + "inherit_credentials_from_existing_server", False, ), ( - 'litellm.proxy.management_endpoints.model_management_endpoints', - '', - '_add_model_to_db', - 'add_model_to_db', + "litellm.proxy.management_endpoints.model_management_endpoints", + "", + "_add_model_to_db", + "add_model_to_db", False, ), ( - 'litellm.proxy.management_endpoints.model_management_endpoints', - '', - '_add_team_model_to_db', - 'add_team_model_to_db', + "litellm.proxy.management_endpoints.model_management_endpoints", + "", + "_add_team_model_to_db", + "add_team_model_to_db", False, ), ( - 'litellm.proxy.management_endpoints.model_management_endpoints', - '', - '_deduplicate_litellm_router_models', - 'deduplicate_litellm_router_models', + "litellm.proxy.management_endpoints.model_management_endpoints", + "", + "_deduplicate_litellm_router_models", + "deduplicate_litellm_router_models", False, ), ( - 'litellm.proxy.management_endpoints.organization_endpoints', - '', - '_verify_org_access', - 'verify_org_access', + "litellm.proxy.management_endpoints.organization_endpoints", + "", + "_verify_org_access", + "verify_org_access", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_all_names_per_competitor', - 'build_all_names_per_competitor', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_all_names_per_competitor", + "build_all_names_per_competitor", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_comparison_blocked_words', - 'build_comparison_blocked_words', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_comparison_blocked_words", + "build_comparison_blocked_words", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_competitor_guardrail_definitions', - 'build_competitor_guardrail_definitions', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_competitor_guardrail_definitions", + "build_competitor_guardrail_definitions", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_name_blocked_words', - 'build_name_blocked_words', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_name_blocked_words", + "build_name_blocked_words", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_recommendation_blocked_words', - 'build_recommendation_blocked_words', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_recommendation_blocked_words", + "build_recommendation_blocked_words", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_build_refinement_prompt', - 'build_refinement_prompt', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_build_refinement_prompt", + "build_refinement_prompt", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_clean_competitor_line', - 'clean_competitor_line', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_clean_competitor_line", + "clean_competitor_line", False, ), ( - 'litellm.proxy.management_endpoints.policy_endpoints.endpoints', - '', - '_parse_variations_response', - 'parse_variations_response', + "litellm.proxy.management_endpoints.policy_endpoints.endpoints", + "", + "_parse_variations_response", + "parse_variations_response", False, ), ( - 'litellm.proxy.management_endpoints.team_endpoints', - '', - '_cleanup_members_with_roles', - 'cleanup_members_with_roles', + "litellm.proxy.management_endpoints.team_endpoints", + "", + "_cleanup_members_with_roles", + "cleanup_members_with_roles", False, ), ( - 'litellm.proxy.management_endpoints.team_endpoints', - '', - '_refresh_cached_team', - 'refresh_cached_team', + "litellm.proxy.management_endpoints.team_endpoints", + "", + "_refresh_cached_team", + "refresh_cached_team", False, ), ( - 'litellm.proxy.management_endpoints.team_endpoints', - 'TeamMemberBudgetHandler', - '_clean_team_member_fields', - 'clean_team_member_fields', + "litellm.proxy.management_endpoints.team_endpoints", + "TeamMemberBudgetHandler", + "_clean_team_member_fields", + "clean_team_member_fields", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - '', - '_sso_return_to_redirect', - 'sso_return_to_redirect', + "litellm.proxy.management_endpoints.ui_sso", + "", + "_sso_return_to_redirect", + "sso_return_to_redirect", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - 'SSOAuthenticationHandler', - '_delete_pkce_verifier', - 'delete_pkce_verifier', + "litellm.proxy.management_endpoints.ui_sso", + "SSOAuthenticationHandler", + "_delete_pkce_verifier", + "delete_pkce_verifier", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - 'SSOAuthenticationHandler', - '_get_cli_state', - 'get_cli_state', + "litellm.proxy.management_endpoints.ui_sso", + "SSOAuthenticationHandler", + "_get_cli_state", + "get_cli_state", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - 'SSOAuthenticationHandler', - '_get_user_email_and_id_from_result', - 'get_user_email_and_id_from_result', + "litellm.proxy.management_endpoints.ui_sso", + "SSOAuthenticationHandler", + "_get_user_email_and_id_from_result", + "get_user_email_and_id_from_result", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - 'SSOAuthenticationHandler', - '_pkce_token_exchange', - 'pkce_token_exchange', + "litellm.proxy.management_endpoints.ui_sso", + "SSOAuthenticationHandler", + "_pkce_token_exchange", + "pkce_token_exchange", False, ), ( - 'litellm.proxy.management_endpoints.ui_sso', - 'SSOAuthenticationHandler', - '_validate_return_to', - 'validate_return_to', + "litellm.proxy.management_endpoints.ui_sso", + "SSOAuthenticationHandler", + "_validate_return_to", + "validate_return_to", False, ), ( - 'litellm.proxy.management_helpers.object_permission_utils', - '', - '_get_allow_all_keys_server_ids', - 'get_allow_all_keys_server_ids', + "litellm.proxy.management_helpers.object_permission_utils", + "", + "_get_allow_all_keys_server_ids", + "get_allow_all_keys_server_ids", False, ), ( - 'litellm.proxy.management_helpers.object_permission_utils', - '', - '_get_team_allowed_mcp_servers', - 'get_team_allowed_mcp_servers', + "litellm.proxy.management_helpers.object_permission_utils", + "", + "_get_team_allowed_mcp_servers", + "get_team_allowed_mcp_servers", False, ), ( - 'litellm.proxy.management_helpers.object_permission_utils', - '', - '_set_object_permission', - 'set_object_permission', + "litellm.proxy.management_helpers.object_permission_utils", + "", + "_set_object_permission", + "set_object_permission", False, ), ( - 'litellm.proxy.openai_files_endpoints.common_utils', - '', - '_is_base64_encoded_unified_file_id', - 'is_base64_encoded_unified_file_id', + "litellm.proxy.openai_files_endpoints.common_utils", + "", + "_is_base64_encoded_unified_file_id", + "is_base64_encoded_unified_file_id", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints', - '', - '_extract_model_from_bedrock_endpoint', - 'extract_model_from_bedrock_endpoint', + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints", + "", + "_extract_model_from_bedrock_endpoint", + "extract_model_from_bedrock_endpoint", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints', - 'BaseOpenAIPassThroughHandler', - '_base_openai_pass_through_handler', - 'base_openai_pass_through_handler', + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints", + "BaseOpenAIPassThroughHandler", + "_base_openai_pass_through_handler", + "base_openai_pass_through_handler", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler', - 'AnthropicPassthroughLoggingHandler', - '_build_complete_streaming_response', - 'build_complete_streaming_response', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler", + "AnthropicPassthroughLoggingHandler", + "_build_complete_streaming_response", + "build_complete_streaming_response", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler', - 'AnthropicPassthroughLoggingHandler', - '_build_usage_only_response_from_chunks', - 'build_usage_only_response_from_chunks', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler", + "AnthropicPassthroughLoggingHandler", + "_build_usage_only_response_from_chunks", + "build_usage_only_response_from_chunks", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler', - 'AnthropicPassthroughLoggingHandler', - '_handle_logging_anthropic_collected_chunks', - 'handle_logging_anthropic_collected_chunks', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler", + "AnthropicPassthroughLoggingHandler", + "_handle_logging_anthropic_collected_chunks", + "handle_logging_anthropic_collected_chunks", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler', - 'AssemblyAIPassthroughLoggingHandler', - '_get_assembly_base_url_from_region', - 'get_assembly_base_url_from_region', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler", + "AssemblyAIPassthroughLoggingHandler", + "_get_assembly_base_url_from_region", + "get_assembly_base_url_from_region", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler', - 'AssemblyAIPassthroughLoggingHandler', - '_get_assembly_region_from_url', - 'get_assembly_region_from_url', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler", + "AssemblyAIPassthroughLoggingHandler", + "_get_assembly_region_from_url", + "get_assembly_region_from_url", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler', - 'AssemblyAIPassthroughLoggingHandler', - '_should_log_request', - 'should_log_request', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.assembly_passthrough_logging_handler", + "AssemblyAIPassthroughLoggingHandler", + "_should_log_request", + "should_log_request", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler', - '', - '_is_openai_compatible_url', - 'is_openai_compatible_url', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler", + "", + "_is_openai_compatible_url", + "is_openai_compatible_url", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler', - 'OpenAIPassthroughLoggingHandler', - '_handle_logging_openai_collected_chunks', - 'handle_logging_openai_collected_chunks', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler", + "OpenAIPassthroughLoggingHandler", + "_handle_logging_openai_collected_chunks", + "handle_logging_openai_collected_chunks", False, ), ( - 'litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler', - 'VertexPassthroughLoggingHandler', - '_handle_logging_vertex_collected_chunks', - 'handle_logging_vertex_collected_chunks', + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler", + "VertexPassthroughLoggingHandler", + "_handle_logging_vertex_collected_chunks", + "handle_logging_vertex_collected_chunks", False, ), ( - 'litellm.proxy.pass_through_endpoints.pass_through_endpoints', - 'HttpPassThroughEndpointHelpers', - '_init_kwargs_for_pass_through_endpoint', - 'init_kwargs_for_pass_through_endpoint', + "litellm.proxy.pass_through_endpoints.pass_through_endpoints", + "HttpPassThroughEndpointHelpers", + "_init_kwargs_for_pass_through_endpoint", + "init_kwargs_for_pass_through_endpoint", False, ), ( - 'litellm.proxy.pass_through_endpoints.pass_through_endpoints', - 'HttpPassThroughEndpointHelpers', - '_update_stream_param_based_on_request_body', - 'update_stream_param_based_on_request_body', + "litellm.proxy.pass_through_endpoints.pass_through_endpoints", + "HttpPassThroughEndpointHelpers", + "_update_stream_param_based_on_request_body", + "update_stream_param_based_on_request_body", False, ), ( - 'litellm.proxy.pass_through_endpoints.streaming_handler', - 'PassThroughStreamingHandler', - '_route_streaming_logging_to_handler', - 'route_streaming_logging_to_handler', + "litellm.proxy.pass_through_endpoints.streaming_handler", + "PassThroughStreamingHandler", + "_route_streaming_logging_to_handler", + "route_streaming_logging_to_handler", False, ), ( - 'litellm.proxy.pass_through_endpoints.success_handler', - 'PassThroughEndpointLogging', - '_handle_logging', - 'handle_logging', + "litellm.proxy.pass_through_endpoints.success_handler", + "PassThroughEndpointLogging", + "_handle_logging", + "handle_logging", False, ), ( - 'litellm.proxy.policy_engine.policy_registry', - 'PolicyRegistry', - '_parse_policy', - 'parse_policy', + "litellm.proxy.policy_engine.policy_registry", + "PolicyRegistry", + "_parse_policy", + "parse_policy", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_apply_uvicorn_max_requests_jitter', - 'apply_uvicorn_max_requests_jitter', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_apply_uvicorn_max_requests_jitter", + "apply_uvicorn_max_requests_jitter", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_configure_dev_reload', - 'configure_dev_reload', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_configure_dev_reload", + "configure_dev_reload", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_echo_litellm_version', - 'echo_litellm_version', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_echo_litellm_version", + "echo_litellm_version", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_get_default_unvicorn_init_args', - 'get_default_unvicorn_init_args', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_get_default_unvicorn_init_args", + "get_default_unvicorn_init_args", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_get_loop_type', - 'get_loop_type', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_get_loop_type", + "get_loop_type", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_init_granian_server', - 'init_granian_server', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_init_granian_server", + "init_granian_server", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_init_hypercorn_server', - 'init_hypercorn_server', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_init_hypercorn_server", + "init_hypercorn_server", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_is_port_in_use', - 'is_port_in_use', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_is_port_in_use", + "is_port_in_use", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_maybe_setup_prometheus_multiproc_dir', - 'maybe_setup_prometheus_multiproc_dir', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_maybe_setup_prometheus_multiproc_dir", + "maybe_setup_prometheus_multiproc_dir", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_run_config_validation', - 'run_config_validation', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_run_config_validation", + "run_config_validation", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_run_gunicorn_server', - 'run_gunicorn_server', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_run_gunicorn_server", + "run_gunicorn_server", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_run_health_check', - 'run_health_check', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_run_health_check", + "run_health_check", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_run_ollama_serve', - 'run_ollama_serve', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_run_ollama_serve", + "run_ollama_serve", False, ), ( - 'litellm.proxy.proxy_cli', - 'ProxyInitializationHelpers', - '_run_test_chat_completion', - 'run_test_chat_completion', + "litellm.proxy.proxy_cli", + "ProxyInitializationHelpers", + "_run_test_chat_completion", + "run_test_chat_completion", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_build_redis_usage_cache', - 'build_redis_usage_cache', + "litellm.proxy.proxy_server", + "", + "_build_redis_usage_cache", + "build_redis_usage_cache", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_ensure_spend_counter_initialized', - 'ensure_spend_counter_initialized', + "litellm.proxy.proxy_server", + "", + "_ensure_spend_counter_initialized", + "ensure_spend_counter_initialized", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_ensure_window_spend_counter_initialized', - 'ensure_window_spend_counter_initialized', + "litellm.proxy.proxy_server", + "", + "_ensure_window_spend_counter_initialized", + "ensure_window_spend_counter_initialized", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_environment_has_redis_connection_target', - 'environment_has_redis_connection_target', + "litellm.proxy.proxy_server", + "", + "_environment_has_redis_connection_target", + "environment_has_redis_connection_target", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_get_model_group_info', - 'get_model_group_info', + "litellm.proxy.proxy_server", + "", + "_get_model_group_info", + "get_model_group_info", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_increment_spend_counter_cache', - 'increment_spend_counter_cache', + "litellm.proxy.proxy_server", + "", + "_increment_spend_counter_cache", + "increment_spend_counter_cache", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_initialize_shared_aiohttp_session', - 'initialize_shared_aiohttp_session', + "litellm.proxy.proxy_server", + "", + "_initialize_shared_aiohttp_session", + "initialize_shared_aiohttp_session", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_invalidate_spend_counter', - 'invalidate_spend_counter', + "litellm.proxy.proxy_server", + "", + "_invalidate_spend_counter", + "invalidate_spend_counter", False, ), ( - 'litellm.proxy.proxy_server', - '', - '_title', - 'title', + "litellm.proxy.proxy_server", + "", + "_title", + "title", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_add_deployment_locked', - 'add_deployment_locked', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_add_deployment_locked", + "add_deployment_locked", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_decrypt_and_set_db_env_variables', - 'decrypt_and_set_db_env_variables', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_decrypt_and_set_db_env_variables", + "decrypt_and_set_db_env_variables", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_decrypt_db_variables', - 'decrypt_db_variables', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_decrypt_db_variables", + "decrypt_db_variables", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_encrypt_env_variables', - 'encrypt_env_variables', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_encrypt_env_variables", + "encrypt_env_variables", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_encrypt_env_variables_for_db', - 'encrypt_env_variables_for_db', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_encrypt_env_variables_for_db", + "encrypt_env_variables_for_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_get_models_from_db', - 'get_models_from_db', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_get_models_from_db", + "get_models_from_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_init_cache', - 'init_cache', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_init_cache", + "init_cache", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_init_semantic_filter_settings_in_db', - 'init_semantic_filter_settings_in_db', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_init_semantic_filter_settings_in_db", + "init_semantic_filter_settings_in_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyConfig', - '_serve_pass_through_endpoints', - 'serve_pass_through_endpoints', + "litellm.proxy.proxy_server", + "ProxyConfig", + "_serve_pass_through_endpoints", + "serve_pass_through_endpoints", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_add_proxy_budget_to_db', - 'add_proxy_budget_to_db', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_add_proxy_budget_to_db", + "add_proxy_budget_to_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_attach_router_to_prompt_injection_detectors', - 'attach_router_to_prompt_injection_detectors', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_attach_router_to_prompt_injection_detectors", + "attach_router_to_prompt_injection_detectors", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_get_transaction_buffer_redis_cache', - 'get_transaction_buffer_redis_cache', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_get_transaction_buffer_redis_cache", + "get_transaction_buffer_redis_cache", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_init_coordination_redis_from_db', - 'init_coordination_redis_from_db', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_init_coordination_redis_from_db", + "init_coordination_redis_from_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_init_dd_tracer', - 'init_dd_tracer', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_init_dd_tracer", + "init_dd_tracer", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_init_pyroscope', - 'init_pyroscope', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_init_pyroscope", + "init_pyroscope", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_initialize_jwt_auth', - 'initialize_jwt_auth', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_initialize_jwt_auth", + "initialize_jwt_auth", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_initialize_semantic_tool_filter', - 'initialize_semantic_tool_filter', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_initialize_semantic_tool_filter", + "initialize_semantic_tool_filter", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_initialize_startup_logging', - 'initialize_startup_logging', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_initialize_startup_logging", + "initialize_startup_logging", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_setup_prisma_client', - 'setup_prisma_client', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_setup_prisma_client", + "setup_prisma_client", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_sync_ui_settings_to_general_settings', - 'sync_ui_settings_to_general_settings', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_sync_ui_settings_to_general_settings", + "sync_ui_settings_to_general_settings", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_update_default_team_member_budget', - 'update_default_team_member_budget', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_update_default_team_member_budget", + "update_default_team_member_budget", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_validate_redis_transaction_buffer_config', - 'validate_redis_transaction_buffer_config', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_validate_redis_transaction_buffer_config", + "validate_redis_transaction_buffer_config", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_warm_global_spend_cache', - 'warm_global_spend_cache', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_warm_global_spend_cache", + "warm_global_spend_cache", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_warn_budget_without_db', - 'warn_budget_without_db', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_warn_budget_without_db", + "warn_budget_without_db", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_warn_fail_closed_rate_limits_without_redis', - 'warn_fail_closed_rate_limits_without_redis', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_warn_fail_closed_rate_limits_without_redis", + "warn_fail_closed_rate_limits_without_redis", False, ), ( - 'litellm.proxy.proxy_server', - 'ProxyStartupEvent', - '_warn_if_mock_testing_params_enabled', - 'warn_if_mock_testing_params_enabled', + "litellm.proxy.proxy_server", + "ProxyStartupEvent", + "_warn_if_mock_testing_params_enabled", + "warn_if_mock_testing_params_enabled", False, ), ( - 'litellm.proxy.spend_tracking.spend_management_endpoints', - '', - '_get_spend_report_for_time_range', - 'get_spend_report_for_time_range', + "litellm.proxy.spend_tracking.spend_management_endpoints", + "", + "_get_spend_report_for_time_range", + "get_spend_report_for_time_range", False, ), ( - 'litellm.proxy.spend_tracking.spend_management_endpoints', - '', - '_is_admin_view_safe', - 'is_admin_view_safe', + "litellm.proxy.spend_tracking.spend_management_endpoints", + "", + "_is_admin_view_safe", + "is_admin_view_safe", False, ), ( - 'litellm.proxy.spend_tracking.spend_tracking_utils', - '', - '_is_master_key', - 'is_master_key', + "litellm.proxy.spend_tracking.spend_tracking_utils", + "", + "_is_master_key", + "is_master_key", False, ), ( - 'litellm.proxy.spend_tracking.spend_tracking_utils', - '', - '_sanitize_error_information_for_spend_logs', - 'sanitize_error_information_for_spend_logs', + "litellm.proxy.spend_tracking.spend_tracking_utils", + "", + "_sanitize_error_information_for_spend_logs", + "sanitize_error_information_for_spend_logs", False, ), ( - 'litellm.proxy.utils', - '', - '_cache_user_row', - 'cache_user_row', + "litellm.proxy.utils", + "", + "_cache_user_row", + "cache_user_row", False, ), ( - 'litellm.proxy.utils', - '', - '_check_and_merge_model_level_guardrails', - 'check_and_merge_model_level_guardrails', + "litellm.proxy.utils", + "", + "_check_and_merge_model_level_guardrails", + "check_and_merge_model_level_guardrails", False, ), ( - 'litellm.proxy.utils', - '', - '_get_docs_url', - 'get_docs_url', + "litellm.proxy.utils", + "", + "_get_docs_url", + "get_docs_url", False, ), ( - 'litellm.proxy.utils', - '', - '_get_openapi_url', - 'get_openapi_url', + "litellm.proxy.utils", + "", + "_get_openapi_url", + "get_openapi_url", False, ), ( - 'litellm.proxy.utils', - '', - '_get_projected_spend_over_limit', - 'get_projected_spend_over_limit', + "litellm.proxy.utils", + "", + "_get_projected_spend_over_limit", + "get_projected_spend_over_limit", False, ), ( - 'litellm.proxy.utils', - '', - '_get_redoc_url', - 'get_redoc_url', + "litellm.proxy.utils", + "", + "_get_redoc_url", + "get_redoc_url", False, ), ( - 'litellm.proxy.utils', - '', - '_hash_token_if_needed', - 'hash_token_if_needed', + "litellm.proxy.utils", + "", + "_hash_token_if_needed", + "hash_token_if_needed", False, ), ( - 'litellm.proxy.utils', - '', - '_is_projected_spend_over_limit', - 'is_projected_spend_over_limit', + "litellm.proxy.utils", + "", + "_is_projected_spend_over_limit", + "is_projected_spend_over_limit", False, ), ( - 'litellm.proxy.utils', - '', - '_is_valid_team_configs', - 'is_valid_team_configs', + "litellm.proxy.utils", + "", + "_is_valid_team_configs", + "is_valid_team_configs", False, ), ( - 'litellm.proxy.utils', - '', - '_monitor_spend_logs_queue', - 'monitor_spend_logs_queue', + "litellm.proxy.utils", + "", + "_monitor_spend_logs_queue", + "monitor_spend_logs_queue", False, ), ( - 'litellm.proxy.utils', - '', - '_premium_user_check', - 'premium_user_check', + "litellm.proxy.utils", + "", + "_premium_user_check", + "premium_user_check", False, ), ( - 'litellm.proxy.utils', - '', - '_raise_failed_update_spend_exception', - 'raise_failed_update_spend_exception', + "litellm.proxy.utils", + "", + "_raise_failed_update_spend_exception", + "raise_failed_update_spend_exception", False, ), ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_arelease_max_parallel_requests_on_disconnect', - 'arelease_max_parallel_requests_on_disconnect', + "litellm.proxy.utils", + "ProxyLogging", + "_arelease_max_parallel_requests_on_disconnect", + "arelease_max_parallel_requests_on_disconnect", False, ), ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_callback_capabilities', - 'callback_capabilities', + "litellm.proxy.utils", + "ProxyLogging", + "_callback_capabilities", + "callback_capabilities", False, ), ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_convert_mcp_hook_response_to_kwargs', - 'convert_mcp_hook_response_to_kwargs', + "litellm.proxy.utils", + "ProxyLogging", + "_convert_mcp_hook_response_to_kwargs", + "convert_mcp_hook_response_to_kwargs", False, ), ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_create_mcp_request_object_from_kwargs', - 'create_mcp_request_object_from_kwargs', + "litellm.proxy.utils", + "ProxyLogging", + "_create_mcp_request_object_from_kwargs", + "create_mcp_request_object_from_kwargs", False, ), ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_fire_deferred_stream_logging', - 'fire_deferred_stream_logging', + "litellm.proxy.utils", + "ProxyLogging", + "_fire_deferred_stream_logging", + "fire_deferred_stream_logging", False, ), ) @@ -5188,17 +5201,16 @@ LLMS_FORWARDER_CASES: Final = ( ) PROXY_FORWARDER_CASES: Final = ( ( - 'litellm.proxy.utils', - 'ProxyLogging', - '_convert_mcp_to_llm_format', - 'convert_mcp_to_llm_format', - 'instance', + "litellm.proxy.utils", + "ProxyLogging", + "_convert_mcp_to_llm_format", + "convert_mcp_to_llm_format", + "instance", False, ), ) - def _get_owner(module_path: str, owner_name: str) -> object: module: Final = import_module(module_path or "litellm") if module_path == "litellm.llms.openai_like.dynamic_config" and owner_name == "JSONProviderConfig": @@ -5292,14 +5304,10 @@ def _make_private_override( def recorded_args(args: tuple[object, ...]) -> tuple[object, ...]: return args[1:] if descriptor in ("instance", "classmethod") else args - def sync_override( - *args: object, **kwargs: object - ) -> tuple[object, tuple[object, ...], dict[str, object]]: + def sync_override(*args: object, **kwargs: object) -> tuple[object, tuple[object, ...], dict[str, object]]: return result, recorded_args(args), kwargs - async def async_override( - *args: object, **kwargs: object - ) -> tuple[object, tuple[object, ...], dict[str, object]]: + async def async_override(*args: object, **kwargs: object) -> tuple[object, tuple[object, ...], dict[str, object]]: return result, recorded_args(args), kwargs implementation: Final = async_override if is_async else sync_override diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index ba928bfb8c3..273bebd98e5 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -20104,11 +20104,13 @@ def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooli async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, surface: str, credential_kind: str ) -> None: + from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.router import AllowedFailsPolicy monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setattr("time.time", lambda: 1767225600.0) + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) def upstream(request: httpx.Request) -> httpx.Response: credential: Final = request.headers.get("authorization", request.headers.get("x-api-key", "")) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 8a5d164fdff..d2b5d2b911b 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -750,9 +750,11 @@ async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attem respx_mock: respx.MockRouter, ): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.responses.streaming_iterator import drain_logging_tasks monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "anthropic_messages", python_messages.anthropic_messages) litellm.in_memory_llm_clients_cache.flush_cache() provider: Final = respx_mock.post(provider_url).mock( side_effect=[