From 8b0a21b05e6d584544255237238b2f0a841d40d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:52:57 -0800 Subject: [PATCH] fix(count_tokens): keep the caller betas a Converse Bedrock deployment's count body needs Filtering a Converse deployment's count betas through the bedrock_converse row dropped betas the count body still needs, because the count keeps fields like tools[].input_examples and a system message's output_config that Converse inference strips. Mantle counts carrying those fields went from 200 to 400, so the count goes back to filtering through the count API's own row --- .../count_tokens/bedrock_token_counter.py | 17 +-- .../test_bedrock_token_counter.py | 141 +++++++++++++----- 2 files changed, 104 insertions(+), 54 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 690e13e9d21..437d6705eca 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -12,7 +12,6 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger -from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.count_tokens.transformation import ( NO_CLAUDE_COUNT_OPTIONS, @@ -138,12 +137,6 @@ def _options_converse_sends( } -def _betas_converse_sends(anthropic_beta: str | None, inference_model: str | None) -> str | None: - if anthropic_beta is None or inference_model is None or not _runs_on_converse(inference_model): - return anthropic_beta - return ",".join(filter_and_transform_beta_headers(anthropic_beta.split(","), "bedrock_converse")) or None - - class BedrockTokenCounter(BaseTokenCounter): """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" @@ -210,20 +203,18 @@ class BedrockTokenCounter(BaseTokenCounter): resolved_model: Final = get_bedrock_base_model(counted_model) primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE litellm_params: Final = _deployment_litellm_params(deployment) - claude_model: Final = is_mantle_claude_model(resolved_model) - route: Final = _deployment_route(litellm_params) if claude_model else None - inference_model: Final = deployment_inference_model(deployment) if claude_model else None + route: Final = _deployment_route(litellm_params) if is_mantle_claude_model(resolved_model) else None claude_options: Final = ( claude_count_options( counted_model, _claude_options_provider(route, use_mantle), messages, tools, - _options_converse_sends(request_options, inference_model), - _betas_converse_sends(anthropic_beta, inference_model), + _options_converse_sends(request_options, deployment_inference_model(deployment)), + anthropic_beta, deployment, ) - if claude_model + if is_mantle_claude_model(resolved_model) else NO_CLAUDE_COUNT_OPTIONS ) if isinstance(claude_options, ChatCountTranslationError): diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 1e552ba07ea..3b97994456b 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -903,7 +903,7 @@ def _message_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: @pytest.mark.asyncio @pytest.mark.usefixtures("local_beta_headers_config") -async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas_converse_inference_sends() -> None: +async def test_mantle_count_for_a_converse_deployment_sends_converse_format_and_messages_with_mantle_betas() -> None: caller_beta: Final = f"{ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER},context-1m-2025-08-07" request_options: Final[Mapping[str, JsonValue]] = {**_CITY_FORMAT, **_THINKING} inference: Final = _Bedrock(runtime=(200, _CONVERSE_REPLY), mantle=(200, _CONVERSE_REPLY)) @@ -942,58 +942,117 @@ async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas "schema": json.loads(sent_schema), } assert _message_texts(counted) == _message_texts(sent) - assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != () + assert _sent_betas(bedrock.requests[0]) == _row_sends(caller_beta.split(","), "bedrock_mantle") != () -_CONVERSE_CALLER_BETAS: Final = ( - _INTERLEAVED_THINKING, - _PER_TURN_CONTROL, - ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value, - ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value, - ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, - "context-1m-2025-08-07", -) +_TOOL_EXAMPLES_BETA: Final = "tool-examples-2025-10-29" +_TOOLS_WITH_INPUT_EXAMPLES: Final[list[dict[str, JsonValue]]] = [ + {**_WEATHER_TOOL, "input_examples": [{"city": "Paris"}]} +] +_EFFORT_DIRECTIVE: Final[dict[str, JsonValue]] = {"effort": "low"} +_DIRECTIVE_ONLY_SYSTEM_TURN: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + {"role": "system", "content": [], "output_config": _EFFORT_DIRECTIVE}, + {"role": "user", "content": "Go on."}, +] +_SYSTEM_TURN_AFTER_USER: Final[list[dict[str, JsonValue]]] = [ + {"role": "user", "content": "Hi"}, + {"role": "system", "content": "Answer briefly", "output_config": _EFFORT_DIRECTIVE}, + {"role": "assistant", "content": "Hello"}, + {"role": "user", "content": "Go on."}, +] +_PER_TURN_CONTROL_OPT_IN: Final[dict[str, JsonValue]] = {"supports_per_turn_control": True} -async def _converse_inference_betas(caller_beta: str) -> tuple[str, ...]: - inference: Final = _Bedrock(runtime=(200, _CONVERSE_REPLY), mantle=(200, _CONVERSE_REPLY)) - await litellm.anthropic_messages( - model=_CONVERSE_DEPLOYMENT_MODEL, - messages=MESSAGES, - max_tokens=4096, - extra_headers={"anthropic-beta": caller_beta}, - client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), - **_THINKING, - **DEPLOYMENT["litellm_params"], - ) - (request,) = inference.requests - assert request.url.host == RUNTIME_HOST - sent: Final = _at(_JSON_BODY.validate_json(request.content), "additionalModelRequestFields", "anthropic_beta") - assert isinstance(sent, list) - return _sorted_betas(sent) +def _carries(body: Mapping[str, JsonValue], key: str, field: str) -> bool: + entries: Final = body.get(key) + return isinstance(entries, list) and any(isinstance(entry, dict) and field in entry for entry in entries) + + +def _counted_body_and_betas(request: httpx.Request) -> tuple[dict[str, JsonValue], tuple[str, ...]]: + if request.url.host == MANTLE_HOST: + return _JSON_BODY.validate_json(request.content), _sent_betas(request) + body: Final = _invoke_body(request) + return body, _sorted_betas(_betas_besides_effort(body)) @pytest.mark.asyncio @pytest.mark.usefixtures("local_beta_headers_config") @pytest.mark.parametrize( - ("model_to_use", "count_host", "count_row"), + ("model_to_use", "count_host", "model_info", "caller_beta", "messages", "tools", "carried"), ( - (_MANTLE_MODEL, MANTLE_HOST, "bedrock_mantle"), - ("us.anthropic.claude-opus-5-5", RUNTIME_HOST, "bedrock"), + ( + _MANTLE_MODEL, + MANTLE_HOST, + {}, + _TOOL_EXAMPLES_BETA, + MESSAGES, + _TOOLS_WITH_INPUT_EXAMPLES, + ("tools", "input_examples"), + ), + ( + _MANTLE_MODEL, + MANTLE_HOST, + {}, + _PER_TURN_CONTROL, + _DIRECTIVE_ONLY_SYSTEM_TURN, + None, + ("messages", "output_config"), + ), + ( + _MANTLE_MODEL, + MANTLE_HOST, + {}, + _PER_TURN_CONTROL, + _SYSTEM_TURN_AFTER_USER, + None, + ("messages", "output_config"), + ), + ( + "us.anthropic.claude-opus-5-5", + RUNTIME_HOST, + _PER_TURN_CONTROL_OPT_IN, + _PER_TURN_CONTROL, + _DIRECTIVE_ONLY_SYSTEM_TURN, + None, + ("messages", "output_config"), + ), + ), + ids=( + "mantle-tool-input-examples", + "mantle-directive-only-system-output-config", + "mantle-system-output-config-after-user", + "runtime-directive-only-system-output-config", ), - ids=("mantle", "runtime"), ) -async def test_count_for_a_converse_deployment_sends_only_the_caller_betas_converse_inference_sends( - model_to_use: str, count_host: str, count_row: str +async def test_count_for_a_converse_deployment_keeps_the_caller_beta_its_body_needs( + model_to_use: str, + count_host: str, + model_info: dict[str, JsonValue], + caller_beta: str, + messages: list[dict[str, JsonValue]], + tools: list[dict[str, JsonValue]] | None, + carried: tuple[str, str], ) -> None: - caller_beta: Final = ",".join(_CONVERSE_CALLER_BETAS) - sent: Final = await _converse_inference_betas(caller_beta) - host, counted = await _counted_betas( - model_to_use, - caller_beta, - {"litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL}}, + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model_to_use, + messages=messages, + contents=None, + deployment={ + "litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL}, + "model_info": model_info, + }, + request_model="claude", + tools=tools, + request_options=_THINKING, + anthropic_beta=caller_beta, ) - assert host == count_host - assert counted == _row_sends(sent, count_row) != () - assert counted != _row_sends(_CONVERSE_CALLER_BETAS, count_row) + assert result is not None and result.error is False + (request,) = bedrock.requests + body, betas = _counted_body_and_betas(request) + assert request.url.host == count_host + assert _carries(body, *carried) + assert caller_beta in betas