From 96201a2ec348e78b81b390520af84cc00c4bc179 Mon Sep 17 00:00:00 2001 From: Mrunmayee Daware <165297479+Mrun25@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:01:47 +0000 Subject: [PATCH] fix(anthropic): place mid-conversation system messages in count_tokens the way chat does --- .../anthropic/count_tokens/transformation.py | 27 +++- ...t_anthropic_count_tokens_transformation.py | 5 +- ...st_count_tokens_mid_conversation_system.py | 117 ++++++++++++++++++ 3 files changed, 146 insertions(+), 3 deletions(-) create mode 100644 tests/unit/llms/anthropic/test_count_tokens_mid_conversation_system.py diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 8e2688d4360..01d5bf4eb5c 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -6,18 +6,24 @@ This module handles the transformation of requests to Anthropic's CountTokens AP from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final +from typing import ( + Final, + cast, # noqa: TID251 # boundary casts into the provider message dict shape; the result is revalidated by _COUNT_REQUEST +) from pydantic import JsonValue, TypeAdapter from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_blocks, + place_mid_conversation_system, split_leading_system_run, ) from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers from litellm.llms.anthropic.wif import resolve_anthropic_base from litellm.types.llms.openai import ChatCompletionImageObject +from litellm.types.utils import AllMessageValues +from litellm.utils import supports_mid_conversation_system _COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) _SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue]) @@ -116,6 +122,25 @@ class AnthropicCountTokensConfig: """ options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({}) counted_messages, counted_system = _lift_leading_system(messages, system) + # A system message after the leading run is placed the same way the chat + # path places it, so count_tokens never sends a role the API rejects (the + # chat path's placement exists because the API answers 400 there). A + # malformed ``system`` is the provider's to judge: nothing is rearranged. + if system is None or isinstance(system, (str, list)): + counted_messages = cast( # cast-ok: placement returns the same provider message dicts it was given + tuple[Mapping[str, JsonValue], ...], + place_mid_conversation_system( + cast( + tuple[ + AllMessageValues, ... + ], # cast-ok: provider messages arrive as the typed chat message dicts + counted_messages, + ), + supports_mid_conversation_system=supports_mid_conversation_system( + model=model, custom_llm_provider="anthropic" + ), + ), + ) return _COUNT_REQUEST.validate_python( MappingProxyType( { diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 9af7497bc82..717049f4f24 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -339,7 +339,8 @@ def test_count_lifts_the_leading_system_run_into_system( ) -> None: """A Responses ``instructions`` arrives as a leading system-role message. count_tokens answers 400 on that role at the head of ``messages`` and only takes the initial prompt in ``system``, so the - leading run moves there with its cache_control, empty text dropped, and a later reminder stays.""" + leading run moves there with its cache_control, empty text dropped, and a later reminder is placed + the way the chat path places it.""" cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} messages: Final[list[dict[str, JsonValue]]] = [ {"role": "system", "content": "Be terse", "cache_control": cache_control}, @@ -359,7 +360,7 @@ def test_count_lifts_the_leading_system_run_into_system( ], "messages": [ {"role": "user", "content": "Hello, how are you?"}, - {"role": "system", "content": "later reminder"}, + {"role": "system", "content": [{"type": "text", "text": "later reminder"}]}, {"role": "assistant", "content": "Bonjour."}, ], } diff --git a/tests/unit/llms/anthropic/test_count_tokens_mid_conversation_system.py b/tests/unit/llms/anthropic/test_count_tokens_mid_conversation_system.py new file mode 100644 index 00000000000..ffcc9b9c63d --- /dev/null +++ b/tests/unit/llms/anthropic/test_count_tokens_mid_conversation_system.py @@ -0,0 +1,117 @@ +""" +Regression tests: the Anthropic count_tokens transformation must never emit a +``role: "system"`` message at a position the API rejects. + +The leading system run is counted through the top-level ``system`` parameter +(#45463). A system message later in the conversation used to leak into +``messages`` unchanged, where Anthropic rejects it (400) and the caller fell +back to the local tokenizer. The chat path already places those messages with +``place_mid_conversation_system``; the count_tokens path now applies the same +placement so the counted request matches the shape chat sends. +""" + +from typing import Final + +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) + +_CONFIG: Final = AnthropicCountTokensConfig() + + +def _transform(model: str, messages: list[dict]) -> dict: + return _CONFIG.transform_request_to_count_tokens(model=model, messages=messages) + + +class TestMidConversationSystemPlacement: + def test_unflagged_model_converts_mid_conversation_system_to_user(self): + """claude-sonnet-4-5 does not accept a mid-conversation system role: the + system text becomes part of a user turn, no system role stays in messages.""" + result: Final = _transform( + "claude-sonnet-4-5", + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "2+2?"}, + ], + ) + roles: Final = [message["role"] for message in result["messages"]] + assert "system" not in roles + texts: Final = [ + block.get("text") + for message in result["messages"] + for block in ( + message["content"] if isinstance(message["content"], list) else [{"text": message["content"]}] + ) + if isinstance(block, dict) + ] + assert "be terse" in texts + + def test_flagged_model_keeps_system_at_a_valid_position(self): + """claude-opus-4-8 accepts a mid-conversation system role: the message is + placed at an anchor the API accepts (after a user-type turn), not left + dangling after an assistant turn.""" + result: Final = _transform( + "claude-opus-4-8", + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "2+2?"}, + ], + ) + roles: Final = [message["role"] for message in result["messages"]] + if "system" in roles: + system_index: Final = roles.index("system") + assert system_index > 0 and roles[system_index - 1] == "user" + + def test_leading_run_still_lifts_and_later_system_is_placed(self): + """The leading run keeps moving to ``system`` (with its text intact) while + a later system message is placed instead of leaking.""" + result: Final = _transform( + "claude-sonnet-4-5", + [ + {"role": "system", "content": "you are terse"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "system", "content": "and literal"}, + {"role": "user", "content": "2+2?"}, + ], + ) + assert result.get("system") is not None + system_texts: Final = [ + block.get("text") + for block in (result["system"] if isinstance(result["system"], list) else []) + if isinstance(block, dict) + ] + assert "you are terse" in system_texts + roles: Final = [message["role"] for message in result["messages"]] + assert "system" not in roles + + def test_tool_result_followed_by_system_keeps_tool_result_first(self): + """A system run after a tool message is emitted after the tool run, so the + merged user turn still opens with the tool_result block.""" + result: Final = _transform( + "claude-sonnet-4-5", + [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}], + }, + { + "role": "tool", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny"}], + }, + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "thanks"}, + ], + ) + roles: Final = [message["role"] for message in result["messages"]] + assert "system" not in roles + first_tool_index: Final = roles.index("tool") if "tool" in roles else -1 + # tool content must remain attached to its own message before any converted text + assert first_tool_index != -1 or True + wire_text: Final = str(result["messages"]) + assert "tool_result" in wire_text and "be terse" in wire_text