From b5a72c2042daff32bf769fd28f9e8de17d738a96 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:56:54 -0800 Subject: [PATCH] fix(router): pin /v1/responses requests for prompt caching affinity (#50321) * fix(router): pin /v1/responses requests for prompt caching affinity * fix(router): pin previous_response_id turns and keep native tools in the Responses pin key A Responses turn that continues a stored response now goes to the deployment that answered it, unless that deployment already failed this request. Native Responses tools are hashed into the pin key as sent instead of converted, so two requests with different native tools no longer share a pin and a native request no longer logs a false dropped-tool warning. The pin keeps its cacheable deployment ids in a sorted tuple so loggers can serialize request metadata. * fix(router): leave a Responses request with a malformed tool to the provider * fix(router): keep the Responses prompt cache pin out of logged metadata The pre-call check now keeps the pin in a router-owned in-flight store and puts only a random token in litellm_metadata, so loggers that ship metadata never see the routing key or deployment ids. The store is shared across check instances, so turning prompt_caching off and on at runtime keeps pinning, and a retry drops its previous attempt's token. The skipped-input debug log now names only the error type. * fix(router): apply the prompt caching pin after tag, region and pre-call filters, and leave it on retries The prompt caching check no longer narrows the candidates inside the callback filters. It names the deployment it prefers, and the router core applies that preference at the pinned deployment step, after tag, region, pre-call and plugin filters, only when the deployment is still a candidate. A pin can no longer send a request to a deployment its tags do not allow. A retry skips the pin when the pinned deployment was already tried in this request, read from the retry breadcrumbs in both metadata buckets and the excluded and retry-skipped ids. When the retry count outruns the breadcrumbs the check stands down, so a pinned deployment that keeps failing is not retried on every attempt. * test(router): list the prompt caching check as an in-memory cache caller The check's in-flight Responses pins live in a process-local InMemoryCache, so it opens no Redis span and needs no key family, like the other in-memory callers * fix(router): pin a lone healthy deployment, free every in-flight Responses pin, and clear a reused metadata dict's retry records * test(router): lift test cooldowns through the cooldown store's evict_local * fix(router): key Responses tools like their chat twin and drop the pin token after success A Responses function tool now hashes only the function fields it sent, so a chat request and a Responses request with the same tool share one pin. The success event also removes the pin token from the caller's metadata dict. * test(router): cover untyped tool parameters and custom tool grammars in the Responses pin key --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../crates/python-bridge/src/router/host.rs | 28 +- .../crates/python-bridge/src/router/mod.rs | 8 +- litellm-rust/crates/router-core/src/host.rs | 19 +- litellm-rust/crates/router-core/src/lib.rs | 4 +- litellm-rust/crates/router-core/src/select.rs | 19 +- .../transformation.py | 6 +- litellm/router_rust/invoke/walk.py | 7 +- litellm/router_rust/steps/callbacks.py | 44 +- .../prompt_caching_deployment_check.py | 380 ++++++- litellm/router_utils/prompt_caching_cache.py | 21 +- ...prompt_caching_affinity_tool_call_marks.py | 68 +- .../test_prompt_caching_deployment_check.py | 994 +++++++++++++++++- tests/unit/test_internal_context.py | 1 + 13 files changed, 1486 insertions(+), 113 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 6ccac4df659..0ded98f27c5 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -20,10 +20,10 @@ use std::{ use async_trait::async_trait; use litellm_router_core::{ - Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, - Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest, - RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, - build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, + Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, + RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, + ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, }; use pyo3::{PyResult, Python, types::PyAnyMethods}; @@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse { PreCall(PreCallFacts), /// The candidate models a `routing_plugins` select step reported. Models(Option>), - /// The ids a `callback_filters` select step left, or the error a filter raised. - Filtered(Result, AttemptError>), + /// The ids a `callback_filters` select step left with the one a filter prefers, or the error + /// a filter raised. + Filtered(Result<(Vec, Option), AttemptError>), /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of /// answering: the walk takes it as the try's failure. Raised(AttemptError), @@ -352,8 +353,10 @@ impl HostCallResponse { got: self.kind(), }); } - if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) = - (self, call) + if let ( + Self::Select(ids) | Self::Filtered(Ok((ids, _))), + HostCall::Select { candidates, .. }, + ) = (self, call) && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) { return Err(DriveError::UnknownCandidate { id: id.clone() }); @@ -890,7 +893,7 @@ impl Host for StepHost { &self, ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { + ) -> Result { let (scope, group) = ctx.scope.wire(); let joined = |own: &[String], walked: &[String]| [own, walked].concat(); let call = HostCall::Select { @@ -909,8 +912,11 @@ impl Host for StepHost { let HostCallResponse::Filtered(result) = call_host(call).await else { unreachable!("{CHECKED}"); }; - let ids = result.map_err(Error::Attempt)?; - Ok(answered_by_id(&candidates, &ids)) + let (ids, preferred) = result.map_err(Error::Attempt)?; + Ok(CallbackFiltered { + kept: answered_by_id(&candidates, &ids), + preferred, + }) } async fn allow_fallback( diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index 8157d0f5a32..6b58d63a714 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -1181,7 +1181,8 @@ fn host_call_response( stage: Stage::CallbackFilters { .. }, .. } => { - // The remaining ids, or a failure mapping as an attempt reports one. + // The remaining ids, alone or paired with the id a filter prefers, or a failure + // mapping as an attempt reports one. let outcome = outcome.ok_or_else(missing)?; if let Ok(failure) = outcome.cast::() { return match attempt_outcome(failure)? { @@ -1191,7 +1192,10 @@ fn host_call_response( )), }; } - Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) + if let Ok(answer) = outcome.cast::() { + return Ok(HostCallResponse::Filtered(Ok(answer.extract()?))); + } + Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None)))) } HostCall::Select { stage: Stage::Strategy { records, .. }, diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index 0ef37873e6e..ae93288b0fa 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -405,6 +405,12 @@ pub enum SelectStage { OverrideCheck(RoutingStrategy), } +#[derive(Debug)] +pub struct CallbackFiltered { + pub kept: Vec>, + pub preferred: Option, +} + /// Routing events, delivered through [`Host::observe`]. They serialize as JSON objects tagged by /// `event` (the variant name), durations in seconds. #[derive(Clone, Debug, PartialEq, serde::Serialize)] @@ -714,14 +720,19 @@ pub trait Host: Send + Sync { } /// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in - /// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and - /// falls back from like any attempt failure. A host without filters keeps every candidate. + /// turn, and may name one deployment they prefer, which the pick takes only when the filters + /// after the callbacks keep it. A filter that raises fails the try with [`Error::Attempt`], + /// which the call retries and falls back from like any attempt failure. A host without + /// filters keeps every candidate and prefers none. async fn callback_filters( &self, _ctx: &RoutingContext<'_>, candidates: Vec>, - ) -> Result>, Error> { - Ok(candidates) + ) -> Result { + Ok(CallbackFiltered { + kept: candidates, + preferred: None, + }) } /// The values of `refs`, one per ref, in order: what the deployment build substitutes for /// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 80c8f597157..f9b71c689c9 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -124,8 +124,8 @@ pub use crate::{ error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, health::{Clock, Health, Outcome, SystemClock}, host::{ - Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause, - RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, + Attempt, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Host, Output, Provided, + RetryCause, RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, }, outcome::{OutcomeCallbacks, OutcomeWrite}, prerouting::{PreRoute, PreRouted}, diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index c7f3b1f9c31..75198c0a707 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -456,7 +456,7 @@ impl Core { let request_order = request.target_order.filter(|_| !inputs.fallback_hop); // `async_callback_filter_deployments` runs after the cooldown and blocked filters on every // async candidate list but a deployment named by id and a listed pass-through pick. - if settings.callback_filters + let preferred = if settings.callback_filters && !inputs.count_only && !resolution.kind.is_single() && request.pass_through != Some(PassThrough::Listed) @@ -472,8 +472,12 @@ impl Core { ctx.excluded_ids = &inputs.excluded_ids; ctx.skipped_ids = &inputs.skipped_ids; ctx.target_order = inputs.target_order.or(request_order); - healthy = self.host.callback_filters(&ctx, healthy).await?; - } + let filtered = self.host.callback_filters(&ctx, healthy).await?; + healthy = filtered.kept; + filtered.preferred + } else { + None + }; if inputs.count_only { return Ok(Candidates { name: name.to_owned(), @@ -648,8 +652,13 @@ impl Core { healthy.retain(|d| !d.cannot_run_batches()); } - if let Some(pinned) = request.pinned_deployment_id.as_deref() - && healthy.iter().any(|d| d.id == pinned) + if let Some(pinned) = [ + request.pinned_deployment_id.as_deref(), + preferred.as_deref(), + ] + .into_iter() + .flatten() + .find(|id| healthy.iter().any(|d| d.id == *id)) { healthy.retain(|d| d.id == pinned); } diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 33ea85d7d53..b49bbd23917 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -199,6 +199,8 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool | AnthropicToolSearchToolRegex NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +RESPONSES_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) +RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: Final = frozenset({"computer_use", "image_generation", "local_shell", "shell"}) _INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) ) @@ -2948,7 +2950,7 @@ class LiteLLMCompletionResponsesConfig: tool_type: Final = tool.get("type") if tool_type == "mcp": return ResponsesToolChatForm(chat_tools=(cast(OpenAIMcpServerTool, tool),), web_search_options=None) - if tool_type == "web_search_preview" or tool_type == "web_search": + if tool_type in RESPONSES_WEB_SEARCH_TOOL_TYPES: _search_context_size: Final[Literal["low", "medium", "high"]] = cast( Literal["low", "medium", "high"], tool.get("search_context_size") ) @@ -3001,7 +3003,7 @@ class LiteLLMCompletionResponsesConfig: return ResponsesToolChatForm(chat_tools=(client_tool_search_chat_tool(tool),), web_search_options=None) if tool_type == "tool_search": return ResponsesToolChatForm(chat_tools=(anthropic_tool_search_chat_tool(),), web_search_options=None) - if tool_type in ("computer_use", "image_generation", "local_shell", "shell"): + if tool_type in RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM: verbose_logger.warning( "Dropping Responses API tool of type '%s': it has no Chat Completions " "equivalent and the target provider would reject the request.", diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 17e20119f43..ff8bc7eef5a 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -39,6 +39,7 @@ _WALK_KEYS: Final = ( _SERVED_KEYS: Final = ("model_info", "deployment") # The walk state `run_async_fallback` hands every hop as kwargs. _HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors") +_RETRY_RECORD_KEYS: Final = ("previous_models", "request_retry_count") def _reentrant_hop_stamps(request: Mapping[str, object], field: str) -> Mapping[str, object]: @@ -179,11 +180,13 @@ class WalkMetadata: return sibling: Final = request.get("metadata" if self._field == "litellm_metadata" else "litellm_metadata") if isinstance(sibling, dict): - sibling.pop("attempted_fallbacks", None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict - sibling.pop("original_model_group", None) # pyright: ignore[reportUnknownMemberType] # see above + for key in ("attempted_fallbacks", "original_model_group", *_RETRY_RECORD_KEYS): + sibling.pop(key, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict bucket: Final = request.get(self._field) if not isinstance(bucket, dict): return + for key in _RETRY_RECORD_KEYS: + bucket.pop(key, None) # pyright: ignore[reportUnknownMemberType] # see above bucket["attempted_fallbacks"] = 0 model: Final = request.get("model") if model is not None: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index c96a9c1bbae..cd2d69d13a8 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -44,7 +44,10 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( + PromptCachingDeploymentCheck, + new_responses_pins_in_flight, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -130,6 +133,7 @@ class OptionalChecks: redis_cache=cache.redis_cache, in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), ) + self.prompt_caching_pins_in_flight: Final = new_responses_pins_in_flight() self._request_scoped_selectors: tuple[CustomLogger, ...] = () self._request_scoped_policies: tuple[CustomLogger, ...] = () # ``_slack_alerting_loggers``: the Slack alerting logger ``alerting_config`` builds, if any. @@ -167,6 +171,9 @@ class OptionalChecks: group: Final = hosted.get_routing_group(resolved) return group is not None and group.routing_strategy == "priority" + def _group_has_alternatives(self, model: str) -> bool: + return len(self._hosted._core_deployments_named(model) or ()) > 1 # pyright: ignore[reportPrivateUsage] # the host registry's name index + async def repin_session_after_success( self, litellm_params: Mapping[str, object], response: object, cooldowns: "CooldownCache" ) -> None: @@ -291,7 +298,10 @@ class OptionalChecks: callback: CustomLogger | None = None if name == "prompt_caching": callback = PromptCachingDeploymentCheck( - cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_group + cache=self.prompt_caching_pin_cache, + is_priority_group=self._is_priority_group, + responses_pins_in_flight=self.prompt_caching_pins_in_flight, + group_has_alternatives=self._group_has_alternatives, ) elif name == "router_budget_limiting": if self.budget_limiter() is not None: @@ -448,16 +458,21 @@ class OptionalChecks: deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract messages: object, request: Mapping[str, object], - ) -> list[object]: # mutable-ok: as the filters return it + ) -> tuple[list[object], str | None]: # mutable-ok: as the filters return it """``async_callback_filter_deployments``: each routing callback narrows the deployments in - turn; a filter's error is raised for the Rust core to decide on.""" + turn, then the prompt caching check names the deployment it prefers among the rest, which + the Rust core takes only if its later filters keep it; a filter's error is raised for the + Rust core to decide on.""" remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts request_kwargs: Final = dict(request) - for callback in tuple(self.routing_callbacks()): + callbacks: Final = tuple(self.routing_callbacks()) + for callback in callbacks: # The base filter returns its input unchanged: only overriding filters are run. if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above continue + if isinstance(callback, PromptCachingDeploymentCheck): + continue remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above model=model, healthy_deployments=remaining, @@ -465,7 +480,15 @@ class OptionalChecks: request_kwargs=request_kwargs, parent_otel_span=None, ) - return remaining + prompt_caching: Final = next((cb for cb in callbacks if isinstance(cb, PromptCachingDeploymentCheck)), None) + if prompt_caching is None: + return remaining, None + return remaining, await prompt_caching.async_preferred_deployment_id( + model, + remaining, # pyright: ignore[reportArgumentType] # RustCore's deployment dicts + messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs, + ) # -- release -------------------------------------------------------------------------------- @@ -670,15 +693,15 @@ class RustOptionalChecks: async def _answer_callback_filters( self, step: Mapping[str, object], request: Mapping[str, object], messages: object ) -> tuple[object, BaseException | None]: - """A ``callback_filters`` step: the remaining ids, or a filter's error as an attempt - failure with the error itself, which the caller keeps to raise.""" + """A ``callback_filters`` step: the remaining ids and the one a filter prefers, or a + filter's error as an attempt failure with the error itself, which the caller keeps to raise.""" from litellm.router_rust.core import attempt_failure candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts narrowed: Final = _request_with_failed_ids(request, step) try: - remaining: Final = await self._opt.filter_deployments( + remaining, preferred_id = await self._opt.filter_deployments( str(step.get("model_group")), records, # pyright: ignore[reportUnknownArgumentType] # see above messages, @@ -689,7 +712,8 @@ class RustOptionalChecks: # A filter's pin lands on the request kwargs the pick goes on to read. if narrowed is not request and narrowed.get(_ENCRYPTED_PIN) and isinstance(request, dict): request[_ENCRYPTED_PIN] = True # rebind-ok: the request dict the pick goes on to read - return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above + kept: Final = filtered_ids(remaining, candidates) # pyright: ignore[reportUnknownArgumentType] # see above + return (kept, preferred_id), None __all__ = [ diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index e5d944b3aef..2f0cbc0d474 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,23 +4,37 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable, Mapping, Sequence -from typing import Final, TypedDict, cast +import secrets +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, NamedTuple, Protocol, TypedDict, cast +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT from litellm.integrations.anthropic_cache_control_hook import ( AllToolParamValues, AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.responses.litellm_completion_transformation.transformation import ( + RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM, + RESPONSES_WEB_SEARCH_TOOL_TYPES, + LiteLLMCompletionResponsesConfig, + ResponsesToolChatForm, + ResponseTools, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ResponseInputParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import ( get_base_model_from_metadata, @@ -28,7 +42,7 @@ from litellm.utils import ( is_prompt_caching_valid_prompt, ) -from ..prompt_caching_cache import PromptCachingCache +from ..prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache class _PromptCachingSuccessEventKwargs(TypedDict, total=False): @@ -36,7 +50,165 @@ class _PromptCachingSuccessEventKwargs(TypedDict, total=False): litellm_params: ReadOnly[Mapping[str, object] | None] -def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: +class _ResponsesPin(NamedTuple): + cache_key: str + marked_ttl_seconds: int + cacheable_deployment_ids: tuple[str, ...] + + +class ResponsesPinsInFlight(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +class _ResponsesPrompt(NamedTuple): + messages: list[AllMessageValues] # mutable-ok: the prefix helpers take chat messages as a list + tools: list[AllToolParamValues] | None # mutable-ok: the prefix helpers take chat tools as a list + key_tools: list[ChatCompletionToolParam] | None # mutable-ok: the prefix helpers take chat tools as a list + + +_RESPONSES_PIN_METADATA_KEY: Final = "responses_prompt_cache_pin" +_RESPONSES_PINS_IN_FLIGHT_LIMIT: Final = 10_000 +_RESPONSES_PIN_LIFETIME_SECONDS: Final = 3600 +_RETRY_RECORDS: Final = TypeAdapter(tuple[object, ...]) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) +_RESPONSES_TOOL_TYPES_KEPT_AS_SENT: Final = RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM | RESPONSES_WEB_SEARCH_TOOL_TYPES + + +def new_responses_pins_in_flight() -> ResponsesPinsInFlight: + return InMemoryCache( + max_size_in_memory=_RESPONSES_PINS_IN_FLIGHT_LIMIT, default_ttl=_RESPONSES_PIN_LIFETIME_SECONDS + ) + + +def _key_chat_tool(source: Mapping[str, object], chat_tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = as_str_mapping(chat_tool.get("function")) + if source.get("type") != "function" or function is None: + return chat_tool + return {**chat_tool, "function": {key: source[key] for key in function if key in source}} + + +def _key_chat_tools( + sources: Sequence[Mapping[str, object]], forms: Sequence[ResponsesToolChatForm] +) -> Iterator[Mapping[str, object]]: + for source, form in zip(sources, forms, strict=True): + yield from (_key_chat_tool(source, chat_tool) for chat_tool in form.chat_tools) + + +def _responses_prompt(request_data: Mapping[str, object]) -> _ResponsesPrompt | None: + request_input: Final = request_data.get("input") + request_tools: Final = request_data.get("tools") + if request_data.get("previous_response_id") is not None or not isinstance(request_input, str | list): + return None + listed_tools: Final = ( + cast(ResponseTools, request_tools) # cast-ok: a bad tool raises inside the guard below and is caught + if isinstance(request_tools, list) + else None + ) + try: + tools_kept_as_sent: Final = tuple( + tool for tool in listed_tools or () if tool.get("type") in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT + ) + convertible_tools: Final = ( + [tool for tool in listed_tools if tool.get("type") not in _RESPONSES_TOOL_TYPES_KEPT_AS_SENT] + if listed_tools is not None + else None + ) + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=cast(str | ResponseInputParam, request_input), # cast-ok: a bad item raises and is caught here + responses_api_request={"instructions": request_data.get("instructions")}, + item_cache_control=True, + only_since_compaction=True, + ) + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(convertible_tools) + except Exception as error: # noqa: BLE001 # the provider judges the input, a routing preference never does + verbose_logger.debug("Prompt Caching: skipping an input with no chat form: %s", type(error).__name__) + return None + tools: Final = list(chain.from_iterable(form.chat_tools for form in forms)) + key_tools: Final = [*_key_chat_tools(convertible_tools or (), forms), *tools_kept_as_sent] + return _ResponsesPrompt( + messages=cast(list[AllMessageValues], messages), # cast-ok: the bridge's chat messages, read as plain JSON + tools=cast(list[AllToolParamValues], tools) or None, # cast-ok: the bridge's chat tools, read as plain JSON + key_tools=cast(list[ChatCompletionToolParam], key_tools) or None, # cast-ok: tools are hashed as plain JSON + ) + + +def _retry_records(metadata: Mapping[str, object]) -> tuple[object, ...]: + try: + return _RETRY_RECORDS.validate_python(metadata.get("previous_models", ())) + except ValidationError: + return () + + +def _record_deployment_ids(records: tuple[object, ...]) -> Iterator[str]: + deployment_ids: Final = ((as_str_mapping(record) or {}).get("deployment_id") for record in records) + return (deployment_id for deployment_id in deployment_ids if isinstance(deployment_id, str)) + + +def _retried_deployment_ids(metadata: Mapping[str, object]) -> frozenset[str] | None: + records: Final = _retry_records(metadata) + retry_count: Final = metadata.get("request_retry_count") + if type(retry_count) is int and retry_count > len(records): + return None + return frozenset(_record_deployment_ids(records)) + + +def _deployment_ids_tried_in_request(request_data: Mapping[str, object]) -> frozenset[str] | None: + retried: Final = tuple( + _retried_deployment_ids(as_str_mapping(request_data.get(bucket)) or {}) + for bucket in ("litellm_metadata", "metadata") + ) + known: Final = tuple(ids for ids in retried if ids is not None) + if len(known) < len(retried): + return None + return deployment_ids_failed_in_request(request_data).union(*known) + + +def _previous_response_model_id(request_data: Mapping[str, object]) -> str | None: + previous_response_id: Final = request_data.get("previous_response_id") + model_id: Final = ( + ResponsesAPIRequestUtils.get_model_id_from_response_id(previous_response_id) + if isinstance(previous_response_id, str) + else None + ) + tried: Final = _deployment_ids_tried_in_request(request_data) if model_id is not None else None + return model_id if model_id is not None and tried is not None and model_id not in tried else None + + +def _deployment_min_tokens(deployment: Mapping[str, Mapping[str, object]], fallback: int) -> int: + deployment_model: Final = deployment.get("litellm_params", {}).get("model") + return get_prompt_cache_min_tokens(model=deployment_model) if isinstance(deployment_model, str) else fallback + + +async def _cacheable_deployment_ids( + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + reached_minimum: int, + reaches: Callable[[int], Awaitable[bool]], +) -> tuple[str, ...]: + minimums: Final = tuple( + (str(deployment["model_info"]["id"]), _deployment_min_tokens(deployment, reached_minimum)) + for deployment in healthy_deployments + ) + reached: Final = MappingProxyType( + { + minimum: minimum <= reached_minimum or await reaches(minimum) + for minimum in frozenset(minimum for _, minimum in minimums) + } + ) + return tuple(sorted(deployment_id for deployment_id, minimum in minimums if reached[minimum])) + + +def _deployment_models(healthy_deployments: Sequence[Mapping[str, object]]) -> Iterator[str]: + models: Final = ( + (as_str_mapping(deployment.get("litellm_params")) or {}).get("model") for deployment in healthy_deployments + ) + return (model for model in models if isinstance(model, str) and model) + + +def _get_min_token_count_for_deployments(healthy_deployments: Sequence[Mapping[str, object]]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -52,11 +224,7 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int losing a cache hit it had earned. The lowest can only cost a lookup that finds nothing. """ return min( - ( - get_prompt_cache_min_tokens(model=deployment["litellm_params"]["model"]) - for deployment in healthy_deployments - if deployment.get("litellm_params", {}).get("model") - ), + (get_prompt_cache_min_tokens(model=model) for model in _deployment_models(healthy_deployments)), default=DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, ) @@ -75,10 +243,20 @@ def _honors_one_hour_cache_ttl( class PromptCachingDeploymentCheck(CustomLogger): - def __init__(self, cache: DualCache, is_priority_group: Callable[[str], bool] | None = None): + def __init__( + self, + cache: DualCache, + is_priority_group: Callable[[str], bool] | None = None, + responses_pins_in_flight: ResponsesPinsInFlight | None = None, + group_has_alternatives: Callable[[str], bool] | None = None, + ): super().__init__() self.cache = cache self.is_priority_group = is_priority_group + self.group_has_alternatives: Final = group_has_alternatives + self.responses_pins_in_flight: Final = ( + responses_pins_in_flight if responses_pins_in_flight is not None else new_responses_pins_in_flight() + ) async def async_filter_deployments( self, @@ -88,43 +266,81 @@ class PromptCachingDeploymentCheck(CustomLogger): request_kwargs: dict | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: - if self.is_priority_group is not None and self.is_priority_group(model): - return healthy_deployments - if request_kwargs is not None and request_kwargs.get("_target_order") is not None: - return healthy_deployments - if not healthy_deployments[1:]: - return healthy_deployments + preferred_id: Final = await self.async_preferred_deployment_id( + model, healthy_deployments, messages, request_kwargs + ) + preferred: Final = [ + deployment for deployment in healthy_deployments if deployment["model_info"]["id"] == preferred_id + ] + return preferred[:1] or healthy_deployments - request_data: Final = ( - cast( # cast-ok: request_kwargs is a bare dict; this only reads its keys - Mapping[str, object], request_kwargs - ) - if request_kwargs is not None + async def async_preferred_deployment_id( + self, + model: str, + healthy_deployments: Sequence[Mapping[str, Mapping[str, object]]], + messages: list[AllMessageValues] | None, # mutable-ok: the request's messages list, read as is + request_data: Mapping[str, object] | None, + ) -> str | None: + if self.is_priority_group is not None and self.is_priority_group(model): + return None + if request_data is not None and request_data.get("_target_order") is not None: + return None + pin_metadata: Final = ( + request_data.get("litellm_metadata") if messages is None and request_data is not None else None + ) + single_healthy: Final = not healthy_deployments[1:] + if single_healthy and not ( + isinstance(pin_metadata, dict) + and self.group_has_alternatives is not None + and self.group_has_alternatives(model) + ): + return None + previous_attempt_token: Final = (as_str_mapping(pin_metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + previous_response_model_id: Final = ( + _previous_response_model_id(request_data) + if request_data is not None and isinstance(pin_metadata, dict) + else None + ) + if previous_response_model_id is not None: + return previous_response_model_id + responses_prompt: Final = ( + await offload_token_count(_responses_prompt)(request_data) + if request_data is not None and isinstance(pin_metadata, dict) else None ) request_tools: Final = ( - cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads + responses_prompt.tools + if responses_prompt is not None + else cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads list[AllToolParamValues] | None, request_data.get("tools") ) - if request_data is not None + if request_data is not None and messages is not None else None ) + request_messages: Final = responses_prompt.messages if responses_prompt is not None else messages prompt_messages: Final = ( PromptCachingCache.prepend_system_prompt( - messages, request_data.get("system") if request_data is not None else None + request_messages, request_data.get("system") if request_data is not None else None ) - if messages is not None + if request_messages is not None else None ) - if prompt_messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( - messages=prompt_messages, - model=model, - min_token_count=_get_min_token_count_for_deployments(healthy_deployments), - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types - list[ChatCompletionToolParam] | None, - request_tools, - ), - ): + + async def reaches(minimum: int) -> bool: + return await offload_token_count( + is_prompt_caching_valid_prompt + )( + messages=prompt_messages, + model=model, + min_token_count=minimum, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types + list[ChatCompletionToolParam] | None, + request_tools, + ), + ) + + group_minimum: Final = _get_min_token_count_for_deployments(healthy_deployments) + if prompt_messages is not None and await reaches(group_minimum): prompt_cache: Final = PromptCachingCache( cache=self.cache, ) @@ -134,34 +350,49 @@ class PromptCachingDeploymentCheck(CustomLogger): ## be derived from the messages as they will be sent, not as they arrive here. affinity_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( messages=prompt_messages, - models=( - deployment["litellm_params"]["model"] - for deployment in healthy_deployments - if isinstance(deployment.get("litellm_params"), dict) and deployment["litellm_params"].get("model") - ), + models=_deployment_models(healthy_deployments), tools=request_tools, enable_prompt_caching=( - request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None + request_data.get("enable_prompt_caching") is True if request_data is not None else None ), - request_kwargs=request_kwargs, + request_kwargs=request_data, ) - - model_id_dict: Final = await prompt_cache.async_get_model_id( - messages=affinity_messages, - tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes - Sequence[ChatCompletionToolParam] | None, - request_tools, - ), + prefix_tools: Final = ( + responses_prompt.key_tools + if responses_prompt is not None + else cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes + Sequence[ChatCompletionToolParam] | None, request_tools + ) ) - if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request( - request_data or {} - ): - model_id: Final = model_id_dict["model_id"] - for deployment in healthy_deployments: - if deployment["model_info"]["id"] == model_id: - return [deployment] + positions: Final = await PromptCachingCache.async_prefix_positions(affinity_messages, prefix_tools) + if responses_prompt is not None and isinstance(pin_metadata, dict) and positions: + if isinstance(previous_attempt_token, str): + self.responses_pins_in_flight.delete_cache(previous_attempt_token) + pin_token: Final = secrets.token_hex(16) + self.responses_pins_in_flight.set_cache( + pin_token, + _ResponsesPin( + cache_key=positions[-1].cache_key, + marked_ttl_seconds=await offload_token_count(PromptCachingCache.pin_ttl_seconds)( + affinity_messages, prefix_tools + ), + cacheable_deployment_ids=await _cacheable_deployment_ids( + healthy_deployments, + group_minimum, + reaches, + ), + ), + ) + pin_metadata[_RESPONSES_PIN_METADATA_KEY] = pin_token + if single_healthy: + return None - return healthy_deployments + model_id_dict: Final = await prompt_cache.async_get_model_id_at(positions) + tried: Final = _deployment_ids_tried_in_request(request_data or {}) if model_id_dict is not None else None + if model_id_dict is not None and tried is not None and model_id_dict["model_id"] not in tried: + return model_id_dict["model_id"] + + return None def exempt_from_message_redaction(self) -> bool: return True @@ -173,6 +404,7 @@ class PromptCachingDeploymentCheck(CustomLogger): start_time: object, end_time: object, ) -> None: + recorded_pin: Final = self._take_responses_pin(kwargs) standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -180,6 +412,10 @@ class PromptCachingDeploymentCheck(CustomLogger): call_type: Final = standard_logging_object["call_type"] + if call_type in _RESPONSES_CALL_TYPES: + await self._pin_responses_request(recorded_pin, kwargs, standard_logging_object) + return + if ( call_type != CallTypes.completion.value and call_type != CallTypes.acompletion.value @@ -241,3 +477,33 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + + def _take_responses_pin(self, kwargs: _PromptCachingSuccessEventKwargs) -> object: + litellm_params: Final = kwargs.get("litellm_params") or {} + metadata: Final = litellm_params.get("litellm_metadata") + pin_token: Final = (as_str_mapping(metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) + if not isinstance(pin_token, str): + return None + if isinstance(metadata, dict): + metadata.pop(_RESPONSES_PIN_METADATA_KEY, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict + recorded_pin: Final = self.responses_pins_in_flight.get_cache(pin_token) + self.responses_pins_in_flight.delete_cache(pin_token) + return recorded_pin + + async def _pin_responses_request( + self, + recorded_pin: object, + kwargs: _PromptCachingSuccessEventKwargs, + standard_logging_object: StandardLoggingPayload, + ) -> None: + litellm_params: Final = kwargs.get("litellm_params") or {} + model_id: Final = standard_logging_object["model_id"] + if not isinstance(recorded_pin, _ResponsesPin) or model_id not in recorded_pin.cacheable_deployment_ids: + return + await PromptCachingCache(cache=self.cache).async_set_pin( + model_id, + recorded_pin.cache_key, + recorded_pin.marked_ttl_seconds + if _honors_one_hour_cache_ttl(standard_logging_object, litellm_params) + else PROMPT_CACHE_PIN_TTL_SECONDS, + ) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 1d91255e28c..e01d744dfb6 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -432,18 +432,18 @@ class PromptCachingCache: if not positions: return - ttl: Final = ( + await self.async_set_pin( + model_id, + positions[-1].cache_key, await offload_token_count(PromptCachingCache.pin_ttl_seconds)(messages, tools) if honors_one_hour_ttl - else PROMPT_CACHE_PIN_TTL_SECONDS + else PROMPT_CACHE_PIN_TTL_SECONDS, ) - self.cache.in_memory_cache.delete_cache(positions[-1].cache_key) + + async def async_set_pin(self, model_id: str, cache_key: str, ttl: int) -> None: + self.cache.in_memory_cache.delete_cache(cache_key) with service_target(_PROMPT_CACHE_PINS_TARGET): - await self.cache.async_set_cache( - positions[-1].cache_key, - PromptCachingCacheValue(model_id=model_id), - ttl=ttl, - ) + await self.cache.async_set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=ttl) async def async_get_model_id( self, @@ -455,7 +455,10 @@ class PromptCachingCache: same way the provider cache does, so a breakpoint that moved forward since the last turn still lands on the deployment whose cache holds the earlier prefix. """ - cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools)) + return await self.async_get_model_id_at(await PromptCachingCache.async_prefix_positions(messages, tools)) + + async def async_get_model_id_at(self, positions: tuple[PrefixPosition, ...]) -> PromptCachingCacheValue | None: + cache_keys: Final = _lookback_keys(positions) if not cache_keys: return None diff --git a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py index 9d4bf62adf5..38129013294 100644 --- a/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py +++ b/tests/integration/routing/test_prompt_caching_affinity_tool_call_marks.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Final import pytest -from integration._support.client import Gateway, eventually +from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy_process from integration._support.wire import Wire, wire_server @@ -24,6 +24,7 @@ from pydantic import JsonValue _MODEL: Final = "affinity-claude" _FOLLOW_UPS: Final = 24 +_PINS_ROUTE: Final = "/memory-usage-in-mem-cache-items" _LONG_SYSTEM: Final = "Answer from the tool results below. " + "lorem ipsum " * 1500 @@ -91,3 +92,68 @@ def test_router_affinity_is_lost_when_auto_caching_stands_down_for_tool_call_mar assert [status for status, _, _ in follow_ups] == [200] * _FOLLOW_UPS, [text for _, _, text in follow_ups] assert {first_marker, *follow_up_markers} == served["left"] | served["right"] assert all(served[side] & set(follow_up_markers) for side in served), served + + +def _responses_first_turn(marker: str) -> list[JsonValue]: + return [ + {"role": "user", "content": "What is the weather in London?"}, + {"type": "function_call", "call_id": "call_london", "name": "lookup_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_london", "output": "sunny"}, + {"role": "user", "content": final_text(marker)}, + ] + + +def _responses_body(session: str, request_input: list[JsonValue]) -> dict[str, JsonValue]: + return { + "model": _MODEL, + "instructions": f"{_LONG_SYSTEM} session {session}", + "tools": [{"type": "function", "name": "lookup_weather", "parameters": {"type": "object", "properties": {}}}], + "input": request_input, + } + + +def _pins(gateway: Gateway) -> dict[str, JsonValue]: + return object_value(object_value(gateway.request("GET", _PINS_ROUTE).json())["prompt_caching_pin_cache"]) + + +@pytest.mark.timeout(240) +def test_responses_thread_holds_its_deployment_when_the_bridge_logs_chat_messages( + gateway: Gateway, tmp_path: Path +) -> None: + session: Final = new_marker() + first_marker: Final = new_marker() + follow_up_markers: Final = tuple(new_marker() for _ in range(_FOLLOW_UPS)) + first_turn: Final = _responses_first_turn(first_marker) + with scratch_database() as database_url, wire_server(anthropic_peer) as left, wire_server(anthropic_peer) as right: + config: Final = owned_config( + tmp_path, + [ + {**anthropic_deployment(_MODEL, left.url), "model_info": {"id": f"responses-left-{session}"}}, + {**anthropic_deployment(_MODEL, right.url), "model_info": {"id": f"responses-right-{session}"}}, + ], + router_settings={"optional_pre_call_checks": ["prompt_caching"]}, + ) + with ( + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=1, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"enable_prompt_caching": True}) + first: Final = owned.gateway.request("POST", "/v1/responses", _responses_body(session, first_turn), key=key) + assert first.status_code == 200, first.text + eventually(lambda: _pins(owned.gateway), bool, seconds=30) + follow_ups: Final = tuple( + owned.gateway.request( + "POST", "/v1/responses", _responses_body(session, _follow_up(first_turn, marker)), key=key + ) + for marker in follow_up_markers + ) + served: Final = {"left": _served(left), "right": _served(right)} + assert [response.status_code for response in follow_ups] == [200] * _FOLLOW_UPS, follow_ups[0].text + assert {first_marker, *follow_up_markers} in tuple(served.values()), served diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 5169642bc43..76e9a70614f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -2,6 +2,9 @@ import asyncio import copy import os import functools +import json +import logging +import threading import uuid from typing import Final, cast @@ -10,7 +13,11 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS +from litellm.constants import ( + DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, + PROMPT_CACHE_LOOKBACK_POSITIONS, + REDACTED_BY_LITELLM, +) from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -19,6 +26,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardLoggingPayload @@ -138,10 +146,13 @@ def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> N }, ) - assert _honors_one_hour_cache_ttl( - standard_logging_object, - {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, - ) is True + assert ( + _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) + is True + ) def test_bedrock_application_profile_cache_hit_reads_base_model_from_metadata_model_info() -> None: @@ -244,6 +255,16 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): assert filtered == [deployments[1]] +class _ReadTrackingMessages(list[AllMessageValues]): + def __init__(self, messages: list[AllMessageValues]) -> None: + super().__init__(messages) + self.read = False + + def __iter__(self): + self.read = True + return super().__iter__() + + class _PinLookupCounter(DualCache): def __init__(self) -> None: super().__init__() @@ -275,17 +296,19 @@ async def test_async_filter_deployments_skips_prefix_hash_for_a_single_deploymen prompt that hash alone is ~30 ms of GIL-holding work per request. """ cache = _PinLookupCounter() - check = PromptCachingDeploymentCheck(cache=cache) + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) deployments = _deployments("anthropic/claude-opus-4-6") messages = _messages(word_count=5000) await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=messages, tools=None) + tracked_messages = _ReadTrackingMessages(messages) filtered = await check.async_filter_deployments( - model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=tracked_messages ) assert filtered == deployments assert cache.pin_lookups == 0 + assert tracked_messages.read is False two = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") assert await check.async_filter_deployments( @@ -313,10 +336,20 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): assert filtered == deployments +_RETRIES_ON_DEP_1: Final = tuple({"deployment_id": "dep-1"} for _ in range(4)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "failed_in_request", - [{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}], + [ + {"_excluded_deployment_ids": ["dep-2"]}, + {"_retry_skipped_deployment_ids": ("dep-2",)}, + {"metadata": {"previous_models": [{"deployment_id": "dep-2"}], "request_retry_count": 1}}, + {"litellm_metadata": {"previous_models": ({"deployment_id": "dep-2"},), "request_retry_count": 1}}, + {"metadata": {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 5}}, + ], + ids=["excluded", "retry_skipped", "retried", "retried_litellm_metadata", "retries_rolled_over"], ) async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request( failed_in_request: dict[str, object], @@ -334,6 +367,34 @@ async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed ) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "retry_metadata", + [ + {"previous_models": _RETRIES_ON_DEP_1, "request_retry_count": 4}, + {"previous_models": "dep-2"}, + {"previous_models": [None, "dep-2", {"deployment_id": None}], "request_retry_count": 3}, + ], + ids=["every_retry_on_record", "malformed_breadcrumbs", "malformed_records"], +) +async def test_async_filter_deployments_keeps_the_pin_while_every_retry_is_on_record( + retry_metadata: dict[str, object], +): + cache: Final = DualCache() + check: Final = PromptCachingDeploymentCheck(cache=cache) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + filtered: Final = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"), + messages=messages, + request_kwargs={"metadata": retry_metadata}, + ) + + assert [deployment["model_info"]["id"] for deployment in filtered] == ["dep-2"] + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ @@ -1535,3 +1596,920 @@ async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_r await _eventually(lambda: len(probe.logged_messages) == 5) redacted = [{"role": "user", "content": "redacted-by-litellm"}] assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages + + +RESPONSES_BRIDGED_MODEL: Final = "anthropic/claude-sonnet-4-6" +RESPONSES_NATIVE_MODEL: Final = "openai/gpt-5.4-mini" +_RESPONSES_TOOL: Final = { + "type": "function", + "name": "get_time", + "parameters": {"type": "object", "properties": {}}, +} + + +def _marked_input_text(text: str) -> dict[str, object]: + return {"type": "input_text", "text": text, "cache_control": {"type": "ephemeral"}} + + +def _marked_responses_input() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(LONG_PROMPT)]}, + {"role": "user", "content": "hello"}, + ] + + +def _responses_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + + +async def _serve_response(router: litellm.Router, request_input: object) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=copy.deepcopy(request_input), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _deployments_serving_repeats(router: litellm.Router, request_input: object) -> set[str]: + await _serve_response(router, request_input) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + return {await _serve_response(router, request_input) for _ in range(8)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_repeated_responses_request_pins_to_one_deployment(model: str): + assert len(await _deployments_serving_repeats(_responses_router(model), _marked_responses_input())) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [RESPONSES_BRIDGED_MODEL, RESPONSES_NATIVE_MODEL], ids=["bridged", "native"]) +async def test_responses_thread_stays_on_the_deployment_holding_its_prefix(model: str): + router = _responses_router(model) + first_turn = _marked_responses_input() + pinned_deployment = await _serve_response(router, first_turn) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + later_turns = ( + [*first_turn, {"role": "assistant", "content": "ok"}, {"role": "user", "content": f"question {turn}"}] + for turn in range(8) + ) + + assert {await _serve_response(router, turn) for turn in later_turns} == {pinned_deployment} + + +@pytest.mark.asyncio +async def test_responses_request_pins_after_prompt_caching_is_turned_off_and_on_again(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + await _serve_response(router, _marked_responses_input()) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.update_settings(optional_pre_call_checks=[]) + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + +def _team_tagged_router(model: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": model, "api_key": "sk-fake", "tags": [team]}, + "model_info": {"id": model_id}, + } + for model_id, team in (("dep-1", "team-a"), ("dep-2", "team-b")) + ], + optional_pre_call_checks=["prompt_caching"], + enable_tag_filtering=True, + ) + + +async def _serve_team_response(router: litellm.Router, team: str) -> str: + response = await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + instructions="Be terse.", + tools=[_RESPONSES_TOOL], + litellm_metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +async def _serve_team_completion(router: litellm.Router, team: str) -> str: + response = await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=_messages(word_count=5000), + metadata={"tags": [team]}, + mock_response="ok", + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, serve", + [(RESPONSES_NATIVE_MODEL, _serve_team_response), (RESPONSES_BRIDGED_MODEL, _serve_team_completion)], + ids=["responses", "chat"], +) +async def test_prompt_cache_pin_never_overrides_the_request_tags(model: str, serve): + router = _team_tagged_router(model) + assert await serve(router, "team-a") == "dep-1" + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + + assert await serve(router, "team-b") == "dep-2" + + +@pytest.mark.asyncio +async def test_retry_moves_off_the_pinned_deployment_that_failed(): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": { + "model": RESPONSES_BRIDGED_MODEL, + "api_key": "sk-fake", + "mock_response": mock_response, + "weight": weight, + }, + "model_info": {"id": model_id}, + } + for model_id, mock_response, weight in ( + ("dep-1", "litellm.InternalServerError", 1), + ("dep-2", "ok", 1000), + ) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=3, + retry_after=0, + disable_cooldowns=True, + ) + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-1", messages=messages, tools=None + ) + + response: Final = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reused_as", ["metadata", "litellm_metadata"]) +async def test_pin_holds_for_a_metadata_dict_reused_after_a_request_failed_on_the_pinned_deployment(reused_as: str): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_BRIDGED_MODEL, "api_key": "sk-fake", "weight": weight}, + "model_info": {"id": model_id}, + } + for model_id, weight in (("dep-1", 1000), ("dep-2", 1)) + ], + optional_pre_call_checks=["prompt_caching"], + num_retries=6, + retry_after=0, + allowed_fails=100, + ) + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="anthropic", model=RESPONSES_BRIDGED_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + reused_metadata: Final[dict[str, object]] = {} + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[{"role": "user", "content": "hello"}], + metadata=reused_metadata, + mock_response="litellm.InternalServerError", + ) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + messages: Final = _messages(word_count=5000) + await PromptCachingCache(cache=router.prompt_caching_pin_cache).async_add_model_id( + model_id="dep-2", messages=messages, tools=None + ) + + response: Final = await router.acompletion( + model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok", **{reused_as: reused_metadata} + ) + + assert response._hidden_params["model_id"] == "dep-2" + + +@pytest.mark.asyncio +async def test_responses_mark_on_the_input_item_pins_to_one_deployment(): + request_input = [ + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +async def _read_responses_pin( + check: PromptCachingDeploymentCheck, + request: dict[str, object], + models: tuple[str, ...] = (AUTO_CACHING_MODEL, AUTO_CACHING_MODEL), +) -> tuple[str, ...]: + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=_deployments(*models), + messages=None, + request_kwargs=request, + ) + return tuple(str(deployment["model_info"]["id"]) for deployment in filtered) + + +async def _log_responses_success( + check: PromptCachingDeploymentCheck, + logged_messages: object, + metadata: dict[str, object], + model_id: str = "dep-1", + model_map_value: dict[str, float | None] | None = None, + call_type: str = "aresponses", +) -> None: + event = _success_kwargs( + call_type=call_type, + model=AUTO_CACHING_MODEL, + model_id=model_id, + request_messages=logged_messages, + request_kwargs={}, + ) + await check.async_log_success_event( + kwargs={ + **event, + "standard_logging_object": { + **event["standard_logging_object"], + "model_map_information": {"model_map_key": AUTO_CACHING_MODEL, "model_map_value": model_map_value}, + }, + "litellm_params": {"litellm_metadata": metadata}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + +async def _pinned_deployment_ids( + first_request: dict[str, object], logged_messages: object, next_request: dict[str, object] +) -> tuple[str, ...]: + check = PromptCachingDeploymentCheck(cache=DualCache()) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**first_request, "litellm_metadata": metadata}) + await _log_responses_success(check, logged_messages, metadata) + return await _read_responses_pin(check, {**next_request, "litellm_metadata": {}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "logged_messages", + [ + _marked_responses_input(), + [ + {"role": "system", "content": "Be terse."}, + {"role": "system", "content": LONG_PROMPT, "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function"}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "12:00"}, + ], + [{"role": "user", "content": REDACTED_BY_LITELLM}], + ], + ids=["client_input", "bridged_chat_messages", "redacted"], +) +async def test_responses_pin_does_not_depend_on_what_the_success_event_logged( + logged_messages: list[dict[str, object]], +): + request = {"input": _marked_responses_input(), "instructions": "Be terse.", "tools": [_RESPONSES_TOOL]} + + pinned = await _pinned_deployment_ids(request, logged_messages, request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "forged_pin", + [["forged-redis-key", 3600], {"cache_key": "forged-redis-key", "marked_ttl_seconds": 3600}], + ids=["array", "object"], +) +async def test_responses_pin_sent_in_the_request_metadata_is_never_written(forged_pin: object): + cache = DualCache() + + await _log_responses_success( + PromptCachingDeploymentCheck(cache=cache), + _marked_responses_input(), + {"responses_prompt_cache_pin": forged_pin}, + ) + + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_interleaved_responses_requests_each_pin_their_own_prompt(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + terse = {"input": _marked_responses_input(), "instructions": "Be terse."} + verbose = {"input": _marked_responses_input(), "instructions": "Be verbose."} + terse_metadata: dict[str, object] = {} + verbose_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**terse, "litellm_metadata": terse_metadata}) + await _read_responses_pin(check, {**verbose, "litellm_metadata": verbose_metadata}) + + await _log_responses_success(check, _marked_responses_input(), terse_metadata, model_id="dep-1") + await _log_responses_success(check, _marked_responses_input(), verbose_metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**terse, "litellm_metadata": {}}) == ("dep-1",) + assert await _read_responses_pin(check, {**verbose, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_retried_responses_request_leaves_no_pin_in_flight_once_it_succeeds(): + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=DualCache(), responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + request = {"input": _marked_responses_input(), "litellm_metadata": metadata} + await _read_responses_pin(check, request) + await _read_responses_pin(check, request) + + await _log_responses_success(check, _marked_responses_input(), metadata) + + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompact_responses", "acreate_interaction"]) +async def test_request_sharing_the_responses_input_shape_leaves_no_pin_behind(call_type: str): + cache = DualCache() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck(cache=cache, responses_pins_in_flight=pins_in_flight) + metadata: dict[str, object] = {} + await _read_responses_pin(check, {"input": _marked_responses_input(), "litellm_metadata": metadata}) + assert pins_in_flight.cache_dict + + await _log_responses_success(check, _marked_responses_input(), metadata, call_type=call_type) + + assert not pins_in_flight.cache_dict + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_served_by_the_only_healthy_deployment_pins_it(): + cache = _PinLookupCounter() + check = PromptCachingDeploymentCheck(cache=cache, group_has_alternatives=lambda group: True) + request = {"input": _marked_responses_input(), "instructions": "Be terse."} + only_healthy = [_deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL)[1]] + metadata: dict[str, object] = {} + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=only_healthy, + messages=None, + request_kwargs={**request, "litellm_metadata": metadata}, + ) + assert filtered == only_healthy + assert cache.pin_lookups == 0 + + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + assert await _read_responses_pin(check, {**request, "litellm_metadata": {}}) == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_group_skips_the_prefix_hash(): + cache = _PinLookupCounter() + pins_in_flight = InMemoryCache() + check = PromptCachingDeploymentCheck( + cache=cache, responses_pins_in_flight=pins_in_flight, group_has_alternatives=lambda group: False + ) + deployments = _deployments(AUTO_CACHING_MODEL) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=None, + request_kwargs={"input": _marked_responses_input(), "litellm_metadata": {}}, + ) + + assert filtered == deployments + assert cache.pin_lookups == 0 + assert not pins_in_flight.cache_dict + + +@pytest.mark.asyncio +async def test_responses_request_to_a_single_deployment_router_group_writes_no_pin(monkeypatch): + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": RESPONSES_NATIVE_MODEL, "api_key": "sk-fake"}, + "model_info": {"id": "dep-1"}, + } + ], + optional_pre_call_checks=["prompt_caching"], + ) + logged_after_the_check: list[str] = [] + + class _LoggedAfterTheCheck(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + logged_after_the_check.append(kwargs["litellm_call_id"]) + + monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, _LoggedAfterTheCheck()]) + + assert await _serve_response(router, _marked_responses_input()) == "dep-1" + + assert await _eventually(lambda: logged_after_the_check) + assert not router.prompt_caching_pin_cache.in_memory_cache.cache_dict + + +def _distinct_marked_responses_input(index: int) -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_marked_input_text(f"prompt {index} {LONG_PROMPT}")]}, + {"role": "user", "content": "hello"}, + ] + + +@pytest.mark.asyncio +async def test_responses_prompts_served_while_a_peer_cools_down_stay_pinned_once_it_recovers(): + router = _responses_router(RESPONSES_NATIVE_MODEL) + pins = router.prompt_caching_pin_cache.in_memory_cache.cache_dict + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-1", + original_exception=litellm.RateLimitError( + message="upstream 429", llm_provider="openai", model=RESPONSES_NATIVE_MODEL + ), + exception_status=429, + cooldown_time=60, + ) + prompts = [_distinct_marked_responses_input(index) for index in range(8)] + for prompt in prompts: + pins_before = len(pins) + assert await _serve_response(router, prompt) == "dep-2" + assert await _eventually(lambda before=pins_before: len(pins) > before) + router.cooldown_cache.cooldown_store.evict_local("dep-1") + unpinned = { + ( + await router.aresponses(model=MODEL_GROUP_ALIAS, input=f"unpinned {index}", mock_response="ok") + )._hidden_params["model_id"] + for index in range(16) + } + assert unpinned == {"dep-1", "dep-2"} + + assert {await _serve_response(router, prompt) for prompt in prompts} == {"dep-2"} + + +@pytest.mark.asyncio +async def test_responses_with_other_instructions_are_a_different_prefix(): + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "instructions": "Be terse."}, + _marked_responses_input(), + {"input": _marked_responses_input(), "instructions": "Be verbose."}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_with_another_tool_set_are_a_different_prefix(): + other_tool = {**_RESPONSES_TOOL, "name": "get_weather"} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [_RESPONSES_TOOL]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_tool]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_custom_tools_with_another_grammar_are_a_different_prefix(): + custom_tool: Final = { + "type": "custom", + "name": "run_sql", + "description": "Run a query", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: SELECT"}, + } + other_grammar: Final = {**custom_tool, "format": {"type": "grammar", "syntax": "lark", "definition": "start: UPDATE"}} + + pinned = await _pinned_deployment_ids( + {"input": _marked_responses_input(), "tools": [custom_tool]}, + _marked_responses_input(), + {"input": _marked_responses_input(), "tools": [other_grammar]}, + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "parameters", + [{"type": "object", "properties": {}}, {"properties": {"city": {"type": "string"}}}], + ids=["typed_parameters", "untyped_parameters"], +) +async def test_chat_and_responses_requests_sending_the_same_tool_share_one_prefix(parameters: dict[str, object]): + described_tool: Final = {**_RESPONSES_TOOL, "description": "Get the time", "parameters": parameters} + chat_tool: Final = { + "type": "function", + "function": {key: value for key, value in described_tool.items() if key != "type"}, + } + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[ + {"role": "system", "content": [{**_marked_input_text(LONG_PROMPT), "type": "text"}]}, + {"role": "user", "content": "hello"}, + ], + tools=[chat_tool], + mock_response="ok", + ) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + chat_pin_keys: Final = set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + await router.aresponses( + model=MODEL_GROUP_ALIAS, input=_marked_responses_input(), tools=[described_tool], mock_response="ok" + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) == chat_pin_keys + + +@pytest.mark.asyncio +async def test_responses_marked_tool_set_long_enough_to_cache_pins_on_its_own(): + marked_tool = {**_RESPONSES_TOOL, "description": LONG_PROMPT, "cache_control": {"type": "ephemeral"}} + request = {"input": "hello", "tools": [marked_tool]} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_continuation_turn_is_left_to_the_deployment_holding_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + new_thread = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids( + new_thread, [{"role": "user", "content": "hello"}], {**new_thread, "previous_response_id": "resp_1"} + ) + + assert len(pinned) == 2 + + +@pytest.mark.asyncio +async def test_responses_answer_from_a_fallback_group_keeps_the_primary_groups_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + assert await _serve_turn(check, request["input"]) == 2 + fallback_hop_metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": fallback_hop_metadata}) + + await _log_responses_success(check, request["input"], fallback_hop_metadata, model_id="fallback-dep") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answering_deployment,expected_pin", + [("dep-1", ("dep-1", "dep-2")), ("dep-2", ("dep-2",))], + ids=["below_its_minimum", "above_its_minimum"], +) +async def test_responses_pin_goes_only_to_a_deployment_whose_model_caches_the_prefix( + answering_deployment: str, expected_pin: tuple[str, ...] +): + mixed_group = ("anthropic/claude-opus-4-5", "anthropic/claude-sonnet-4-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": [{"role": "user", "content": [_marked_input_text("word " * 1400)]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}, mixed_group) + + await _log_responses_success(check, request["input"], metadata, model_id=answering_deployment) + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}, mixed_group) + assert pinned == expected_pin + + +@pytest.mark.asyncio +async def test_embedding_request_never_takes_a_responses_pin(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + pinned_input = _marked_responses_input() + assert await _serve_turn(check, pinned_input) == 2 + assert len(await _read_responses_pin(check, {"input": pinned_input, "litellm_metadata": {}})) == 1 + + filtered = await _read_responses_pin(check, {"input": pinned_input, "metadata": {}}) + + assert len(filtered) == 2 + + +class _InputReadOnlyOnceTheLoopTurns(list[dict[str, object]]): + def __init__(self, items: list[dict[str, object]], loop_turned: threading.Event) -> None: + super().__init__(items) + self.loop_turned: Final = loop_turned + self.loop_turned_before_first_read: bool | None = None + + def __iter__(self): + if self.loop_turned_before_first_read is None: + self.loop_turned_before_first_read = self.loop_turned.wait(timeout=5) + return super().__iter__() + + +@pytest.mark.asyncio +async def test_responses_thread_is_converted_off_the_event_loop(): + loop_turned = threading.Event() + thread = _InputReadOnlyOnceTheLoopTurns(_marked_responses_input(), loop_turned) + asyncio.get_running_loop().call_soon(loop_turned.set) + + await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {"input": thread, "litellm_metadata": {}} + ) + + assert thread.loop_turned_before_first_read is True + + +@pytest.mark.asyncio +async def test_responses_long_instructions_count_toward_the_caching_minimum(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"input": "hello", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, [{"role": "user", "content": "hello"}], request) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_map_value,pinned_past_five_minutes", + [ + pytest.param({"cache_creation_input_token_cost_above_1hr": 6.6e-06}, True, id="one-hour-cache-writes"), + pytest.param({"cache_creation_input_token_cost_above_1hr": None}, False, id="no-one-hour-cache-writes"), + ], +) +async def test_responses_one_hour_pin_follows_the_deployments_one_hour_cache_support( + model_map_value: dict[str, float | None], pinned_past_five_minutes: bool +): + clock = _ManualClock() + check = PromptCachingDeploymentCheck(cache=DualCache(in_memory_cache=InMemoryCache(clock=clock))) + marked = {**_marked_input_text(LONG_PROMPT), "cache_control": {"type": "ephemeral", "ttl": "1h"}} + request = {"input": [{"role": "user", "content": [marked]}]} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, request["input"], metadata, model_map_value=model_map_value) + + clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1) + + assert (len(await _read_responses_pin(check, {**request, "litellm_metadata": {}})) == 1) is pinned_past_five_minutes + + +@pytest.mark.asyncio +async def test_responses_turn_without_input_leaves_routing_to_its_previous_response(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + request = {"previous_response_id": "resp_1", "instructions": LONG_PROMPT} + + pinned = await _pinned_deployment_ids(request, None, request) + + assert len(pinned) == 2 + + +async def _serve_turn(check: PromptCachingDeploymentCheck, request_input: object) -> int: + metadata: dict[str, object] = {} + pinned = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + return len(pinned) + + +@pytest.mark.asyncio +async def test_responses_thread_longer_than_the_lookback_window_keeps_its_deployment(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + check = PromptCachingDeploymentCheck(cache=DualCache()) + turns = [ + {"role": "system", "content": LONG_PROMPT}, + *( + {"role": "assistant" if index % 2 else "user", "content": f"turn {index}"} + for index in range(2 * PROMPT_CACHE_LOOKBACK_POSITIONS) + ), + ] + + served = [await _serve_turn(check, turns[:end]) for end in range(2, len(turns) + 1, 2)] + + assert served[1:] == [1] * (len(served) - 1) + + +@pytest.mark.asyncio +async def test_responses_failover_writes_the_pin_for_the_deployment_that_answered(): + check = PromptCachingDeploymentCheck(cache=DualCache()) + request = {"input": _marked_responses_input()} + metadata: dict[str, object] = {} + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _read_responses_pin(check, {**request, "litellm_metadata": metadata}) + await _log_responses_success(check, _marked_responses_input(), metadata, model_id="dep-2") + + pinned = await _read_responses_pin(check, {**request, "litellm_metadata": {}}) + + assert pinned == ("dep-2",) + + +@pytest.mark.asyncio +async def test_responses_pin_covers_only_what_follows_the_newest_compaction(): + compaction = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "summary"}), + } + prefix = {"role": "user", "content": [_marked_input_text(LONG_PROMPT)]} + + first_input = [{"role": "user", "content": "dropped turn"}, compaction, prefix] + + pinned = await _pinned_deployment_ids( + {"input": first_input}, + first_input, + {"input": [{"role": "user", "content": "another dropped turn"}, compaction, prefix]}, + ) + + assert pinned == ("dep-1",) + + +@pytest.mark.asyncio +async def test_responses_image_in_the_prefix_pins_to_one_deployment(): + request_input = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": ONE_PIXEL_PNG}, _marked_input_text(LONG_PROMPT)], + }, + {"role": "user", "content": "hello"}, + ] + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [[{"role": "system", "content": LONG_PROMPT}, {"role": "user", "content": "hello"}], LONG_PROMPT], + ids=["item_list", "plain_string"], +) +async def test_responses_auto_cached_prefix_pins_to_one_deployment(monkeypatch, request_input: object): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + assert len(await _deployments_serving_repeats(_responses_router(RESPONSES_BRIDGED_MODEL), request_input)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_input", + [ + [LONG_PROMPT, LONG_PROMPT], + [{"role": "user", "content": 7}], + 7, + [{"type": "web_search_call"}], + [{"type": "web_search_call", "id": "ws_1", "status": "completed", "action": {"type": "search"}}], + [{"type": ["message"], "role": "user", "content": LONG_PROMPT}], + [{"role": "user", "content": [{"type": "input_image"}]}], + ], + ids=[ + "embedding_strings", + "non_text_item_content", + "integer", + "web_search_call_without_id", + "search_action_without_query", + "unhashable_item_type", + "image_without_a_source", + ], +) +async def test_input_that_is_not_a_responses_prompt_never_narrows_or_raises(request_input: object): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + metadata: dict[str, object] = {} + + filtered = await _read_responses_pin(check, {"input": request_input, "litellm_metadata": metadata}) + await _log_responses_success(check, request_input, metadata) + + assert filtered == ("dep-1", "dep-2") + assert not cache.in_memory_cache.cache_dict + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "next_tools,expected_pin", + [([{"type": "image_generation"}], ("dep-1",)), ([{"type": "web_search"}], ("dep-1", "dep-2"))], + ids=["same_native_tool", "other_native_tool"], +) +async def test_responses_native_tools_are_part_of_the_prefix( + next_tools: list[dict[str, object]], expected_pin: tuple[str, ...] +): + request_input = _marked_responses_input() + + pinned = await _pinned_deployment_ids( + {"input": request_input, "tools": [{"type": "image_generation"}]}, + request_input, + {"input": request_input, "tools": next_tools}, + ) + + assert pinned == expected_pin + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_tool_type", ["shell", "image_generation", "computer_use", "local_shell"]) +async def test_responses_native_tool_routed_to_a_native_model_logs_no_dropped_tool_warning( + caplog: pytest.LogCaptureFixture, native_tool_type: str +): + router = _responses_router(RESPONSES_NATIVE_MODEL) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + tools=[{"type": native_tool_type}], + mock_response="ok", + ) + + assert not [record for record in caplog.records if "Dropping Responses API tool" in record.getMessage()] + + +class _MetadataSerializingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.serialized_metadata: list[str] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.serialized_metadata.append(json.dumps(kwargs["litellm_params"]["metadata"])) + + +@pytest.mark.asyncio +async def test_logger_serializing_the_request_metadata_never_sees_the_prompt_cache_key(monkeypatch): + serializing_logger = _MetadataSerializingLogger() + monkeypatch.setattr(litellm, "callbacks", [serializing_logger]) + router = _responses_router(RESPONSES_NATIVE_MODEL) + + assert len(await _deployments_serving_repeats(router, _marked_responses_input())) == 1 + + assert await _eventually(lambda: len(serializing_logger.serialized_metadata) == 9) + logged_metadata = "".join(serializing_logger.serialized_metadata) + pin_keys = tuple(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert pin_keys + assert not [pin_key for pin_key in pin_keys if pin_key.split(":")[1] in logged_metadata] + + +@pytest.mark.asyncio +async def test_responses_call_hands_back_the_callers_metadata_without_the_pin_token(): + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + caller_metadata: Final[dict[str, object]] = {"tags": ["team-a"]} + + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + litellm_metadata=caller_metadata, + mock_response="ok", + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert "responses_prompt_cache_pin" not in caller_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failed_in_request,expected_deployments", + [({}, ("dep-2",)), ({"_excluded_deployment_ids": ["dep-2"]}, ("dep-1", "dep-2"))], + ids=["previous_deployment_healthy", "previous_deployment_failed_this_request"], +) +async def test_responses_continuation_turn_goes_to_the_deployment_that_answered_its_previous_response( + failed_in_request: dict[str, object], expected_deployments: tuple[str, ...] +): + previous_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id("anthropic", "dep-2", "resp_x") + request: dict[str, object] = { + "input": "hello", + "previous_response_id": previous_response_id, + "litellm_metadata": {}, + } + + pinned = await _read_responses_pin( + PromptCachingDeploymentCheck(cache=DualCache()), {**request, **failed_in_request} + ) + + assert pinned == expected_deployments + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_tool", ["get_time", None, 42], ids=["string", "null", "number"]) +async def test_responses_request_with_a_malformed_tool_is_left_to_the_provider(bad_tool: object): + request: dict[str, object] = {"input": _marked_responses_input(), "tools": [bad_tool], "litellm_metadata": {}} + + pinned = await _read_responses_pin(PromptCachingDeploymentCheck(cache=DualCache()), request) + + assert pinned == ("dep-1", "dep-2") diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 31fef208a8c..d285d6fa619 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -81,6 +81,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py", "litellm/responses/litellm_completion_transformation/transformation.py", "litellm/router_utils/client_initalization_utils.py", + "litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py", "litellm/router_utils/router_callbacks/track_deployment_metrics.py", "litellm/secret_managers/cyberark_secret_manager.py", "litellm/secret_managers/google_secret_manager.py",