diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 29d2f2201e7..06f879f3d40 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -12,6 +12,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import DeploymentBetaFilter from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits @@ -46,6 +47,7 @@ from litellm.types.utils import ( RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" _RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404}) +_CONVERSE_BETA_FILTER: Final = DeploymentBetaFilter(provider="bedrock_converse", api_base=None, unlisted="drop") class _CountTokensReply(LiteLLMPydanticObjectBase): @@ -167,6 +169,12 @@ def _options_converse_sends( } +def _beta_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 _CONVERSE_BETA_FILTER.filtered({"anthropic-beta": anthropic_beta}, keep=frozenset()).get("anthropic-beta") + + def _request_data( model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object ) -> Mapping[str, object]: @@ -284,14 +292,17 @@ class BedrockTokenCounter(BaseTokenCounter): return translated.response(request_model, model_to_use, primary_tokenizer) sent_messages: Final = messages if translated is None else translated.messages sent_tools: Final = tools if translated is None else translated.tools + inference_model: Final = ( + deployment_inference_model(deployment) if is_mantle_claude_model(resolved_model) else None + ) claude_options: Final = ( claude_count_options( counted_model, _claude_options_provider(route, use_mantle), sent_messages, sent_tools, - _options_converse_sends(request_options, deployment_inference_model(deployment)), - anthropic_beta, + _options_converse_sends(request_options, inference_model), + _beta_converse_sends(anthropic_beta, inference_model), deployment, ) if is_mantle_claude_model(resolved_model) 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 383d2093612..bdd40ed4f78 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 @@ -2,7 +2,7 @@ import base64 import functools import itertools import json -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from typing import Final import httpx @@ -769,7 +769,7 @@ _MANTLE_COUNT_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" @pytest.mark.asyncio @pytest.mark.usefixtures("local_beta_headers_config", "aws_credentials_in_the_environment") -async def test_runtime_count_without_a_deployment_drops_the_caller_betas_invoke_inference_drops() -> None: +async def test_runtime_count_without_a_deployment_sends_the_caller_betas_invoke_inference_sends() -> None: sent: Final = await _invoke_inference_betas(_INTERLEAVED_THINKING, {}) bedrock: Final = await _count_with_caller_beta("us.anthropic.claude-opus-5-5", _INTERLEAVED_THINKING, None) @@ -866,15 +866,19 @@ def _at(value: JsonValue, *path: str) -> JsonValue: return functools.reduce(lambda node, key: node.get(key) if isinstance(node, dict) else None, path, value) -def _content_texts(message: JsonValue) -> tuple[str, ...]: +def _texts(message: JsonValue) -> tuple[str, ...]: content: Final = message.get("content") if isinstance(message, dict) else None - return (content,) if isinstance(content, str) else _texts(message) + if isinstance(content, str): + return (content,) + if not isinstance(content, list): + return () + return tuple(text for block in content if isinstance(block, dict) and isinstance(text := block.get("text"), str)) def _message_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: messages: Final = body.get("messages") assert isinstance(messages, list) - return tuple(itertools.chain.from_iterable(map(_content_texts, messages))) + return tuple(itertools.chain.from_iterable(map(_texts, messages))) @pytest.mark.asyncio @@ -901,6 +905,78 @@ async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != () +_CONVERSE_DEPLOYMENT_MODEL: Final = "bedrock/converse/us.anthropic.claude-opus-5-5" +_CALLER_BETAS_THE_ROWS_SPLIT: Final = ",".join( + ( + _INTERLEAVED_THINKING, + "token-efficient-tools-2025-02-19", + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, + ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value, + "context-1m-2025-08-07", + ) +) +_BETAS: Final = TypeAdapter(frozenset[str]) + + +async def _converse_inference_betas(caller_beta: str) -> frozenset[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"], + ) + assert inference.posted_hosts() == (RUNTIME_HOST,) + return frozenset(_sent_betas(inference.requests[0])) + + +async def _invoke_api_inference_betas(caller_beta: str) -> frozenset[str]: + return _BETAS.validate_python(await _invoke_inference_betas(caller_beta, {})) + + +def _invoke_api_counted_betas(bedrock: _Bedrock) -> frozenset[str]: + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + return _BETAS.validate_python(_betas_besides_effort(_invoke_body(bedrock.requests[0]))) + + +async def _mantle_api_inference_betas(caller_beta: str) -> frozenset[str]: + return await _mantle_inference_betas(caller_beta, {"model": _MANTLE_COUNT_MODEL}) + + +def _mantle_api_counted_betas(bedrock: _Bedrock) -> frozenset[str]: + assert bedrock.posted_hosts() == (MANTLE_HOST,) + return frozenset(_sent_betas(bedrock.requests[0])) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("model_to_use", "count_api_inference_betas", "counted_betas"), + ( + ("us.anthropic.claude-opus-5-5", _invoke_api_inference_betas, _invoke_api_counted_betas), + (_MANTLE_COUNT_MODEL, _mantle_api_inference_betas, _mantle_api_counted_betas), + ), + ids=("runtime", "mantle"), +) +async def test_count_for_a_converse_deployment_sends_only_the_caller_betas_converse_inference_sends( + model_to_use: str, + count_api_inference_betas: Callable[[str], Awaitable[frozenset[str]]], + counted_betas: Callable[[_Bedrock], frozenset[str]], +) -> None: + converse_sent: Final = await _converse_inference_betas(_CALLER_BETAS_THE_ROWS_SPLIT) + bedrock: Final = await _count_with_caller_beta( + model_to_use, + _CALLER_BETAS_THE_ROWS_SPLIT, + {"litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL}}, + ) + + assert await count_api_inference_betas(_CALLER_BETAS_THE_ROWS_SPLIT) != converse_sent + assert counted_betas(bedrock) == await count_api_inference_betas(",".join(sorted(converse_sent))) + + @pytest.mark.asyncio async def test_count_on_a_stand_in_model_counts_the_messages_and_output_config_the_deployment_model_sends() -> None: inference, bedrock = await _messages_inference_and_count( @@ -1055,13 +1131,6 @@ async def test_runtime_count_sends_the_tools_invoke_inference_sends() -> None: assert _invoke_body(bedrock.requests[0])["tools"] == sent_tools -def _texts(message: JsonValue) -> tuple[str, ...]: - content: Final = message.get("content") if isinstance(message, dict) else None - if not isinstance(content, list): - return () - return tuple(text for block in content if isinstance(block, dict) and isinstance(text := block.get("text"), str)) - - @pytest.mark.asyncio @pytest.mark.parametrize( ("native_output", "request_options", "schema_inlined"), diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 4b6150cda40..9072d28d32b 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1961,7 +1961,7 @@ async def test_anthropic_count_tokens_does_not_count_again_after_a_rate_limit( _BEDROCK_RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" _BEDROCK_MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" -_BEDROCK_CLAUDE_DEPLOYMENT: Final[dict[str, str]] = { +_BEDROCK_CLAUDE_DEPLOYMENT: Final[Mapping[str, str]] = { "model": "bedrock/us.anthropic.claude-count-probe-v1:0", "aws_access_key_id": "AKIATESTACCESSKEY", "aws_secret_access_key": "test-secret",