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",