diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e6a68de07e9..5e6df8d9d95 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -764,7 +764,16 @@ def _format_function_definitions(tools): lines.append("namespace functions {") lines.append("") for tool in tools: + if not isinstance(tool, dict): + continue function = tool.get("function") + if function is None: + # Anthropic tool shape → OpenAI function dict for token counting. + function = { + "name": tool.get("name"), + "description": tool.get("description"), + "parameters": tool.get("input_schema", tool.get("parameters", {})), + } if function_description := function.get("description"): lines.append(f"// {function_description}") function_name = function.get("name") diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 8ed6126d2eb..6908921a43f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -18,6 +18,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, ) +from litellm.types.llms.anthropic import AppliedEdit from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -28,14 +29,10 @@ if TYPE_CHECKING: pass -# Anthropic-only fields that the translator above already maps into the -# OpenAI-format completion_kwargs (output_config → reasoning_effort / -# response_format, etc.). They must be filtered out of the raw -# extra_kwargs re-merge below or non-Anthropic backends reject the call -# with 400 "Extra inputs are not permitted". Add new entries here when -# extending AnthropicMessagesRequestOptionalParams with another Anthropic- -# specific key. -ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset({"output_config"}) +# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. +ANTHROPIC_ONLY_REQUEST_KEYS: frozenset[str] = frozenset( + {"output_config", "context_management", "_polyfill_applied_edits"} +) ######################################################## # init adapter @@ -306,6 +303,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: Optional[int] = None, top_p: Optional[float] = None, output_format: Optional[Dict] = None, + _polyfill_applied_edits: Optional[List[AppliedEdit]] = None, **kwargs, ) -> Union[AnthropicMessagesResponse, AsyncIterator]: """Handle non-Anthropic models asynchronously using the adapter""" @@ -338,6 +336,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_response, model=model, tool_name_mapping=tool_name_mapping, + applied_edits=_polyfill_applied_edits, ) ) if transformed_stream is not None: @@ -347,6 +346,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params( cast(ModelResponse, completion_response), tool_name_mapping=tool_name_mapping, + applied_edits=_polyfill_applied_edits, ) if anthropic_response is not None: return anthropic_response @@ -369,6 +369,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_p: Optional[float] = None, output_format: Optional[Dict] = None, _is_async: bool = False, + _polyfill_applied_edits: Optional[List[AppliedEdit]] = None, **kwargs, ) -> Union[ AnthropicMessagesResponse, @@ -392,6 +393,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k=top_k, top_p=top_p, output_format=output_format, + _polyfill_applied_edits=_polyfill_applied_edits, **kwargs, ) @@ -424,6 +426,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_response, model=model, tool_name_mapping=tool_name_mapping, + applied_edits=_polyfill_applied_edits, ) ) if transformed_stream is not None: @@ -433,6 +436,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params( cast(ModelResponse, completion_response), tool_name_mapping=tool_name_mapping, + applied_edits=_polyfill_applied_edits, ) if anthropic_response is not None: return anthropic_response diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index c65dfb22730..468283f804c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -3,11 +3,24 @@ import json import traceback from collections import deque -from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Dict, + Iterator, + List, + Literal, + Optional, +) -from litellm import verbose_logger +from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.types.llms.anthropic import UsageDelta +from litellm.types.llms.anthropic import ( + AppliedEdit, + ContextManagementResponse, + UsageDelta, +) from litellm.types.utils import AdapterCompletionStreamWrapper if TYPE_CHECKING: @@ -48,11 +61,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): completion_stream: Any, model: str, tool_name_mapping: Optional[Dict[str, str]] = None, + applied_edits: Optional[List[AppliedEdit]] = None, ): super().__init__(completion_stream) self.model = model # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) self.tool_name_mapping = tool_name_mapping or {} + # Polyfill applied_edits on final message_delta. + self.applied_edits: List[AppliedEdit] = list(applied_edits or []) def _create_initial_usage_delta(self) -> UsageDelta: """ @@ -125,6 +141,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, + applied_edits=self.applied_edits or None, ) if should_start_new_block and not self.sent_content_block_finish: @@ -266,6 +283,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): processed_chunk = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_response_to_anthropic( response=chunk, current_content_block_index=self.current_content_block_index, + applied_edits=self.applied_edits or None, ) # Check if this is a usage chunk and we have a held stop_reason chunk @@ -312,6 +330,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): chunk.usage._cache_read_input_tokens ) merged_chunk["usage"] = usage_dict + if self.applied_edits and "context_management" not in merged_chunk: + merged_chunk["context_management"] = ContextManagementResponse( + applied_edits=list(self.applied_edits) + ) # Queue the merged chunk and reset self.chunk_queue.append(merged_chunk) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 51a1e739a0f..9fc88874b13 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -87,11 +87,13 @@ from litellm.types.llms.anthropic import ( AnthropicResponseContentBlockText, AnthropicResponseContentBlockThinking, AnthropicResponseContentBlockToolUse, + AppliedEdit, ContentBlockDelta, ContentJsonBlockDelta, ContentTextBlockDelta, ContentThinkingBlockDelta, ContentThinkingSignatureBlockDelta, + ContextManagementResponse, MessageBlockDelta, MessageDelta, UsageDelta, @@ -195,6 +197,7 @@ class AnthropicAdapter: self, response: ModelResponse, tool_name_mapping: Optional[Dict[str, str]] = None, + applied_edits: Optional[List[AppliedEdit]] = None, ) -> Optional[AnthropicMessagesResponse]: """ Translate OpenAI response to Anthropic format. @@ -204,10 +207,12 @@ class AnthropicAdapter: tool_name_mapping: Optional mapping of truncated tool names to original names. Used to restore original names for tools that exceeded OpenAI's 64-char limit. + applied_edits: Polyfill AppliedEdit list for response context_management. """ return LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( response=response, tool_name_mapping=tool_name_mapping, + applied_edits=applied_edits, ) def translate_completion_output_params_streaming( @@ -215,6 +220,7 @@ class AnthropicAdapter: completion_stream: Any, model: str, tool_name_mapping: Optional[Dict[str, str]] = None, + applied_edits: Optional[List[AppliedEdit]] = None, ) -> Union[AsyncIterator[bytes], None]: """ Translate OpenAI streaming response to Anthropic format. @@ -223,11 +229,13 @@ class AnthropicAdapter: completion_stream: The OpenAI streaming response model: The model name tool_name_mapping: Optional mapping of truncated tool names to original names. + applied_edits: Polyfill AppliedEdit list on final message_delta. """ anthropic_wrapper = AnthropicStreamWrapper( completion_stream=completion_stream, model=model, tool_name_mapping=tool_name_mapping, + applied_edits=applied_edits, ) # Return the SSE-wrapped version for proper event formatting return anthropic_wrapper.async_anthropic_sse_wrapper() @@ -1342,6 +1350,7 @@ class LiteLLMAnthropicMessagesAdapter: self, response: ModelResponse, tool_name_mapping: Optional[Dict[str, str]] = None, + applied_edits: Optional[List[AppliedEdit]] = None, ) -> AnthropicMessagesResponse: """ Translate OpenAI response to Anthropic format. @@ -1351,6 +1360,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_name_mapping: Optional mapping of truncated tool names to original names. Used to restore original names for tools that exceeded OpenAI's 64-char limit. + applied_edits: Polyfill AppliedEdit list for response context_management. """ ## translate content block anthropic_content = self._translate_openai_content_to_anthropic( @@ -1396,6 +1406,11 @@ class LiteLLMAnthropicMessagesAdapter: stop_reason=anthropic_finish_reason, ) + if applied_edits: + translated_obj["context_management"] = ContextManagementResponse( + applied_edits=list(applied_edits) + ) + return translated_obj def _translate_streaming_openai_chunk_to_anthropic_content_block( @@ -1528,7 +1543,10 @@ class LiteLLMAnthropicMessagesAdapter: return "text_delta", ContentTextBlockDelta(type="text_delta", text=text) def translate_streaming_openai_response_to_anthropic( - self, response: ModelResponse, current_content_block_index: int + self, + response: ModelResponse, + current_content_block_index: int, + applied_edits: Optional[List[AppliedEdit]] = None, ) -> Union[ContentBlockDelta, MessageBlockDelta]: ## base case - final chunk w/ finish reason if response.choices[0].finish_reason is not None: @@ -1578,9 +1596,14 @@ class LiteLLMAnthropicMessagesAdapter: usage_delta["cache_read_input_tokens"] = cached_tokens else: usage_delta = UsageDelta(input_tokens=0, output_tokens=0) - return MessageBlockDelta( + message_block = MessageBlockDelta( type="message_delta", delta=delta, usage=usage_delta # type: ignore ) + if applied_edits: + message_block["context_management"] = ContextManagementResponse( + applied_edits=list(applied_edits) + ) + return message_block ( type_of_content, content_block_delta, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py new file mode 100644 index 00000000000..f9bd9f4c66d --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py @@ -0,0 +1,4 @@ +from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER +from .dispatcher import apply_context_management + +__all__ = ["apply_context_management", "CLEARED_TOOL_RESULT_PLACEHOLDER"] diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py new file mode 100644 index 00000000000..c5775315856 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/constants.py @@ -0,0 +1,10 @@ +"""Constants for the in-gateway context-management polyfill.""" + +CLEAR_TOOL_USES_EDIT_TYPE = "clear_tool_uses_20250919" +CLEAR_THINKING_EDIT_TYPE = "clear_thinking_20251015" +COMPACT_EDIT_TYPE = "compact_20260112" + +DEFAULT_INPUT_TOKENS_TRIGGER = 100_000 +DEFAULT_KEEP_TOOL_USES = 3 + +CLEARED_TOOL_RESULT_PLACEHOLDER = "[Cleared by context management]" diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py new file mode 100644 index 00000000000..17a510b8aba --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -0,0 +1,56 @@ +"""Dispatch ``context_management`` edits to registered polyfill editors.""" + +from typing import Any, Callable, Dict, List, Optional, Tuple + +from litellm._logging import verbose_logger +from litellm.types.llms.anthropic import AppliedEdit + +from .constants import CLEAR_TOOL_USES_EDIT_TYPE +from .editors import apply_clear_tool_uses_20250919 + +EditorFn = Callable[..., Tuple[List[Dict[str, Any]], Optional[AppliedEdit]]] + +_EDITOR_REGISTRY: Dict[str, EditorFn] = { + CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919, +} + + +def apply_context_management( + *, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + system: Any, + context_management_spec: Dict[str, Any], +) -> Tuple[List[Dict[str, Any]], List[AppliedEdit]]: + """Run edits in order; return (messages, applied_edits that fired).""" + edits = context_management_spec.get("edits") if context_management_spec else None + if not edits or not isinstance(edits, list): + return messages, [] + + applied_edits: List[AppliedEdit] = [] + current_messages = messages + + for edit_spec in edits: + if not isinstance(edit_spec, dict): + continue + edit_type = edit_spec.get("type") + editor = _EDITOR_REGISTRY.get(edit_type) if isinstance(edit_type, str) else None + if editor is None: + verbose_logger.debug( + "context_management polyfill: unknown edit type '%s' — skipping", + edit_type, + ) + continue + + current_messages, applied = editor( + model=model, + messages=current_messages, + tools=tools, + system=system, + edit_spec=edit_spec, + ) + if applied is not None: + applied_edits.append(applied) + + return current_messages, applied_edits diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/__init__.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/__init__.py new file mode 100644 index 00000000000..d402632380d --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/__init__.py @@ -0,0 +1,3 @@ +from .clear_tool_uses import apply_clear_tool_uses_20250919 + +__all__ = ["apply_clear_tool_uses_20250919"] diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py new file mode 100644 index 00000000000..c2193406b32 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py @@ -0,0 +1,198 @@ +"""``clear_tool_uses_20250919`` polyfill (v0: ``trigger`` and ``keep`` only).""" + +from typing import Any, Dict, List, Optional, Tuple, cast + +import litellm +from litellm._logging import verbose_logger +from litellm.types.llms.anthropic import AppliedEdit + +from ..constants import ( + CLEAR_TOOL_USES_EDIT_TYPE, + DEFAULT_INPUT_TOKENS_TRIGGER, + DEFAULT_KEEP_TOOL_USES, +) +from ..placeholders import build_cleared_tool_result_content + + +def _count_tool_uses(messages: List[Dict[str, Any]]) -> int: + """Return the number of tool_use content blocks across all messages.""" + count = 0 + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + count += 1 + return count + + +def _collect_tool_use_ids_in_order(messages: List[Dict[str, Any]]) -> List[str]: + """Return tool_use ids in the chronological order they appear in messages.""" + ids: List[str] = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + block_id = block.get("id") + if isinstance(block_id, str): + ids.append(block_id) + return ids + + +def _trigger_met( + trigger: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], +) -> Tuple[bool, Optional[int]]: + """Return (trigger_met, input_tokens if counted for reuse).""" + trigger_type = trigger.get("type", "input_tokens") + threshold = trigger.get("value") + + if trigger_type == "tool_uses": + if not isinstance(threshold, int): + return False, None + return _count_tool_uses(messages) > threshold, None + + if not isinstance(threshold, int): + threshold = DEFAULT_INPUT_TOKENS_TRIGGER + current_tokens = litellm.token_counter( + model=model, + messages=messages, + tools=cast(Any, tools), + ) + verbose_logger.debug( + f"context_management polyfill: current_tokens: {current_tokens}" + ) + verbose_logger.debug(f"context_management polyfill: threshold: {threshold}") + return current_tokens > threshold, current_tokens + + +def _resolve_keep_count(keep: Dict[str, Any]) -> int: + keep_type = keep.get("type", "tool_uses") + if keep_type != "tool_uses": + return DEFAULT_KEEP_TOOL_USES + value = keep.get("value") + if not isinstance(value, int) or value < 0: + return DEFAULT_KEEP_TOOL_USES + return value + + +def _last_completed_tool_use_id( + messages: List[Dict[str, Any]], +) -> Optional[str]: + """Latest completed tool_result id; never cleared.""" + last_id: Optional[str] = None + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + block_id = block.get("tool_use_id") + if isinstance(block_id, str): + last_id = block_id + return last_id + + +def _clear_tool_results( + messages: List[Dict[str, Any]], ids_to_clear: set +) -> Tuple[List[Dict[str, Any]], int]: + """Clear matching tool_result content; return (messages, cleared_count).""" + cleared = 0 + new_messages: List[Dict[str, Any]] = [] + for msg in messages: + content = msg.get("content") + if not isinstance(content, list): + new_messages.append(msg) + continue + + new_blocks: List[Any] = [] + mutated = False + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "tool_result" + and block.get("tool_use_id") in ids_to_clear + ): + new_block = { + **block, + "content": build_cleared_tool_result_content(block.get("content")), + } + new_blocks.append(new_block) + mutated = True + cleared += 1 + else: + new_blocks.append(block) + + if mutated: + new_messages.append({**msg, "content": new_blocks}) + else: + new_messages.append(msg) + + return new_messages, cleared + + +def apply_clear_tool_uses_20250919( + *, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + system: Any, + edit_spec: Dict[str, Any], +) -> Tuple[List[Dict[str, Any]], Optional[AppliedEdit]]: + """Apply clear_tool_uses; return (messages, AppliedEdit or None).""" + for ignored_knob in ("clear_at_least", "exclude_tools", "clear_tool_inputs"): + if ignored_knob in edit_spec: + verbose_logger.debug( + "context_management polyfill: ignoring '%s' on %s " + "(supported only on Anthropic-family forwarding path in v0)", + ignored_knob, + CLEAR_TOOL_USES_EDIT_TYPE, + ) + + trigger = edit_spec.get("trigger") or { + "type": "input_tokens", + "value": DEFAULT_INPUT_TOKENS_TRIGGER, + } + keep = edit_spec.get("keep") or { + "type": "tool_uses", + "value": DEFAULT_KEEP_TOOL_USES, + } + + met, tokens_before = _trigger_met(trigger, model, messages, tools) + if not met: + return messages, None + + keep_count = _resolve_keep_count(keep) + tool_use_ids = _collect_tool_use_ids_in_order(messages) + if len(tool_use_ids) <= keep_count: + return messages, None + + ids_to_clear = set(tool_use_ids[: len(tool_use_ids) - keep_count]) + + # Never clear the latest completed tool_result (reply context). + last_completed_id = _last_completed_tool_use_id(messages) + if last_completed_id is not None: + ids_to_clear.discard(last_completed_id) + + edited, cleared_count = _clear_tool_results(messages, ids_to_clear) + verbose_logger.debug(f"context_management polyfill: edited: {edited}") + if cleared_count == 0: + return messages, None + + if tokens_before is None: + tokens_before = litellm.token_counter( + model=model, messages=messages, tools=cast(Any, tools) + ) + tokens_after = litellm.token_counter( + model=model, messages=edited, tools=cast(Any, tools) + ) + cleared_input_tokens = max(tokens_before - tokens_after, 0) + + applied: AppliedEdit = { + "type": CLEAR_TOOL_USES_EDIT_TYPE, + "cleared_tool_uses": cleared_count, + "cleared_input_tokens": cleared_input_tokens, + } + return edited, applied diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py b/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py new file mode 100644 index 00000000000..f684d970df4 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/placeholders.py @@ -0,0 +1,14 @@ +"""Placeholder content for cleared ``tool_result`` blocks (string or block list).""" + +from typing import Any, List, Union + +from .constants import CLEARED_TOOL_RESULT_PLACEHOLDER + + +def build_cleared_tool_result_content( + original_content: Any, +) -> Union[str, List[dict]]: + """Return a string or single text block list, matching ``original_content`` shape.""" + if isinstance(original_content, list): + return [{"type": "text", "text": CLEARED_TOOL_RESULT_PLACEHOLDER}] + return CLEARED_TOOL_RESULT_PLACEHOLDER diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 14e06e047ea..c767cc6bb45 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -20,6 +20,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.llms.anthropic import AppliedEdit from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -456,9 +457,29 @@ def anthropic_messages_handler( return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( **_shared_kwargs ) + + # In-gateway context_management polyfill on the chat-completions adapter + # (native on Anthropic/Responses paths). Skipped when drop_params is on. + context_management_spec = _shared_kwargs.pop("context_management", None) + polyfill_applied_edits: List[AppliedEdit] = [] + if context_management_spec and not litellm.drop_params: + from litellm.llms.anthropic.experimental_pass_through.context_management import ( + apply_context_management, + ) + + edited_messages, polyfill_applied_edits = apply_context_management( + model=model, + messages=_shared_kwargs["messages"], + tools=_shared_kwargs.get("tools"), + system=_shared_kwargs.get("system"), + context_management_spec=context_management_spec, + ) + _shared_kwargs["messages"] = edited_messages + return ( LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - **_shared_kwargs + _polyfill_applied_edits=polyfill_applied_edits, + **_shared_kwargs, ) ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index efc890d9ee2..7d43e13ff0b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -573,6 +573,9 @@ class AmazonConverseConfig(BaseConfig): ): supported_params.append("thinking") supported_params.append("reasoning_effort") + + if base_model.startswith("anthropic"): + supported_params.append("context_management") return supported_params def map_tool_choice_values( @@ -932,6 +935,13 @@ class AmazonConverseConfig(BaseConfig): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params ) + if param == "context_management" and isinstance(value, (dict, list)): + # Forward context_management; later filter keeps compact_20260112 only. + optional_params["context_management"] = ( + AnthropicConfig.map_openai_context_management_to_anthropic( + cast(Union[dict, list], value) + ) + ) if param == "requestMetadata": if value is not None and isinstance(value, dict): self._validate_request_metadata(value) # type: ignore @@ -1430,6 +1440,11 @@ class AmazonConverseConfig(BaseConfig): if ANTHROPIC_EFFORT_BETA_HEADER not in anthropic_beta_list: anthropic_beta_list.append(ANTHROPIC_EFFORT_BETA_HEADER) + # Bedrock Converse: compact_20260112 edits only (+ beta header). + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional_request_params, anthropic_beta_list + ) + # Set anthropic_beta in additional_request_params if we have any beta features # ONLY apply to Anthropic/Claude models - other models (e.g., Qwen, Llama) don't support this field if anthropic_beta_list and base_model.startswith("anthropic"): @@ -1437,6 +1452,35 @@ class AmazonConverseConfig(BaseConfig): return bedrock_tools, anthropic_beta_list + @staticmethod + def _filter_context_management_for_bedrock_converse( + additional_request_params: dict, + anthropic_beta_list: list, + ) -> None: + """Keep only compact_20260112 edits for Bedrock; add beta header or drop field.""" + cm = additional_request_params.get("context_management") + if not isinstance(cm, dict): + return + edits = cm.get("edits") + if not isinstance(edits, list): + additional_request_params.pop("context_management", None) + return + + compact_edits = [ + e + for e in edits + if isinstance(e, dict) and e.get("type") == "compact_20260112" + ] + if compact_edits: + if "compact-2026-01-12" not in anthropic_beta_list: + anthropic_beta_list.append("compact-2026-01-12") + additional_request_params["context_management"] = { + **cm, + "edits": compact_edits, + } + else: + additional_request_params.pop("context_management", None) + def _transform_request_helper( self, model: str, diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 1c4d31d21ad..4655951ae24 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -2,7 +2,7 @@ from enum import Enum from typing import Any, Dict, Iterable, List, Optional, Union from pydantic import BaseModel, ConfigDict -from typing_extensions import Literal, Required, TypedDict +from typing_extensions import Literal, NotRequired, Required, TypedDict from .openai import ( ChatCompletionCachedContent, @@ -514,6 +514,21 @@ class UsageDelta(TypedDict, total=False): cache_read_input_tokens: int +class AppliedEdit(TypedDict, total=False): + """One applied context_management edit (Anthropic response shape).""" + + type: str + cleared_input_tokens: int + cleared_tool_uses: int + cleared_thinking_turns: int + + +class ContextManagementResponse(TypedDict, total=False): + """Response ``context_management`` with ``applied_edits``.""" + + applied_edits: List[AppliedEdit] + + class MessageBlockDelta(TypedDict): """ Anthropic @@ -523,6 +538,7 @@ class MessageBlockDelta(TypedDict): type: Literal["message_delta"] delta: MessageDelta usage: UsageDelta + context_management: NotRequired[ContextManagementResponse] class MessageChunk(TypedDict, total=False): diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 1eab1b37e06..85a2b3fee7c 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -1,10 +1,11 @@ from typing import Any, Dict, List, Literal, Optional, Union -from typing_extensions import TypeAlias, TypedDict +from typing_extensions import NotRequired, TypeAlias, TypedDict from litellm.types.llms.anthropic import ( AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, + ContextManagementResponse, ) @@ -94,3 +95,4 @@ class AnthropicMessagesResponse(TypedDict, total=False): stop_sequence: Optional[str] type: Optional[Literal["message"]] usage: Optional[AnthropicUsage] + context_management: NotRequired[ContextManagementResponse] diff --git a/tests/pass_through_unit_tests/test_context_management_polyfill.py b/tests/pass_through_unit_tests/test_context_management_polyfill.py new file mode 100644 index 00000000000..564dbe36f66 --- /dev/null +++ b/tests/pass_through_unit_tests/test_context_management_polyfill.py @@ -0,0 +1,272 @@ +"""Integration tests for context_management polyfill on /v1/messages adapter path.""" + +import json +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.anthropic.experimental_pass_through.context_management.constants import ( + CLEARED_TOOL_RESULT_PLACEHOLDER, +) +from litellm.types.utils import ( + Choices, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + Delta, + Usage, +) + +MODEL = "xai/grok-4" + + +def _make_history(n_pairs: int, result_filler: str = "x" * 50): + messages = [{"role": "user", "content": "Compare weather across cities."}] + for i in range(n_pairs): + messages.append( + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": f"toolu_{i:02d}", + "name": "get_weather", + "input": {"location": f"City{i}"}, + } + ], + } + ) + messages.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{i:02d}", + "content": f"Result {i}: {result_filler}", + } + ], + } + ) + return messages + + +def _mock_completion_response() -> ModelResponse: + return ModelResponse( + id="chatcmpl-test", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(role="assistant", content="ok"), + ) + ], + created=0, + model="grok-4", + object="chat.completion", + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ) + + +async def _mock_streaming_chunks(): + yield ModelResponseStream( + id="chatcmpl-test", + created=0, + model="grok-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content="ok"), + ) + ], + ) + yield ModelResponseStream( + id="chatcmpl-test", + created=0, + model="grok-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage(prompt_tokens=10, completion_tokens=2, total_tokens=12), + ) + + +@pytest.mark.asyncio +async def test_polyfill_round_trip_non_streaming(): + captured = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + return _mock_completion_response() + + with patch("litellm.acompletion", side_effect=fake_acompletion): + response = await litellm.anthropic.messages.acreate( + model=MODEL, + messages=_make_history(n_pairs=5), + max_tokens=128, + api_key="sk-test", + context_management={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 2}, + } + ] + }, + ) + + # 1. Downstream got the edited messages — older tool_result.content cleared. + downstream_messages = captured.get("messages") + assert downstream_messages is not None + cleared_ids = {"toolu_00", "toolu_01", "toolu_02"} + kept_ids = {"toolu_03", "toolu_04"} + found_cleared = 0 + for msg in downstream_messages: + # The adapter may have translated the messages out of Anthropic shape; + # we accept either Anthropic-shape (tool_result block) or OpenAI-shape + # (tool-role message whose content is the placeholder). + if isinstance(msg, dict) and msg.get("role") == "tool": + if msg.get("tool_call_id") in cleared_ids: + content = msg.get("content") + if isinstance(content, str): + if CLEARED_TOOL_RESULT_PLACEHOLDER in content: + found_cleared += 1 + elif isinstance(content, list): + text = "".join( + b.get("text", "") for b in content if isinstance(b, dict) + ) + if CLEARED_TOOL_RESULT_PLACEHOLDER in text: + found_cleared += 1 + elif msg.get("tool_call_id") in kept_ids: + content = msg.get("content") + if isinstance(content, str): + assert CLEARED_TOOL_RESULT_PLACEHOLDER not in content + assert found_cleared == 3 + + # 2. context_management must not leak into downstream kwargs. + assert "context_management" not in captured + + # 3. Response carries the applied_edits in Anthropic's documented shape. + assert isinstance(response, dict) + cm = response.get("context_management") + assert cm is not None, f"context_management missing from response: {response}" + edits = cm.get("applied_edits") + assert isinstance(edits, list) and len(edits) == 1 + edit = edits[0] + assert edit["type"] == "clear_tool_uses_20250919" + assert edit["cleared_tool_uses"] == 3 + assert "cleared_input_tokens" in edit + + +@pytest.mark.asyncio +async def test_polyfill_trigger_not_met_passes_through_unchanged(): + captured = {} + + async def fake_acompletion(**kwargs): + captured.update(kwargs) + return _mock_completion_response() + + with patch("litellm.acompletion", side_effect=fake_acompletion): + response = await litellm.anthropic.messages.acreate( + model=MODEL, + messages=_make_history(n_pairs=2), + max_tokens=128, + api_key="sk-test", + context_management={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 10_000_000}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + ) + + # Downstream still got the request, but no edits applied. + assert captured.get("messages") is not None + assert "context_management" not in captured + + # Response shouldn't carry context_management when nothing fired. + assert isinstance(response, dict) + assert ( + response.get("context_management") is None + or response.get("context_management") == {"applied_edits": []} + or "context_management" not in response + ) + + +@pytest.mark.asyncio +async def test_polyfill_streaming_attaches_to_message_delta(): + async def fake_acompletion(**kwargs): + return _mock_streaming_chunks() + + with patch("litellm.acompletion", side_effect=fake_acompletion): + response = await litellm.anthropic.messages.acreate( + model=MODEL, + messages=_make_history(n_pairs=5), + max_tokens=128, + api_key="sk-test", + stream=True, + context_management={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 2}, + } + ] + }, + ) + + # Collect all SSE bytes. + collected = [] + async for chunk in response: # type: ignore[union-attr] + if isinstance(chunk, (bytes, bytearray)): + collected.append(chunk.decode("utf-8")) + else: + collected.append(str(chunk)) + sse_text = "".join(collected) + + # Find the message_delta event payload and check it carries context_management + # as a sibling of `usage` per Anthropic's spec. + found_delta_with_cm = False + for block in sse_text.split("\n\n"): + if "message_delta" not in block: + continue + data_line = next( + ( + line[len("data:") :].strip() + for line in block.splitlines() + if line.startswith("data:") + ), + None, + ) + if data_line is None: + continue + payload = json.loads(data_line) + if payload.get("type") != "message_delta": + continue + cm = payload.get("context_management") + if cm is None: + continue + assert "applied_edits" in cm + assert len(cm["applied_edits"]) == 1 + assert cm["applied_edits"][0]["type"] == "clear_tool_uses_20250919" + assert cm["applied_edits"][0]["cleared_tool_uses"] == 3 + found_delta_with_cm = True + break + assert found_delta_with_cm, ( + "Expected `context_management` on the message_delta SSE event. " + f"SSE text was: {sse_text!r}" + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py new file mode 100644 index 00000000000..9e450eb9082 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_clear_tool_uses.py @@ -0,0 +1,281 @@ +""" +Unit tests for the in-gateway `clear_tool_uses_20250919` polyfill editor. +""" + +from copy import deepcopy + +from litellm.llms.anthropic.experimental_pass_through.context_management.constants import ( + CLEARED_TOOL_RESULT_PLACEHOLDER, +) +from litellm.llms.anthropic.experimental_pass_through.context_management.editors.clear_tool_uses import ( + apply_clear_tool_uses_20250919, +) + +MODEL = "xai/grok-4" + + +def _make_pair(tool_use_id: str, result_text: str, location: str = "Mumbai"): + """Return an (assistant, user) message pair with one tool_use + tool_result.""" + assistant_msg = { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": tool_use_id, + "name": "get_weather", + "input": {"location": location}, + } + ], + } + user_msg = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_use_id, + "content": result_text, + } + ], + } + return assistant_msg, user_msg + + +def _make_history(n_pairs: int, result_filler: str = "x" * 200): + messages = [{"role": "user", "content": "Compare weather across cities."}] + for i in range(n_pairs): + assistant_msg, user_msg = _make_pair( + tool_use_id=f"toolu_{i:02d}", + result_text=f"Result {i}: {result_filler}", + location=f"City{i}", + ) + messages.append(assistant_msg) + messages.append(user_msg) + return messages + + +def test_below_trigger_returns_unchanged(): + """If trigger threshold isn't exceeded, editor is a no-op.""" + messages = _make_history(n_pairs=2) + original = deepcopy(messages) + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 10_000_000}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is None + assert new_messages == original + + +def test_keep_preserves_most_recent_pairs(): + """With keep=2 and 5 pairs, the 3 oldest pairs are cleared.""" + messages = _make_history(n_pairs=5) + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 2}, + }, + ) + assert applied is not None + assert applied["type"] == "clear_tool_uses_20250919" + assert applied["cleared_tool_uses"] == 3 + + # Tool results for the first 3 pairs should be the placeholder, last 2 untouched. + cleared_ids = {"toolu_00", "toolu_01", "toolu_02"} + kept_ids = {"toolu_03", "toolu_04"} + for msg in new_messages: + if msg.get("role") != "user": + continue + content = msg.get("content") + if not isinstance(content, list): + continue + for block in content: + if block.get("type") != "tool_result": + continue + if block["tool_use_id"] in cleared_ids: + assert block["content"] == CLEARED_TOOL_RESULT_PLACEHOLDER + elif block["tool_use_id"] in kept_ids: + assert "Result" in block["content"] + + +def test_tool_use_input_is_not_cleared(): + """clear_tool_inputs defaults to false — tool_use.input must remain intact.""" + messages = _make_history(n_pairs=3) + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 0}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is not None + # Every tool_use block still has its original `input`. + for msg in new_messages: + if msg.get("role") != "assistant": + continue + for block in msg.get("content", []): + if block.get("type") == "tool_use": + assert block["input"] == {"location": block["input"]["location"]} + assert block["input"]["location"].startswith("City") + + +def test_message_array_length_and_roles_preserved(): + messages = _make_history(n_pairs=4) + original_roles = [m["role"] for m in messages] + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 0}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is not None + assert len(new_messages) == len(messages) + assert [m["role"] for m in new_messages] == original_roles + + +def test_defaults_applied_when_knobs_omitted(): + """No trigger/keep specified — defaults are 100k input_tokens / 3 tool_uses.""" + messages = _make_history(n_pairs=2) + # Below 100k tokens; should not fire. + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={"type": "clear_tool_uses_20250919"}, + ) + assert applied is None + assert new_messages == messages + + +def test_tool_uses_trigger_variant(): + """Trigger by raw count of tool_use blocks, not tokens.""" + messages = _make_history(n_pairs=4) + _, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 2}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is not None + # 4 total - 1 kept = 3 cleared + assert applied["cleared_tool_uses"] == 3 + + +def test_cleared_input_tokens_is_nonnegative(): + messages = _make_history(n_pairs=4) + _, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is not None + assert applied["cleared_input_tokens"] >= 0 + + +def test_ignored_knobs_do_not_alter_behavior(): + """clear_at_least / exclude_tools / clear_tool_inputs are accepted but ignored in v0.""" + messages = _make_history(n_pairs=3) + _, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 0}, + "keep": {"type": "tool_uses", "value": 1}, + "clear_at_least": {"type": "input_tokens", "value": 999_999_999}, + "exclude_tools": ["get_weather"], + "clear_tool_inputs": True, + }, + ) + # Despite clear_at_least being huge, polyfill still applies (knob ignored). + # Despite clear_tool_inputs=True, inputs are NOT cleared (knob ignored). + assert applied is not None + assert applied["cleared_tool_uses"] == 2 + + +def test_tool_result_list_content_shape_preserved(): + """When tool_result.content is a list of blocks, replacement returns a list shape.""" + messages = [ + {"role": "user", "content": "Hi"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_a", "name": "f", "input": {}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_a", + "content": [{"type": "text", "text": "huge result"}], + } + ], + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_b", "name": "f", "input": {}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_b", + "content": [{"type": "text", "text": "keep me"}], + } + ], + }, + ] + new_messages, applied = apply_clear_tool_uses_20250919( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 0}, + "keep": {"type": "tool_uses", "value": 1}, + }, + ) + assert applied is not None + cleared_block = new_messages[2]["content"][0] + assert isinstance(cleared_block["content"], list) + assert cleared_block["content"][0]["type"] == "text" + assert cleared_block["content"][0]["text"] == CLEARED_TOOL_RESULT_PLACEHOLDER diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py new file mode 100644 index 00000000000..924f79d0902 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py @@ -0,0 +1,131 @@ +""" +Unit tests for the context_management polyfill dispatcher. +""" + +from litellm.llms.anthropic.experimental_pass_through.context_management import ( + apply_context_management, +) + +MODEL = "xai/grok-4" + + +def _history_with_two_tool_pairs(): + return [ + {"role": "user", "content": "Hi"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": "first result", + } + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t2", "name": "f", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t2", + "content": "second result", + } + ], + }, + ] + + +def test_unknown_edit_type_is_noop(): + messages = _history_with_two_tool_pairs() + new_messages, applied = apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={ + "edits": [{"type": "totally_not_a_real_edit_20999999"}] + }, + ) + assert applied == [] + assert new_messages == messages + + +def test_known_edit_is_applied(): + messages = _history_with_two_tool_pairs() + _, applied = apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 1}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + ) + assert len(applied) == 1 + assert applied[0]["type"] == "clear_tool_uses_20250919" + assert applied[0]["cleared_tool_uses"] == 1 + + +def test_mixed_known_unknown_only_known_applied(): + messages = _history_with_two_tool_pairs() + _, applied = apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={ + "edits": [ + {"type": "unknown_foo"}, + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 0}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "another_unknown"}, + ] + }, + ) + assert len(applied) == 1 + assert applied[0]["type"] == "clear_tool_uses_20250919" + + +def test_empty_or_missing_edits_list(): + messages = _history_with_two_tool_pairs() + for spec in [{}, {"edits": None}, {"edits": []}, None]: + new_messages, applied = apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec=spec, # type: ignore[arg-type] + ) + assert applied == [] + assert new_messages == messages + + +def test_malformed_edit_entries_are_skipped(): + """Non-dict entries in `edits` list should be silently skipped.""" + messages = _history_with_two_tool_pairs() + new_messages, applied = apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={"edits": ["not a dict", 42, None, {"type": None}]}, + ) + assert applied == [] + assert new_messages == messages diff --git a/tests/test_litellm/llms/bedrock/test_converse_context_management.py b/tests/test_litellm/llms/bedrock/test_converse_context_management.py new file mode 100644 index 00000000000..709fc4e8b39 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/test_converse_context_management.py @@ -0,0 +1,114 @@ +"""Bedrock Converse context_management forwarding (compact_20260112 only).""" + +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + +CLAUDE_MODEL = "anthropic.claude-opus-4-7-20250115-v1:0" + + +def test_supported_params_include_context_management_for_anthropic(): + cfg = AmazonConverseConfig() + params = cfg.get_supported_openai_params(CLAUDE_MODEL) + assert "context_management" in params + + +def test_supported_params_exclude_context_management_for_non_anthropic(): + cfg = AmazonConverseConfig() + params = cfg.get_supported_openai_params("meta.llama3-70b-instruct-v1:0") + assert "context_management" not in params + + +def test_map_openai_params_forwards_anthropic_shape(): + cfg = AmazonConverseConfig() + optional_params: dict = {} + cfg.map_openai_params( + non_default_params={ + "context_management": {"edits": [{"type": "compact_20260112"}]} + }, + optional_params=optional_params, + model=CLAUDE_MODEL, + drop_params=False, + ) + assert optional_params.get("context_management") == { + "edits": [{"type": "compact_20260112"}] + } + + +def test_map_openai_params_normalizes_openai_list_shape(): + """OpenAI Responses-API style list of {type: "compaction"} normalizes to Anthropic dict.""" + cfg = AmazonConverseConfig() + optional_params: dict = {} + cfg.map_openai_params( + non_default_params={"context_management": [{"type": "compaction"}]}, + optional_params=optional_params, + model=CLAUDE_MODEL, + drop_params=False, + ) + forwarded = optional_params.get("context_management") + assert isinstance(forwarded, dict) + edits = forwarded.get("edits") + assert isinstance(edits, list) and len(edits) == 1 + assert edits[0].get("type") == "compact_20260112" + + +def test_filter_keeps_only_compact_edits_and_adds_beta_header(): + additional = { + "context_management": { + "edits": [ + {"type": "clear_tool_uses_20250919"}, + {"type": "compact_20260112"}, + {"type": "clear_thinking_20251015"}, + ] + } + } + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert additional["context_management"]["edits"] == [{"type": "compact_20260112"}] + assert "compact-2026-01-12" in betas + + +def test_filter_drops_field_when_no_compact_edit_remains(): + additional = { + "context_management": { + "edits": [ + {"type": "clear_tool_uses_20250919"}, + {"type": "clear_thinking_20251015"}, + ] + } + } + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert "context_management" not in additional + assert betas == [] + + +def test_filter_is_noop_when_field_absent(): + additional: dict = {} + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert additional == {} + assert betas == [] + + +def test_filter_drops_malformed_edits_list(): + additional = {"context_management": {"edits": "not a list"}} + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert "context_management" not in additional + assert betas == [] + + +def test_filter_does_not_duplicate_beta_header(): + additional = {"context_management": {"edits": [{"type": "compact_20260112"}]}} + betas: list = ["compact-2026-01-12"] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert betas.count("compact-2026-01-12") == 1