diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 66e2754d5ad8..f4ab5d081bec 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -114,6 +114,7 @@ def __init__( self.default_webhook_url = default_webhook_url self.flush_lock = asyncio.Lock() self.periodic_started = False + self._periodic_flush_task: asyncio.Task[None] | None = None self.hanging_request_check = AlertingHangingRequestCheck( slack_alerting_object=self, ) @@ -125,6 +126,12 @@ def __init__( self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) + def _ensure_periodic_flush_task(self) -> None: + if self.periodic_started and (self._periodic_flush_task is None or not self._periodic_flush_task.done()): + return + self._periodic_flush_task = asyncio.create_task(self.periodic_flush()) + self.periodic_started = True + def update_values( self, alerting: list | None = None, @@ -137,17 +144,14 @@ def update_values( ): if alerting is not None: self.alerting = alerting - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alerting_threshold is not None: self.alerting_threshold = alerting_threshold if alert_types is not None: self.alert_types = alert_types if alerting_args is not None: self.alerting_args = SlackAlertingArgs(**alerting_args) - if not self.periodic_started: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alert_type_config is not None: for key, val in alert_type_config.items(): self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val @@ -1442,9 +1446,8 @@ async def send_alert( return # Start periodic flush if not already started - if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + if self.alerting is not None and len(self.alerting) > 0: + self._ensure_periodic_flush_task() if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None: await self.send_webhook_alert(webhook_event=user_info) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 494d9e0935aa..0d6cbc2232e3 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -36,6 +36,7 @@ CacheControlMessageInjectionPoint, ) from litellm.types.llms.anthropic import ( + ANTHROPIC_TOOL_SEARCH_TOOL_TYPES, AllAnthropicToolsValues, AnthropicSystemMessageContent, ) @@ -124,6 +125,16 @@ def _carries_cache_breakpoint(block: object) -> bool: return isinstance(block, dict) and any(block.get(key) is not None for key in CACHE_BREAKPOINT_KEYS) +def _tool_carries_cache_breakpoint(tool: object) -> bool: + return _carries_cache_breakpoint(tool) or ( + isinstance(tool, dict) and _carries_cache_breakpoint(tool.get("function")) + ) + + +def _chat_transform_drops_tool_cache_control(tool: object) -> bool: + return isinstance(tool, dict) and tool.get("type") in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES + + def _accepts_prompt_cache_breakpoint(block: object) -> bool: return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES @@ -134,6 +145,8 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool: # rather than spending them on a list that is still missing some of their targets. CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_points" +EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints" + class AnthropicCacheControlHook(CustomPromptManagement): @staticmethod @@ -199,19 +212,13 @@ def get_chat_completion_prompt( # Create a deep copy of messages to avoid modifying the original list processed_messages = copy.deepcopy(messages) - # Separate message-level and non-message-level injection points - message_points: Final[list[CacheControlMessageInjectionPoint]] = [] - remaining_points: Final[list[CacheControlInjectionPoint]] = [] - for point in injection_points: - if point.get("location") == "message": - message_points.append(cast(CacheControlMessageInjectionPoint, point)) - else: - remaining_points.append(point) - - # Non-message points (currently Bedrock tool_config) are handled in the - # provider transform, where each tool_config point appends at most one - # cachePoint to the tools. That block also counts toward Anthropic's - # limit, so reserve a slot for it here to leave room. + message_points: Final = tuple( + cast(CacheControlMessageInjectionPoint, point) + for point in injection_points + if point.get("location") == "message" + ) + remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message") + stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect") openai_dialect: Final = ( stamped_dialect @@ -236,8 +243,10 @@ def get_chat_completion_prompt( if carry_unmatched else tuple(message_points) ) - reserved_blocks: Final = ( - 1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0 + stamped_external: Final = injection_points[0].get(EXTERNAL_BREAKPOINTS_STAMP) + external_breakpoints: Final = stamped_external if isinstance(stamped_external, int) else 0 + reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages( + remaining_points, external_breakpoints, openai_dialect ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) processed_messages = self._apply_message_injections( @@ -254,14 +263,19 @@ def get_chat_completion_prompt( # Points this pass did not place: non-message ones for the provider transform, and # the deferred role-targeted ones. Deferring is what reaches the Responses API's - # `instructions`, which is only a system message once the bridge builds one. The - # judged stamp is what makes it safe: the next pass must not re-judge points - # against messages this pass already marked (see `_should_stand_down`). - carried_points: Final[Sequence[CacheControlInjectionPoint]] = (*remaining_points, *carried_message_points) + # `instructions`, which is only a system message once the bridge builds one. A later + # pass re-applies them safely: a target that already carries a mark is skipped and + # the census counts every mark on the wire, litellm's own included. + carried_points: Final[Sequence[CacheControlInjectionPoint]] = ( + *AnthropicCacheControlHook._points_with_a_slot_left( + remaining_points, + AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) + external_breakpoints, + openai_dialect, + ), + *carried_message_points, + ) if carried_points: - non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged( - carried_points - ) + non_default_params["cache_control_injection_points"] = list(carried_points) return model, processed_messages, non_default_params @@ -296,6 +310,72 @@ def count_request_cache_breakpoints(messages: Iterable[object], system: object = ) return system_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages) + @staticmethod + def count_external_cache_breakpoints( + tools: Iterable[object] | None, cache_control: object = None, request_kwargs: object = None + ) -> int: + """Client breakpoints outside messages and system that the provider cap still counts. + + A tool carries its mark at the top level (Anthropic shape) or under ``function`` + (OpenAI shape). A top-level ``cache_control`` is Anthropic's automatic caching, + which places one breakpoint of its own on top of the explicit ones. The + ``extra_body`` envelope of ``request_kwargs`` is merged over the request on the + wire, so a ``tools`` or ``cache_control`` it carries replaces the direct value + and is counted in its place. Callers pass only the tools whose mark reaches the + provider on their path. + """ + extra_body: Final = ( + _validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {} + ) + wire_cache_control: Final = extra_body.get("cache_control", cache_control) + wire_tools: Final = _validated_object_list(extra_body["tools"]) if "tools" in extra_body else tools + tool_blocks: Final = sum(1 for tool in wire_tools or () if _tool_carries_cache_breakpoint(tool)) + envelope_blocks: Final = AnthropicCacheControlHook.count_request_cache_breakpoints( + _validated_object_list(extra_body.get("messages")) or (), extra_body.get("system") + ) + return int(wire_cache_control is not None) + tool_blocks + envelope_blocks + + @staticmethod + def count_external_cache_breakpoints_on_messages_route( + tools: Iterable[object] | None, cache_control: object, request_kwargs: object + ) -> int: + """The /v1/messages census before the route splits. + + The native messages transforms drop the ``extra_body`` envelope while the + chat bridge merges it, so the cap reserves for whichever census is larger + rather than letting an envelope that unmarks a direct tool free a slot the + provider still counts. + """ + return max( + AnthropicCacheControlHook.count_external_cache_breakpoints(tools, cache_control), + AnthropicCacheControlHook.count_external_cache_breakpoints(tools, cache_control, request_kwargs), + ) + + @staticmethod + def _blocks_reserved_outside_messages( + remaining_points: Sequence[CacheControlInjectionPoint], external_breakpoints: int, openai_dialect: bool + ) -> int: + """Slots of the provider cap that the message census cannot see. + + The client's breakpoints on tools and its automatic top-level ``cache_control`` + are already on the wire, and a ``tool_config`` point becomes one more cachePoint + in the Bedrock converse transform. OpenAI's cap counts only its own block markers. + """ + if openai_dialect: + return 0 + tool_config_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 + return external_breakpoints + tool_config_blocks + + @staticmethod + def _points_with_a_slot_left( + remaining_points: Sequence[CacheControlInjectionPoint], breakpoints_on_wire: int, openai_dialect: bool + ) -> tuple[CacheControlInjectionPoint, ...]: + """A ``tool_config`` point becomes a cachePoint the Bedrock converse transform never + counts against the cap, so it is forwarded only while the wire still has a slot.""" + if openai_dialect or breakpoints_on_wire < MAX_CACHE_CONTROL_BLOCKS: + return tuple(remaining_points) + return tuple(point for point in remaining_points if point.get("location") != "tool_config") + @staticmethod def _apply_message_injections( points: Sequence[CacheControlMessageInjectionPoint], @@ -476,11 +556,16 @@ def _system_block_with_breakpoint( def apply_to_anthropic_messages_request( messages: list[dict], system: str | list | None, - injection_points: list[CacheControlInjectionPoint], + injection_points: Sequence[CacheControlInjectionPoint], openai_dialect: bool = False, + external_breakpoints: int = 0, ) -> tuple[list[dict], str | list | None, list[CacheControlInjectionPoint]]: """Apply cache control injection for the Anthropic-native v1/messages endpoint. + ``external_breakpoints`` is the client's breakpoint count outside ``messages`` and + ``system`` (see ``count_external_cache_breakpoints``); it shrinks the budget so + the request never exceeds the provider cap. + Returns (messages, system, remaining_non_message_points). """ if not injection_points: @@ -489,22 +574,17 @@ def apply_to_anthropic_messages_request( processed_messages: list[dict] = copy.deepcopy(messages) processed_system = copy.deepcopy(system) if system is not None else None - message_points: Final[list[CacheControlMessageInjectionPoint]] = [] - system_points: Final[list[CacheControlMessageInjectionPoint]] = [] - remaining_points: Final[list[CacheControlInjectionPoint]] = [] - - for point in injection_points: - if point.get("location") == "message": - msg_point = cast(CacheControlMessageInjectionPoint, point) - if msg_point.get("role") == "system": - system_points.append(msg_point) - else: - message_points.append(msg_point) - else: - remaining_points.append(point) - - reserved_blocks: Final = ( - 1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0 + role_points: Final = tuple( + cast(CacheControlMessageInjectionPoint, point) + for point in injection_points + if point.get("location") == "message" + ) + system_points: Final = tuple(point for point in role_points if point.get("role") == "system") + message_points: Final = tuple(point for point in role_points if point.get("role") != "system") + remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message") + + reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages( + remaining_points, external_breakpoints, openai_dialect ) max_blocks: Final = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks @@ -541,8 +621,14 @@ def apply_to_anthropic_messages_request( max_blocks=max_blocks - system_blocks, openai_dialect=openai_dialect, ) + forwarded_points: Final = AnthropicCacheControlHook._points_with_a_slot_left( + remaining_points, + AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages, processed_system) + + external_breakpoints, + openai_dialect, + ) - return processed_messages, processed_system, remaining_points + return processed_messages, processed_system, list(forwarded_points) @staticmethod def _default_control() -> ChatCompletionCachedContent: @@ -559,31 +645,26 @@ def _default_control() -> ChatCompletionCachedContent: return ChatCompletionCachedContent(type="ephemeral") @staticmethod - def _stamped_as_judged(points: Sequence[CacheControlInjectionPoint]) -> Sequence[Mapping[str, object]]: - """Mark written-back points as having passed the client cache_control judgment. - - Builds copies because config-owned point dicts are shared across - requests; mutating them would leak the stamp into future requests. - """ - return AnthropicCacheControlHook._stamped(points, "_litellm_judged", True) - - @staticmethod - def _judged_configured_points( + def _stamped_for_prompt_hook( points: Sequence[CacheControlInjectionPoint], - messages: list[AllMessageValues], - tools: list[object] | None, - cache_control: object, + external_breakpoints: int, model: str, custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, - request_kwargs: object, - ) -> Sequence[Mapping[str, object]] | None: - if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control, request_kwargs): - return None - return AnthropicCacheControlHook._stamped_with_dialect( + ) -> Sequence[Mapping[str, object]]: + """Carry onto the points what the prompt-management hook never receives. + + The hook sees neither the tools nor the request kwargs, so the target dialect + and the client's breakpoint count outside the message list ride on the points. + Builds copies because config-owned point dicts are shared across requests. + """ + with_dialect: Final = AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options ) + if external_breakpoints == 0: + return with_dialect + return AnthropicCacheControlHook._stamped(with_dialect, EXTERNAL_BREAKPOINTS_STAMP, external_breakpoints) @staticmethod def _stamped_with_dialect( @@ -604,35 +685,9 @@ def _stamped_with_dialect( ) @staticmethod - def _stamped( - points: Sequence[CacheControlInjectionPoint], key: str, value: object - ) -> Sequence[Mapping[str, object]]: + def _stamped(points: Sequence[Mapping[str, object]], key: str, value: object) -> Sequence[Mapping[str, object]]: return [{**point, key: value} for point in points] - @staticmethod - def _should_stand_down( - points: Sequence[CacheControlInjectionPoint], - messages: list[AllMessageValues], - system: str | list | None, - tools: list | None, - cache_control: object = None, - request_kwargs: object = None, - ) -> bool: - """Whether configured injection points must yield to client-set cache_control. - - Points that a prior pass over this request already judged and wrote - back carry the internal judged stamp; any re-entry (acompletion - re-entering completion, the async-to-sync /v1/messages dispatch, - interceptor sub-calls reusing the request kwargs) must not re-judge - them, because by then the messages carry litellm's own injected marks - and the judgment would misread those as client breakpoints. - """ - if all(point.get("_litellm_judged") for point in points): - return False - return AnthropicCacheControlHook._request_has_cache_control( - messages, system, tools, cache_control, request_kwargs - ) - @staticmethod def _request_has_cache_control( messages: list[AllMessageValues], @@ -641,27 +696,18 @@ def _request_has_cache_control( cache_control: object = None, request_kwargs: object = None, ) -> bool: - """Client breakpoints own caching in both the request and its extra_body envelope.""" - bodies: Final = ( - {"messages": messages, "system": system, "tools": tools, "cache_control": cache_control}, - _validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {}, - ) - return any( - body.get("cache_control") is not None - or AnthropicCacheControlHook.count_request_cache_breakpoints( - _validated_object_list(body.get("messages")) or (), body.get("system") - ) - > 0 - or any( - AnthropicCacheControlHook._request_value(tool, "cache_control") is not None - or AnthropicCacheControlHook._request_value( - AnthropicCacheControlHook._request_value(tool, "function"), "cache_control" - ) - is not None - for tool in (_validated_object_list(body.get("tools")) or ()) - ) - for body in bodies - ) + """Return True if the request already carries any client-supplied cache_control. + + Only the automatic defaults stand down on it: a client that marks its own + breakpoints (Claude Code does) has a caching strategy the defaults would + clash with, whether the marks sit in the request or in its ``extra_body`` + envelope. Configured injection points are an explicit instruction and are + applied alongside the client's marks, bounded by the provider cap. + """ + return ( + AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) + + AnthropicCacheControlHook.count_external_cache_breakpoints(tools, cache_control, request_kwargs) + ) > 0 @staticmethod def get_default_injection_points( @@ -769,34 +815,30 @@ def maybe_seed_default_injection_points( ) -> None: """For /chat/completions: resolve the injection points the request should carry. - Configured injection points win over the automatic defaults, but stand - down entirely when the client already marked its own cache_control - breakpoints (messages or tools): injecting alongside them clashes with - the client's caching strategy and can exceed the provider's four-block - limit. The judgment happens once per request; points a prior pass - wrote back carry the judged stamp and are never re-judged (see - ``_should_stand_down``). Seeding the param lets the existing - prompt-management gate and the AnthropicCacheControlHook run - unchanged. + Configured injection points win over the automatic defaults and are applied + even when the client marked its own cache_control elsewhere in the request; + the provider's four-block cap bounds them, counting the client's marks on + messages, tools and the top-level ``cache_control``. Only the defaults stand + down on client marks. Seeding the param lets the existing prompt-management + gate and the AnthropicCacheControlHook run unchanged. """ import litellm - if non_default_params.get("cache_control_injection_points"): - judged: Final = AnthropicCacheControlHook._judged_configured_points( - non_default_params["cache_control_injection_points"], - messages, - tools, - non_default_params.get("cache_control"), + configured: Final = non_default_params.get("cache_control_injection_points") + if configured: + tools_keeping_marks: Final = tuple( + tool for tool in tools or () if not _chat_transform_drops_tool_cache_control(tool) + ) + non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_for_prompt_hook( + configured, + AnthropicCacheControlHook.count_external_cache_breakpoints( + tools_keeping_marks, non_default_params.get("cache_control"), non_default_params + ), model, custom_llm_provider, api_base, non_default_params.get("prompt_cache_options"), - non_default_params, ) - if judged is None: - non_default_params.pop("cache_control_injection_points") - else: - non_default_params["cache_control_injection_points"] = judged return points: Final = AnthropicCacheControlHook.get_default_injection_points( messages=messages, @@ -897,15 +939,14 @@ def maybe_inject_cache_control( ) -> tuple[list[dict], str | list | None]: """Extract cache_control_injection_points from kwargs and apply if present. - Configured points stand down entirely when the client already marked - its own cache_control breakpoints anywhere in the request. The - judgment happens once per request; points a prior pass wrote back - carry the judged stamp and are never re-judged (see - ``_should_stand_down``). When none are configured but + Configured points are applied even when the client marked its own + cache_control elsewhere in the request, bounded by the provider cap, + which counts the client's marks on messages, system, tools and the + top-level ``cache_control``. When none are configured but ``litellm.enable_anthropic_prompt_caching`` or the per-request ``enable_prompt_caching`` kwarg (stamped from key metadata) is on, - synthesize default breakpoints for the native /v1/messages path. Pops - both keys from kwargs; + synthesize default breakpoints for the native /v1/messages path; those + defaults alone stand down on client marks. Pops both keys from kwargs; if remaining (non-message) points exist they are written back so downstream transforms can handle them. """ @@ -917,13 +958,8 @@ def maybe_inject_cache_control( configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) - if configured and AnthropicCacheControlHook._should_stand_down( - configured, typed_messages, system, tools, cache_control, kwargs - ): - return messages, system - injection_points: list[CacheControlInjectionPoint] = configured or [] - if not injection_points and model is not None: - injection_points = AnthropicCacheControlHook.get_default_injection_points( + injection_points: Final[Sequence[CacheControlInjectionPoint]] = configured or ( + AnthropicCacheControlHook.get_default_injection_points( messages=typed_messages, system=system, tools=tools, @@ -933,6 +969,9 @@ def maybe_inject_cache_control( cache_control=cache_control, request_kwargs=kwargs, ) + if model is not None + else () + ) if not injection_points: return messages, system @@ -945,6 +984,9 @@ def maybe_inject_cache_control( system=system, injection_points=injection_points, openai_dialect=openai_dialect, + external_breakpoints=AnthropicCacheControlHook.count_external_cache_breakpoints_on_messages_route( + tools, cache_control, kwargs + ), ) breakpoints_added: Final = ( AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) - breakpoints_before @@ -953,7 +995,7 @@ def maybe_inject_cache_control( if openai_dialect and breakpoints_added > 0: kwargs.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit")) if remaining: - kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining) + kwargs["cache_control_injection_points"] = remaining return messages, system @property diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6c1b7946394c..bf37b1be2e45 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -46,6 +46,8 @@ AllMessageValues, ChatCompletionDocumentObject, ChatCompletionNamedToolChoiceParam, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ChatCompletionToolParam, OpenAIMessageContentListBlock, ) @@ -854,6 +856,8 @@ def _count_content_list( content_list: str | Iterable[ OpenAIMessageContentListBlock + | ChatCompletionThinkingBlock + | ChatCompletionRedactedThinkingBlock | AnthropicMessagesTextParam | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam @@ -898,9 +902,9 @@ def _count_content_list( use_default_image_token_count, default_token_count, ) - elif c["type"] == "thinking": + elif c["type"] in ("thinking", "redacted_thinking"): # Claude extended thinking content block - # Count the thinking text and skip signature (opaque signature blob) + # Count the thinking text and skip the opaque blobs (signature, redacted data) thinking_text = str(c.get("thinking", "")) if thinking_text: num_tokens += count_function(thinking_text) @@ -920,7 +924,8 @@ def _count_content_list( raise ValueError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " - f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)." + f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " + f"tool_reference)." ) return num_tokens except Exception as e: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 5326cf3415f4..01107c435dfb 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -959,7 +959,7 @@ async def get_mcp_tools( mcp_server_auth_headers=None, ) tools: Final = listing.tools - dumped_tools: Final = [dict(tool) for tool in tools] + dumped_tools: Final = [tool.model_dump(by_alias=True) for tool in tools] return {"tools": dumped_tools} diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 16e8ac93d59d..bff27a403538 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -68,16 +68,6 @@ async def create_mcp_list_tools_events( # Use the pre-processed MCP tools that were already fetched, filtered, and deduplicated by the parent filtered_mcp_tools: Final = pre_processed_mcp_tools - # Convert tools to dict format for the event - _mcp_tools_dict: Final = [ - tool.model_dump() - if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump", None)) - else tool.__dict__ - if hasattr(tool, "__dict__") - else {"name": getattr(tool, "name", str(tool))} - for tool in filtered_mcp_tools - ] - # Emit list tools completed event completed_event: Final = MCPListToolsCompletedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED, diff --git a/litellm/router.py b/litellm/router.py index 300b069a4641..8eca83f9d5e5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -117,6 +117,7 @@ MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler @@ -1350,6 +1351,9 @@ def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: `_init_routing_groups`) so repeated `update_settings` calls don't accumulate dead selectors that keep receiving callback events. """ + for selector in selectors: + if isinstance(selector, BaseRoutingStrategy): + selector.retire() selector_ids: Final = {id(s) for s in selectors if s is not None} if not selector_ids: return @@ -12033,7 +12037,7 @@ def update_settings(self, **kwargs): ) rebuild_routing_groups = True elif var == "routing_strategy_args": - routing_args_updated = True + routing_args_updated = value != self.routing_strategy_args setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 686d57e2b777..79d457f28369 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -40,10 +40,24 @@ def setup_sync_task(self, default_sync_interval: float | None): self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval) ) + def cancel_sync_task(self) -> None: + if self._sync_task is not None: + self._sync_task.cancel() + + def retire(self) -> None: + self.cancel_sync_task() + if not self.redis_increment_operation_queue: + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + loop.create_task(self._push_in_memory_increments_to_redis()) + async def cleanup(self): """Cleanup method to be called when shutting down""" if self._sync_task is not None: - self._sync_task.cancel() + self.cancel_sync_task() try: await self._sync_task except asyncio.CancelledError: diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 8794ff2db957..666b010c7426 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -60,7 +60,7 @@ def matches(self, context: Context) -> bool: RULES: Final[Rules] = ( Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), Rule(Route.OCR, Rollout.RUST_OPT_OUT), - Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), + Rule(Route.MESSAGES, Rollout.PYTHON_ONLY), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/types/integrations/anthropic_cache_control_hook.py b/litellm/types/integrations/anthropic_cache_control_hook.py index ef414f22c3b7..20e7885a2bf2 100644 --- a/litellm/types/integrations/anthropic_cache_control_hook.py +++ b/litellm/types/integrations/anthropic_cache_control_hook.py @@ -17,8 +17,8 @@ class CacheControlMessageInjectionPoint(TypedDict): role: Literal["user", "system", "assistant"] | None # Optional: target by role (user, system, assistant) index: int | str | None # Optional: target by specific index control: ChatCompletionCachedContent | None - _litellm_judged: NotRequired[bool] # Internal: written back by litellm once the client cache_control judgment ran _litellm_openai_dialect: NotRequired[ReadOnly[bool]] + _litellm_external_breakpoints: NotRequired[ReadOnly[int]] class CacheControlToolConfigInjectionPoint(TypedDict): @@ -26,8 +26,8 @@ class CacheControlToolConfigInjectionPoint(TypedDict): location: Literal["tool_config"] control: ChatCompletionCachedContent | None - _litellm_judged: NotRequired[bool] # Internal: written back by litellm once the client cache_control judgment ran _litellm_openai_dialect: NotRequired[ReadOnly[bool]] + _litellm_external_breakpoints: NotRequired[ReadOnly[int]] CacheControlInjectionPoint = CacheControlMessageInjectionPoint | CacheControlToolConfigInjectionPoint diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bcd24695f252..43a7b0e0e9cc 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -753,6 +753,10 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) ANTHROPIC_TOOL_SEARCH_BETA_HEADER: Final = "advanced-tool-use-2025-11-20" +ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: Final = frozenset( + {"tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119"} +) + # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 4bc6c08bd630..2be75b43eea2 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -434,3 +434,29 @@ async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch): alert_type=AlertType.budget_alerts, alerting_metadata={}, ) + + +def _periodic_flush_tasks() -> list[asyncio.Task[object]]: + return [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None and t.get_coro().__qualname__ == "SlackAlerting.periodic_flush" + ] + + +@pytest.mark.asyncio +async def test_update_values_repeated_alerting_reload_keeps_single_periodic_flush_task() -> None: + slack_alerting: Final = SlackAlerting(alerting=["slack"]) + try: + for _ in range(5): + slack_alerting.update_values(alerting=["slack"]) + await asyncio.sleep(0) + flush_tasks: Final = _periodic_flush_tasks() + assert len(flush_tasks) == 1, f"expected 1 periodic_flush task, found {len(flush_tasks)}" + finally: + for t in _periodic_flush_tasks(): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 83649c3386ad..7bf4533979af 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -4,10 +4,11 @@ import subprocess import sys import textwrap -from typing import List, Optional, Tuple +from typing import Final, List, Optional, Tuple from unittest.mock import MagicMock, patch import pytest +from pydantic import BaseModel, ConfigDict import litellm from litellm.integrations.anthropic_cache_control_hook import ( @@ -1276,11 +1277,7 @@ def test_cache_control_hook_reserves_slot_for_tool_config_point(): ) assert _count_cache_control(processed) == 3 - # The tool_config point is passed through for the provider transform, - # stamped so re-entries never re-judge it against litellm's own marks. - assert non_default_params["cache_control_injection_points"] == [ - {"location": "tool_config", "_litellm_judged": True} - ] + assert non_default_params["cache_control_injection_points"] == [{"location": "tool_config"}] @pytest.mark.asyncio @@ -1338,18 +1335,8 @@ async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(mo client=client, ) - request_body = json.loads(mock_post.call_args.kwargs["data"]) - - cache_points = sum( - 1 for block in request_body.get("system", []) if isinstance(block, dict) and "cachePoint" in block - ) - for msg in request_body.get("messages", []): - content = msg.get("content", []) - if isinstance(content, list): - cache_points += sum(1 for block in content if isinstance(block, dict) and "cachePoint" in block) - for tool in request_body.get("toolConfig", {}).get("tools", []): - if isinstance(tool, dict) and "cachePoint" in tool: - cache_points += 1 + request_body = _ConverseBody.model_validate_json(mock_post.call_args.kwargs["data"]) + cache_points = _count_converse_cache_points(request_body) assert cache_points <= 4, ( f"Bedrock payload exceeded Anthropic's 4 cache_control block limit " @@ -1357,6 +1344,97 @@ async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(mo ) +class _ConverseMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + content: tuple[dict[str, object], ...] = () + + +class _ConverseToolConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + tools: tuple[dict[str, object], ...] = () + + +class _ConverseBody(BaseModel): + model_config = ConfigDict(frozen=True) + + system: tuple[dict[str, object], ...] = () + messages: tuple[_ConverseMessage, ...] = () + toolConfig: _ConverseToolConfig = _ConverseToolConfig() + + +def _count_converse_cache_points(request_body: _ConverseBody) -> int: + blocks: Final = ( + *request_body.system, + *(block for message in request_body.messages for block in message.content), + *request_body.toolConfig.tools, + ) + return sum(1 for block in blocks if "cachePoint" in block) + + +@pytest.mark.asyncio +async def test_cache_control_hook_bedrock_tool_config_point_stands_down_when_client_marks_fill_the_cap( + monkeypatch: pytest.MonkeyPatch, +): + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "fake_access_key_id", + "AWS_SECRET_ACCESS_KEY": "fake_secret_access_key", + "AWS_REGION_NAME": "us-east-1", + }, + ): + monkeypatch.setattr(litellm, "callbacks", [AnthropicCacheControlHook()]) + + mock_response = MagicMock() + mock_response.json.return_value = { + "output": {"message": {"role": "assistant", "content": "ok"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104}, + } + mock_response.status_code = 200 + + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=mock_response) as mock_post: + marked = {"type": "ephemeral"} + messages = [ + {"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": marked}]}, + *( + {"role": "user", "content": [{"type": "text", "text": f"turn {i}", "cache_control": marked}]} + for i in range(3) + ), + {"role": "user", "content": "What is the weather?"}, + ] + + await litellm.acompletion( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + max_tokens=32, + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + cache_control_injection_points=[{"location": "tool_config"}], + client=client, + ) + + request_body = _ConverseBody.model_validate_json(mock_post.call_args.kwargs["data"]) + + assert _count_converse_cache_points(request_body) == 4 + assert not any("cachePoint" in tool for tool in request_body.toolConfig.tools) + + class TestApplyToAnthropicMessagesRequest: """Tests for apply_to_anthropic_messages_request (v1/messages cache control).""" @@ -1683,13 +1761,17 @@ def test_new_transports_preserve_client_controls(self, monkeypatch, local_model_ result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( messages, system, kwargs, model, provider, tools=tools, ) - if client_control != "none": + if client_control != "none" and not configured: assert (result_messages, result_system, tools) == original assert kwargs["metadata"] == {} else: assert kwargs["metadata"]["litellm_gateway_injected_cache"] == "selected-deployment" assert sum(AnthropicCacheControlHook._count_cache_control_blocks(m) for m in result_messages) == 1 assert result_system[0]["cache_control"] == control + assert result_messages[-1]["content"][-1]["cache_control"] == control + assert tools == original[2] + assert (result_messages == original[0]) == (envelope == "request" and client_control == "message") + assert (result_system == original[1]) == (envelope == "request" and client_control == "system") if provider == "vertex_ai": wire = VertexAIAnthropicConfig().transform_request( model=model, messages=[{"role": "system", "content": result_system}, *result_messages], @@ -1706,7 +1788,7 @@ def test_new_transports_preserve_client_controls(self, monkeypatch, local_model_ AnthropicCacheControlHook.maybe_seed_default_injection_points( seeded, [{"role": "system", "content": original[1]}, *original[0]], model, provider, tools=tools, ) - assert bool(seeded.get("cache_control_injection_points")) == (client_control == "none") + assert bool(seeded.get("cache_control_injection_points")) == (client_control == "none" or configured) @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @@ -2257,13 +2339,11 @@ def test_v1_messages_is_noop_when_disabled(self): assert result_msgs == messages -class TestConfiguredInjectionPointsStandDown: - """Configured cache_control_injection_points must stand down entirely when the - client already set its own cache_control anywhere in the request (LIT-4582); - injecting alongside client breakpoints clashes with the client's caching - strategy and can push the request past Anthropic's four-block limit.""" - +class TestConfiguredInjectionPointsSurviveClientMarks: CONFIGURED = [{"location": "message", "role": "system"}] + TAIL_POINT = [{"location": "message", "index": -1}] + TOOL_CONFIG_POINT = [{"location": "tool_config"}] + EPHEMERAL = {"type": "ephemeral"} CLEAN_MESSAGES: List[AllMessageValues] = [ {"role": "system", "content": "sys"}, @@ -2277,6 +2357,37 @@ class TestConfiguredInjectionPointsStandDown: V1_MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + MARKED_TOOL_TOP_LEVEL = { + "type": "function", + "function": {"name": "t", "parameters": {}}, + "cache_control": {"type": "ephemeral"}, + } + MARKED_TOOL_NESTED = { + "type": "function", + "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}, + } + UNMARKED_TOOL = {"type": "function", "function": {"name": "t", "parameters": {}}} + MARKED_V1_TOOL = {"name": "t", "input_schema": {}, "cache_control": {"type": "ephemeral"}} + UNMARKED_V1_TOOL = {"name": "t", "input_schema": {}} + MARKED_SYSTEM = [{"type": "text", "text": "sys", "cache_control": EPHEMERAL}] + MARKED_TOOL_SEARCH_REGEX = { + "type": "tool_search_tool_regex_20251119", + "name": "tool_search", + "cache_control": {"type": "ephemeral"}, + } + MARKED_TOOL_SEARCH_BM25 = { + "type": "tool_search_tool_bm25_20251119", + "name": "tool_search", + "cache_control": {"type": "ephemeral"}, + } + + @staticmethod + def _marked_user_turns(count: int) -> List[AllMessageValues]: + return [ + {"role": "user", "content": [{"type": "text", "text": f"turn {i}", "cache_control": {"type": "ephemeral"}}]} + for i in range(count) + ] + def _seed(self, params, messages, tools=None): AnthropicCacheControlHook.maybe_seed_default_injection_points( non_default_params=params, @@ -2286,6 +2397,17 @@ def _seed(self, params, messages, tools=None): tools=tools, ) + def _chat(self, params: dict[str, object], messages: List[AllMessageValues]) -> List[AllMessageValues]: + _, processed, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model="claude-sonnet-4-5", + messages=messages, + non_default_params=params, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + return processed + def _inject(self, messages, kwargs, system="sys", tools=None): return AnthropicCacheControlHook.maybe_inject_cache_control( messages, @@ -2296,23 +2418,79 @@ def _inject(self, messages, kwargs, system="sys", tools=None): tools=tools, ) - def test_configured_points_dropped_when_messages_carry_cache_control(self): + def test_chat_tail_point_applies_when_client_marked_the_system_block(self): + messages: List[AllMessageValues] = [ + {"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "history"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "question"}, + ] + params = {"cache_control_injection_points": copy.deepcopy(self.TAIL_POINT)} + self._seed(params, messages) + processed = self._chat(params, messages) + assert processed[0] == messages[0] + assert processed[-1] == {"role": "user", "content": "question", "cache_control": self.EPHEMERAL} + assert _count_cache_control(processed) == 2 + + def test_chat_configured_points_apply_when_messages_carry_cache_control(self): params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} self._seed(params, copy.deepcopy(self.MARKED_MESSAGES)) - assert "cache_control_injection_points" not in params + processed = self._chat(params, copy.deepcopy(self.MARKED_MESSAGES)) + assert processed[0] == {"role": "system", "content": "sys", "cache_control": self.EPHEMERAL} + assert processed[1] == self.MARKED_MESSAGES[1] @pytest.mark.parametrize( - "tool", - [ - {"type": "function", "function": {"name": "t", "parameters": {}}, "cache_control": {"type": "ephemeral"}}, - {"type": "function", "function": {"name": "t", "parameters": {}, "cache_control": {"type": "ephemeral"}}}, - ], - ids=["top_level", "nested_in_function"], + "tool", [MARKED_TOOL_TOP_LEVEL, MARKED_TOOL_NESTED], ids=["top_level", "nested_in_function"] ) - def test_configured_points_dropped_when_tools_carry_cache_control(self, tool): + def test_chat_configured_points_apply_when_tools_carry_cache_control(self, tool): params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES), tools=[tool]) - assert "cache_control_injection_points" not in params + processed = self._chat(params, copy.deepcopy(self.CLEAN_MESSAGES)) + assert processed[0] == {"role": "system", "content": "sys", "cache_control": self.EPHEMERAL} + + @pytest.mark.parametrize( + "tool,injected", + [(MARKED_TOOL_TOP_LEVEL, 0), (MARKED_TOOL_NESTED, 0), (UNMARKED_TOOL, 1)], + ids=["marked_top_level", "marked_nested_in_function", "unmarked"], + ) + def test_chat_cap_counts_client_marked_tools(self, tool, injected): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(3)] + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + self._seed(params, copy.deepcopy(messages), tools=[tool]) + processed = self._chat(params, copy.deepcopy(messages)) + assert _count_cache_control(processed) == 3 + injected + + @pytest.mark.parametrize("tool", [MARKED_TOOL_SEARCH_REGEX, MARKED_TOOL_SEARCH_BM25], ids=["regex", "bm25"]) + def test_chat_cap_ignores_marked_tool_search_tools(self, tool): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(3)] + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + self._seed(params, copy.deepcopy(messages), tools=[tool]) + processed = self._chat(params, copy.deepcopy(messages)) + assert _count_cache_control(processed) == 4 + + @pytest.mark.parametrize("marked_turns,forwarded", [(3, ["tool_config"]), (4, [])], ids=["slot_left", "cap_full"]) + def test_chat_forwards_tool_config_point_only_while_a_slot_is_left(self, marked_turns, forwarded): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(marked_turns)] + params = {"cache_control_injection_points": copy.deepcopy(self.TOOL_CONFIG_POINT)} + self._seed(params, copy.deepcopy(messages), tools=[self.UNMARKED_TOOL]) + self._chat(params, copy.deepcopy(messages)) + assert [p["location"] for p in params.get("cache_control_injection_points", [])] == forwarded + + @pytest.mark.parametrize("marked_turns,forwarded", [(3, ["tool_config"]), (4, [])], ids=["slot_left", "cap_full"]) + def test_v1_messages_forwards_tool_config_point_only_while_a_slot_is_left(self, marked_turns, forwarded): + kwargs = {"cache_control_injection_points": copy.deepcopy(self.TOOL_CONFIG_POINT)} + self._inject(self._marked_user_turns(marked_turns), kwargs, tools=[self.UNMARKED_V1_TOOL]) + assert [p["location"] for p in kwargs.get("cache_control_injection_points", [])] == forwarded + + @pytest.mark.parametrize("marked_turns,injected", [(2, 1), (3, 0)]) + def test_chat_root_cache_control_reserves_a_slot(self, marked_turns, injected): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(marked_turns)] + root_cache_control = {"type": "ephemeral"} + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), "cache_control": root_cache_control} + self._seed(params, copy.deepcopy(messages)) + processed = self._chat(params, copy.deepcopy(messages)) + assert _count_cache_control(processed) == marked_turns + injected + assert params["cache_control"] is root_cache_control def test_configured_points_kept_when_request_is_unmarked(self): configured = copy.deepcopy(self.CONFIGURED) @@ -2320,43 +2498,59 @@ def test_configured_points_kept_when_request_is_unmarked(self): self._seed(params, copy.deepcopy(self.CLEAN_MESSAGES)) assert params["cache_control_injection_points"] is configured - def test_judged_remainder_survives_reentry_despite_injected_marks(self): - """acompletion() re-enters completion() after injection ran, with only the - stamped non-message points written back; the re-entry must not misread - litellm's own marks as client ones and drop that remainder.""" - remainder = [{"location": "tool_config", "_litellm_judged": True}] - params = {"cache_control_injection_points": remainder} - self._seed(params, copy.deepcopy(self.MARKED_MESSAGES)) - assert params["cache_control_injection_points"] is remainder - - def test_v1_messages_stand_down_when_content_block_marked(self): + def test_chat_reentry_over_injected_messages_adds_no_duplicate_marks(self): + points = [{"location": "message", "role": "system"}, {"location": "tool_config"}] + first_params = {"cache_control_injection_points": copy.deepcopy(points)} + self._seed(first_params, copy.deepcopy(self.MARKED_MESSAGES)) + first = self._chat(first_params, copy.deepcopy(self.MARKED_MESSAGES)) + assert _count_cache_control(first) == 2 + assert first_params["cache_control_injection_points"] == [{"location": "tool_config"}] + + second_params = {"cache_control_injection_points": copy.deepcopy(points)} + self._seed(second_params, copy.deepcopy(first)) + second = self._chat(second_params, copy.deepcopy(first)) + assert second == first + assert second_params["cache_control_injection_points"] == [{"location": "tool_config"}] + + def test_v1_messages_configured_point_applies_when_content_block_marked(self): messages = [ {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]} ] kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} result_msgs, result_sys = self._inject(copy.deepcopy(messages), kwargs) assert result_msgs == messages - assert result_sys == "sys" + assert result_sys == [{"type": "text", "text": "sys", "cache_control": self.EPHEMERAL}] assert "cache_control_injection_points" not in kwargs - def test_v1_messages_stand_down_when_system_block_marked(self): - """A configured point targeting a message must not fire when the client - marked the system prompt; the old behavior injected into the message - because only the exact targeted position was guarded.""" + def test_v1_messages_tail_point_applies_when_system_block_marked(self): system = [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}] - kwargs = {"cache_control_injection_points": [{"location": "message", "role": "user"}]} + kwargs = {"cache_control_injection_points": copy.deepcopy(self.TAIL_POINT)} result_msgs, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs, system=system) - assert result_msgs == self.V1_MESSAGES + assert result_msgs == [ + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": self.EPHEMERAL}]} + ] assert result_sys == system - assert "cache_control_injection_points" not in kwargs - def test_v1_messages_stand_down_when_tools_marked(self): - tools = [{"name": "t", "input_schema": {}, "cache_control": {"type": "ephemeral"}}] + def test_v1_messages_configured_point_applies_when_tools_marked(self): kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} - result_msgs, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs, tools=tools) + result_msgs, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs, tools=[self.MARKED_V1_TOOL]) assert result_msgs == self.V1_MESSAGES - assert result_sys == "sys" - assert "cache_control_injection_points" not in kwargs + assert result_sys == [{"type": "text", "text": "sys", "cache_control": self.EPHEMERAL}] + + @pytest.mark.parametrize( + "tool,expected_system", + [ + (MARKED_V1_TOOL, "sys"), + (MARKED_TOOL_SEARCH_REGEX, "sys"), + (MARKED_TOOL_SEARCH_BM25, "sys"), + (UNMARKED_V1_TOOL, [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]), + ], + ids=["marked", "marked_tool_search_regex", "marked_tool_search_bm25", "unmarked"], + ) + def test_v1_messages_cap_counts_client_marked_tools(self, tool, expected_system): + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} + _, result_sys = self._inject(self._marked_user_turns(3), kwargs, tools=[tool]) + assert result_sys == expected_system def test_v1_messages_configured_points_apply_when_unmarked(self): kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED)} @@ -2364,16 +2558,73 @@ def test_v1_messages_configured_points_apply_when_unmarked(self): assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] @pytest.mark.parametrize( - "configured", - [None, CONFIGURED], - ids=["automatic_defaults", "configured_points"], + "extra_body,injected", + [ + ({"tools": [MARKED_TOOL_TOP_LEVEL]}, 0), + ({"cache_control": {"type": "ephemeral"}}, 0), + ({"tools": [UNMARKED_TOOL]}, 1), + ], + ids=["marked_tool", "root_cache_control", "unmarked_tool"], ) - def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured): + def test_chat_cap_counts_client_marks_sent_through_extra_body(self, extra_body, injected): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(3)] + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), "extra_body": extra_body} + self._seed(params, copy.deepcopy(messages)) + processed = self._chat(params, copy.deepcopy(messages)) + assert _count_cache_control(processed) == 3 + injected + + @pytest.mark.parametrize( + "extra_body,expected_system", + [ + ({"cache_control": {"type": "ephemeral"}}, "sys"), + ({"tools": [MARKED_V1_TOOL]}, "sys"), + ({"tools": [UNMARKED_V1_TOOL]}, [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]), + ], + ids=["root_cache_control", "marked_tool", "unmarked_tool"], + ) + def test_v1_messages_cap_counts_client_marks_sent_through_extra_body(self, extra_body, expected_system): + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), "extra_body": extra_body} + _, result_sys = self._inject(self._marked_user_turns(3), kwargs) + assert result_sys == expected_system + + @pytest.mark.parametrize( + "params,tools,marked_turns,injected", + [ + ({"extra_body": {"tools": [MARKED_TOOL_TOP_LEVEL]}}, [MARKED_TOOL_TOP_LEVEL], 2, 1), + ({"extra_body": {"tools": [UNMARKED_TOOL]}}, [MARKED_TOOL_TOP_LEVEL], 3, 1), + ({"extra_body": {"tools": [MARKED_TOOL_TOP_LEVEL]}}, [UNMARKED_TOOL], 3, 0), + ({"extra_body": {"cache_control": EPHEMERAL}, "cache_control": EPHEMERAL}, None, 2, 1), + ], + ids=["same_marked_tool_both_ways", "extra_body_unmarks", "extra_body_marks", "root_cache_control_both_ways"], + ) + def test_chat_cap_counts_extra_body_fields_in_place_of_the_direct_ones(self, params, tools, marked_turns, injected): + messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(marked_turns)] + params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), **copy.deepcopy(params)} + self._seed(params, copy.deepcopy(messages), tools=tools) + processed = self._chat(params, copy.deepcopy(messages)) + assert _count_cache_control(processed) == marked_turns + injected + + @pytest.mark.parametrize( + "kwargs,tools,marked_turns,expected_system", + [ + ({"extra_body": {"tools": [MARKED_V1_TOOL]}}, [MARKED_V1_TOOL], 2, MARKED_SYSTEM), + ({"extra_body": {"tools": [UNMARKED_V1_TOOL]}}, [MARKED_V1_TOOL], 3, "sys"), + ({"extra_body": {"tools": [MARKED_V1_TOOL]}}, [UNMARKED_V1_TOOL], 3, "sys"), + ({"extra_body": {"cache_control": EPHEMERAL}, "cache_control": EPHEMERAL}, None, 2, MARKED_SYSTEM), + ], + ids=["same_marked_tool_both_ways", "extra_body_unmarks", "extra_body_marks", "root_cache_control_both_ways"], + ) + def test_v1_messages_cap_reserves_for_the_larger_of_direct_and_extra_body_marks( + self, kwargs, tools, marked_turns, expected_system + ): + kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), **copy.deepcopy(kwargs)} + _, result_sys = self._inject(self._marked_user_turns(marked_turns), kwargs, tools=tools) + assert result_sys == expected_system + + def test_v1_messages_automatic_defaults_stand_down_for_root_cache_control(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) root_cache_control = {"type": "ephemeral"} kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}} - if configured is not None: - kwargs["cache_control_injection_points"] = copy.deepcopy(configured) result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) @@ -2382,17 +2633,28 @@ def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, confi assert kwargs["cache_control"] is root_cache_control assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + @pytest.mark.parametrize( + "marked_turns,expected_system", + [(2, [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]), (3, "sys")], + ) + def test_v1_messages_configured_points_apply_with_root_cache_control_reserving_a_slot( + self, marked_turns, expected_system + ): + root_cache_control = {"type": "ephemeral"} + kwargs = { + "cache_control": root_cache_control, + "cache_control_injection_points": copy.deepcopy(self.CONFIGURED), + } + _, result_system = self._inject(self._marked_user_turns(marked_turns), kwargs) + assert result_system == expected_system + assert kwargs["cache_control"] is root_cache_control + def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self): - """The advisor interceptor re-enters anthropic_messages() with the outer - request's kwargs and post-injection messages. The first pass applies the - message point and writes back a stamped tool_config remainder; the - re-entry must keep that remainder even though the messages and system - now carry litellm's own marks.""" points = [{"location": "message", "role": "system"}, {"location": "tool_config"}] kwargs = {"cache_control_injection_points": copy.deepcopy(points)} msgs1, sys1 = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) assert sys1[0]["cache_control"] == {"type": "ephemeral"} - expected_remainder = [{"location": "tool_config", "_litellm_judged": True}] + expected_remainder = [{"location": "tool_config"}] assert kwargs["cache_control_injection_points"] == expected_remainder msgs2, sys2 = self._inject(msgs1, kwargs, system=sys1) @@ -2631,22 +2893,26 @@ def test_v1_messages_older_openai_model_keeps_cache_control(self): assert system == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] assert kwargs == {} - def test_v1_messages_client_content_breakpoint_makes_configured_points_stand_down(self): - messages = [{"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}]}] + def test_v1_messages_configured_points_apply_beside_client_content_breakpoint(self): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}]} + ] kwargs = {"cache_control_injection_points": copy.deepcopy(self.SYSTEM_POINT)} result, system = self._inject(messages, "sys", kwargs) assert result == messages - assert system == "sys" - assert kwargs == {} + assert system == [{"type": "text", "text": "sys", "prompt_cache_breakpoint": self.EXPLICIT}] + assert kwargs == {"prompt_cache_options": self.EXPLICIT} - def test_v1_messages_client_system_breakpoint_makes_configured_points_stand_down(self): + def test_v1_messages_tail_point_applies_beside_client_system_breakpoint(self): system = [{"type": "text", "text": "sys", "prompt_cache_breakpoint": self.EXPLICIT}] messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] kwargs = {"cache_control_injection_points": [{"location": "message", "index": -1}]} result, result_system = self._inject(messages, system, kwargs) - assert result == messages + assert result == [ + {"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}]} + ] assert result_system == system - assert kwargs == {} + assert kwargs == {"prompt_cache_options": self.EXPLICIT} def test_chat_system_string_wrapped_with_block_breakpoint(self): params = {"cache_control_injection_points": copy.deepcopy(self.SYSTEM_POINT)} @@ -2710,18 +2976,25 @@ def test_chat_other_targets_keep_message_level_cache_control(self, model): assert processed[0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert params == {} - def test_chat_client_breakpoint_makes_seeded_points_stand_down(self): + def test_chat_seeded_points_apply_beside_client_breakpoint(self): params = {"cache_control_injection_points": copy.deepcopy(self.SYSTEM_POINT)} + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}]}, + ] AnthropicCacheControlHook.maybe_seed_default_injection_points( non_default_params=params, - messages=[ - {"role": "system", "content": "sys"}, - {"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": self.EXPLICIT}]}, - ], + messages=messages, model="openai/gpt-5.6", custom_llm_provider="openai", ) - assert params == {} + assert params["cache_control_injection_points"] == [ + {"location": "message", "role": "system", "_litellm_openai_dialect": True} + ] + _, processed, _ = self._chat(messages, params) + assert processed[0]["content"] == [{"type": "text", "text": "sys", "prompt_cache_breakpoint": self.EXPLICIT}] + assert processed[1] == messages[1] + assert params["prompt_cache_options"] == self.EXPLICIT def test_cap_counts_client_breakpoints_of_both_kinds(self): messages = [ @@ -3315,7 +3588,6 @@ def test_v1_messages_reentry_keeps_the_marker(self, monkeypatch): assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT def test_configured_points_skipping_a_marked_target_record_nothing(self): - """Configured injection stands down on client breakpoints, so no marker lands.""" kwargs: dict = { "litellm_metadata": {}, "cache_control_injection_points": [{"location": "message", "role": "system", "index": None}], diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index ba3a6be609f3..f19a88916094 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1257,6 +1257,25 @@ def test_token_counter_with_thinking_content(): ), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}" + +def test_token_counter_with_redacted_thinking_content(): + """ + A replayed redacted_thinking block (Anthropic redacted reasoning, or the /v1/messages bridge's stand-in + for a reasoning item with no summary) counts zero tokens for its encrypted payload, like a thinking + block with no text. It used to raise, which made is_prompt_caching_valid_prompt return False and the + prompt_caching pre-call check stop pinning the deployment that held the cached prefix. + """ + model = "anthropic/claude-sonnet-4-5-20250929" + reply = {"type": "text", "text": "Draw from the box labeled Mixed, because that label must be wrong."} + redacted_block = {"type": "redacted_thinking", "data": "EqQBCkYIBRgCKkBjZ2xhc3M" * 30} + user_turn = {"role": "user", "content": [{"type": "text", "text": "Which box do you draw from?"}]} + follow_up = {"role": "user", "content": [{"type": "text", "text": "Restate that in one sentence."}]} + + without_block = [user_turn, {"role": "assistant", "content": [reply]}, follow_up] + with_block = [user_turn, {"role": "assistant", "content": [redacted_block, reply]}, follow_up] + + assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block) + def test_token_counter_with_tool_reference_block(): """ Regression test: a message containing an Anthropic tool-search diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index afadd6f3d193..622ab790ffe7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -7725,3 +7725,25 @@ def auth_user(user_id: str): assert result.terminated_sessions == 2 assert {s.user_id for s in result.sessions} == {"bob"} assert "sk-live-bob" not in result.model_dump_json() + + +class TestGetMcpToolsWireShape: + @pytest.mark.asyncio + async def test_get_mcp_tools_returns_each_tool_in_mcp_wire_spelling(self): + from mcp.types import ListToolsResult, Tool + + add_schema = {"type": "object", "properties": {"a": {"type": "integer"}}, "required": ["a"]} + listed = ListToolsResult( + tools=[Tool(name="add", description="Add", inputSchema=add_schema, outputSchema={"type": "integer"})] + ) + with patch( # test-quality-ok: get_mcp_tools has no injection seam for the MCP tool lister on this line + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + AsyncMock(return_value=listed), + ): + result = await mgmt_endpoints.get_mcp_tools(user_api_key_dict=generate_mock_user_api_key_auth()) + + (tool,) = result["tools"] + assert tool["inputSchema"] == add_schema + assert tool["outputSchema"] == {"type": "integer"} + assert "_meta" in tool + assert not {"input_schema", "output_schema", "meta"} & tool.keys() diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 506563a82fbb..983e1cfe80ec 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -10,6 +10,7 @@ import time import uuid from collections.abc import Callable +from typing import Final from unittest.mock import patch import pytest @@ -17,6 +18,7 @@ import litellm from litellm import Router +from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -1676,3 +1678,107 @@ async def test_group_call_429_cools_down_member_across_retries(): ) cooldown_ids = await _call_and_get_cooldowns(router, "quality") assert "deploy-3" in cooldown_ids + + +def _sync_task_count() -> int: + return sum( + 1 + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ) + + +@pytest.mark.asyncio +async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tasks(monkeypatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + assert _sync_task_count() == 1 + selector_before: Final = router.lowesttpm_logger_v2 + + for _ in range(5): + router.update_settings(routing_strategy_args={"ttl": 60}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2 is selector_before, "same routing_strategy_args must not rebuild the selector" + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2.routing_args.ttl == 120 + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + + +class _RecordingRedisCache: + def __init__(self) -> None: + self.increment_lists: list[list[RedisPipelineIncrementOperation]] = [] + + async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: + self.increment_lists.append(list(increment_list)) + return [float(op["increment_value"]) for op in increment_list] + + +@pytest.mark.asyncio +async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue( + monkeypatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + redis_cache: Final = _RecordingRedisCache() + replaced: Final = router.lowesttpm_logger_v2 + replaced.dual_cache.redis_cache = redis_cache + replaced.redis_increment_operation_queue.append( + RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60) + ) + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + await asyncio.gather( + *( + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy._push_in_memory_increments_to_redis" + ) + ) + + assert router.lowesttpm_logger_v2 is not replaced + assert redis_cache.increment_lists == [ + [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)] + ] + assert replaced.redis_increment_operation_queue == [] + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 333e7b2ff318..84cfc708d162 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -197,6 +197,58 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" +@pytest.mark.asyncio +async def test_replayed_redacted_thinking_block_still_records_and_pins(): + """ + A model that returns no reasoning summary (gpt-5.x through the /v1/messages bridge, Anthropic with + redacted reasoning) hands the client a `redacted_thinking` block, and the client replays it on every + later turn. The token count behind `is_prompt_caching_valid_prompt` raised on that block, the helper + swallowed it to False, and the check neither recorded the serving deployment nor pinned it, so the + conversation bounced across the group and paid a cache write on each deployment. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + model = "openai/gpt-5.6-sol" + deployments = _deployments(model, model, model) + messages = cast( + list[AllMessageValues], + [ + *_messages(word_count=3000), + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:" + "Z" * 400}, + {"type": "text", "text": "Draw from the box labeled Mixed."}, + ], + }, + {"role": "user", "content": "Restate that in one sentence."}, + ], + ) + + assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True + + await check.async_log_success_event( + kwargs={ + "standard_logging_object": { + "call_type": "anthropic_messages", + "model": model, + "messages": messages, + "model_id": "dep-2", + } + }, + response_obj=None, + start_time=None, + end_time=None, + ) + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + ) + + assert filtered == [deployments[1]] + + def _auto_caching_messages() -> List[AllMessageValues]: """A prompt over the model minimum that carries no client cache_control.""" return cast( diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 147e863baf5d..269a8271e5f5 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -40,10 +40,6 @@ def test_shipped_decisions( enabled: Final = environment == "1" if environment is not None else process is not False assert catalog.rollout(context) is Rollout.RUST_OPT_OUT assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) - elif route is Route.MESSAGES: - enabled: Final = environment == "1" if environment is not None else process is True - assert catalog.rollout(context) is Rollout.RUST_OPT_IN - assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) elif route is Route.TRANSCRIPTION and provider == "bedrock": assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED