diff --git a/litellm/llms/anthropic/messages/context_management/constants.py b/litellm/llms/anthropic/messages/context_management/constants.py index 5d69f070ce8..0592ee61a95 100644 --- a/litellm/llms/anthropic/messages/context_management/constants.py +++ b/litellm/llms/anthropic/messages/context_management/constants.py @@ -3,6 +3,7 @@ from typing import Final CLEAR_TOOL_USES_EDIT_TYPE: Final = "clear_tool_uses_20250919" +CLEAR_THINKING_EDIT_TYPE: Final = "clear_thinking_20251015" DEFAULT_INPUT_TOKENS_TRIGGER: Final = 100_000 DEFAULT_KEEP_TOOL_USES: Final = 3 diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 21af97e523e..a886b948ac0 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1977,8 +1977,9 @@ class AmazonConverseConfig(BaseConfig): additional_request_params: dict, anthropic_beta_list: list, ) -> None: - """Keep only compact_20260112 edits for Bedrock; add beta header or drop field.""" + """Keep compact edits, and clear_thinking edits while thinking is on. Add their betas or drop the field.""" from litellm.llms.anthropic.messages.context_management.constants import ( + CLEAR_THINKING_EDIT_TYPE, COMPACT_EDIT_TYPE, ) from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES @@ -1992,17 +1993,25 @@ class AmazonConverseConfig(BaseConfig): additional_request_params.pop("context_management", None) return - compact_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") == COMPACT_EDIT_TYPE] - if compact_edits: - compact_beta: Final = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value - if compact_beta not in anthropic_beta_list: - anthropic_beta_list.append(compact_beta) - additional_request_params["context_management"] = { - **cm, - "edits": compact_edits, - } - else: + thinking: Final = additional_request_params.get("thinking") + thinking_is_on: Final = isinstance(thinking, dict) and thinking.get("type") in ("enabled", "adaptive") + beta_by_edit_type: Final = { + COMPACT_EDIT_TYPE: ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + **( + {CLEAR_THINKING_EDIT_TYPE: ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value} + if thinking_is_on + else {} + ), + } + kept_edits: Final = [e for e in edits if isinstance(e, dict) and e.get("type") in beta_by_edit_type] + if not kept_edits: additional_request_params.pop("context_management", None) + return + + for beta in dict.fromkeys(beta_by_edit_type[e["type"]] for e in kept_edits): + if beta not in anthropic_beta_list: + anthropic_beta_list.append(beta) + additional_request_params["context_management"] = {**cm, "edits": kept_edits} def _transform_request_helper( self, diff --git a/tests/unit/llms/bedrock/test_converse_context_management.py b/tests/unit/llms/bedrock/test_converse_context_management.py index 709fc4e8b39..71d8420f4ae 100644 --- a/tests/unit/llms/bedrock/test_converse_context_management.py +++ b/tests/unit/llms/bedrock/test_converse_context_management.py @@ -1,4 +1,6 @@ -"""Bedrock Converse context_management forwarding (compact_20260112 only).""" +"""Bedrock Converse context_management forwarding (compact_20260112, and clear_thinking_20251015 while thinking is on).""" + +import pytest from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig @@ -112,3 +114,72 @@ def test_filter_does_not_duplicate_beta_header(): additional, betas ) assert betas.count("compact-2026-01-12") == 1 + + +CLEAR_THINKING_EDIT = { + "type": "clear_thinking_20251015", + "keep": {"type": "thinking_turns", "value": 1}, +} + + +@pytest.mark.parametrize( + "thinking", [{"type": "enabled", "budget_tokens": 2048}, {"type": "adaptive"}] +) +def test_filter_keeps_clear_thinking_edit_and_adds_beta_when_thinking_is_on(thinking): + additional = { + "thinking": thinking, + "context_management": {"edits": [CLEAR_THINKING_EDIT]}, + } + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert additional["context_management"] == {"edits": [CLEAR_THINKING_EDIT]} + assert betas == ["context-management-2025-06-27"] + + +@pytest.mark.parametrize("thinking", [None, {"type": "disabled"}]) +def test_filter_drops_clear_thinking_edit_when_thinking_is_off(thinking): + additional = {"context_management": {"edits": [CLEAR_THINKING_EDIT]}} + if thinking is not None: + additional["thinking"] = thinking + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert "context_management" not in additional + assert betas == [] + + +def test_filter_keeps_compact_and_clear_thinking_together_with_both_betas(): + additional = { + "thinking": {"type": "adaptive"}, + "context_management": { + "edits": [ + CLEAR_THINKING_EDIT, + {"type": "compact_20260112"}, + {"type": "clear_tool_uses_20250919"}, + ] + }, + } + betas: list = [] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert additional["context_management"]["edits"] == [ + CLEAR_THINKING_EDIT, + {"type": "compact_20260112"}, + ] + assert sorted(betas) == ["compact-2026-01-12", "context-management-2025-06-27"] + + +def test_filter_does_not_duplicate_context_management_beta(): + additional = { + "thinking": {"type": "adaptive"}, + "context_management": {"edits": [CLEAR_THINKING_EDIT]}, + } + betas: list = ["context-management-2025-06-27"] + AmazonConverseConfig._filter_context_management_for_bedrock_converse( + additional, betas + ) + assert betas == ["context-management-2025-06-27"]