diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 437d6705eca..690e13e9d21 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 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, @@ -137,6 +138,12 @@ 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""" @@ -203,18 +210,20 @@ 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) - route: Final = _deployment_route(litellm_params) if is_mantle_claude_model(resolved_model) else None + 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 claude_options: Final = ( claude_count_options( counted_model, _claude_options_provider(route, use_mantle), messages, tools, - _options_converse_sends(request_options, deployment_inference_model(deployment)), - anthropic_beta, + _options_converse_sends(request_options, inference_model), + _betas_converse_sends(anthropic_beta, inference_model), deployment, ) - if is_mantle_claude_model(resolved_model) + if claude_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 796ce478735..ddc6421e9fb 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 @@ -1,8 +1,7 @@ import base64 -import functools import itertools import json -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from typing import Final import httpx @@ -10,6 +9,7 @@ import pytest from pydantic import JsonValue, TypeAdapter import litellm +from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.llms.anthropic import ( @@ -765,9 +765,22 @@ _INTERLEAVED_THINKING: Final = "interleaved-thinking-2025-05-14" _PER_TURN_CONTROL: Final = ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value _INVOKE_DEPLOYMENT_MODEL: Final = "bedrock/invoke/us.anthropic.claude-opus-5-5" _MANTLE_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" +_BETAS: Final = TypeAdapter(tuple[str, ...]) -async def _inference_betas(model: str, caller_beta: str, model_info: Mapping[str, JsonValue]) -> frozenset[JsonValue]: +def _sorted_betas(betas: Sequence[JsonValue]) -> tuple[str, ...]: + return tuple(sorted(_BETAS.validate_python(betas))) + + +def _row_sends(betas: Sequence[str], row: str) -> tuple[str, ...]: + return tuple(sorted(filter_and_transform_beta_headers(list(betas), row))) + + +def _row_keeps(beta: str, row: str) -> bool: + return beta in _row_sends((beta,), row) + + +async def _inference_betas(model: str, caller_beta: str, model_info: Mapping[str, JsonValue]) -> tuple[str, ...]: inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY)) await litellm.anthropic_messages( model=model, @@ -781,13 +794,13 @@ async def _inference_betas(model: str, caller_beta: str, model_info: Mapping[str ) request: Final = inference.requests[0] if request.url.host == MANTLE_HOST: - return frozenset(_sent_betas(request)) - return frozenset(_betas_besides_effort(_JSON_BODY.validate_json(request.content))) + return _sent_betas(request) + return _sorted_betas(_betas_besides_effort(_JSON_BODY.validate_json(request.content))) async def _counted_betas( model_to_use: str, caller_beta: str, deployment: Mapping[str, JsonValue] | None -) -> tuple[str, frozenset[JsonValue]]: +) -> tuple[str, tuple[str, ...]]: 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, @@ -801,14 +814,14 @@ async def _counted_betas( assert result is not None and result.error is False (request,) = bedrock.requests if request.url.host == MANTLE_HOST: - return request.url.host, frozenset(_sent_betas(request)) - return request.url.host, frozenset(_betas_besides_effort(_invoke_body(request))) + return request.url.host, _sent_betas(request) + return request.url.host, _sorted_betas(_betas_besides_effort(_invoke_body(request))) @pytest.mark.asyncio @pytest.mark.usefixtures("local_beta_headers_config") @pytest.mark.parametrize( - ("caller_beta", "model_info", "deployment_route", "kept"), + ("caller_beta", "model_info", "deployment_route", "opted_in"), ( (_INTERLEAVED_THINKING, {}, {"model": _INVOKE_DEPLOYMENT_MODEL}, False), (_PER_TURN_CONTROL, {"supports_per_turn_control": True}, {"model": _INVOKE_DEPLOYMENT_MODEL}, True), @@ -828,7 +841,7 @@ async def _counted_betas( ), ) async def test_mantle_count_for_a_bedrock_deployment_sends_the_caller_betas_invoke_inference_sends( - caller_beta: str, model_info: Mapping[str, JsonValue], deployment_route: Mapping[str, str], kept: bool + caller_beta: str, model_info: Mapping[str, JsonValue], deployment_route: Mapping[str, str], opted_in: bool ) -> None: sent: Final = await _inference_betas(_INVOKE_DEPLOYMENT_MODEL, caller_beta, model_info) deployment: Final[Mapping[str, JsonValue]] = { @@ -839,21 +852,21 @@ async def test_mantle_count_for_a_bedrock_deployment_sends_the_caller_betas_invo assert host == MANTLE_HOST assert counted == sent - assert (caller_beta in counted) is kept + assert (caller_beta in counted) is (opted_in or _row_keeps(caller_beta, "bedrock")) @pytest.mark.asyncio @pytest.mark.usefixtures("local_beta_headers_config") @pytest.mark.parametrize( - ("model_to_use", "inference_model", "count_host", "kept"), + ("model_to_use", "inference_model", "count_host", "count_row"), ( - ("us.anthropic.claude-opus-5-5", _INVOKE_DEPLOYMENT_MODEL, RUNTIME_HOST, False), - (_MANTLE_MODEL, _MANTLE_MODEL, MANTLE_HOST, True), + ("us.anthropic.claude-opus-5-5", _INVOKE_DEPLOYMENT_MODEL, RUNTIME_HOST, "bedrock"), + (_MANTLE_MODEL, _MANTLE_MODEL, MANTLE_HOST, "bedrock_mantle"), ), ids=("runtime", "mantle"), ) async def test_count_without_a_deployment_sends_the_caller_beta_inference_on_its_route_sends( - monkeypatch: pytest.MonkeyPatch, model_to_use: str, inference_model: str, count_host: str, kept: bool + monkeypatch: pytest.MonkeyPatch, model_to_use: str, inference_model: str, count_host: str, count_row: str ) -> None: monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") @@ -863,7 +876,7 @@ async def test_count_without_a_deployment_sends_the_caller_beta_inference_on_its assert host == count_host assert counted == sent - assert (_INTERLEAVED_THINKING in counted) is kept + assert (_INTERLEAVED_THINKING in counted) is _row_keeps(_INTERLEAVED_THINKING, count_row) _CONVERSE_DEPLOYMENT_MODEL: Final = "bedrock/converse/us.anthropic.claude-opus-5-5" @@ -871,7 +884,9 @@ _NATIVE_OUTPUT_OPT_IN: Final[Mapping[str, JsonValue]] = {"supports_bedrock_nativ def _at(value: JsonValue, *path: str) -> JsonValue: - return functools.reduce(lambda node, key: node.get(key) if isinstance(node, dict) else None, path, value) + if not path: + return value + return _at(value.get(path[0]), *path[1:]) if isinstance(value, dict) else None def _content_texts(message: JsonValue) -> tuple[str, ...]: @@ -927,3 +942,57 @@ async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas } assert _message_texts(counted) == _message_texts(sent) assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != () + + +_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", +) + + +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) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("model_to_use", "count_host", "count_row"), + ( + (_MANTLE_MODEL, MANTLE_HOST, "bedrock_mantle"), + ("us.anthropic.claude-opus-5-5", RUNTIME_HOST, "bedrock"), + ), + 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 +) -> 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}}, + ) + + assert host == count_host + assert counted == _row_sends(sent, count_row) != () + assert counted != _row_sends(_CONVERSE_CALLER_BETAS, count_row) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 6d360dc776c..2fa3a2d28e8 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1583,15 +1583,18 @@ _BEDROCK_CLAUDE_DEPLOYMENT: Final[Mapping[str, str]] = { _THINKING_BODY: Final[Mapping[str, JsonValue]] = {"thinking": {"type": "enabled", "budget_tokens": 2000}} _EFFORT_BODY: Final[Mapping[str, JsonValue]] = {"output_config": {"effort": "low"}} _DROPS_THE_EFFORT: Final[Mapping[str, JsonValue]] = {"additional_drop_params": ["output_config.effort"]} +_ANTHROPIC_HOST: Final = "api.anthropic.com" _ANTHROPIC_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { "model": "anthropic/claude-haiku-5-5", "api_key": "sk-ant-test", } +_AZURE_AI_HOST: Final = "example-resource.services.ai.azure.com" _AZURE_AI_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { "model": "azure_ai/claude-haiku-5-5", - "api_base": "https://example-resource.services.ai.azure.com/anthropic", + "api_base": f"https://{_AZURE_AI_HOST}/anthropic", "api_key": "sk-azure-test", } +_VERTEX_HOST: Final = "us-east5-aiplatform.googleapis.com" _VERTEX_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { "model": "vertex_ai/claude-haiku-5-5", "vertex_project": "test-project", @@ -1601,18 +1604,38 @@ _VERTEX_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { @pytest.mark.asyncio @pytest.mark.parametrize( - ("deployment_params", "body", "caller_beta", "provider_calls"), + ("deployment_params", "body", "caller_beta", "counted_hosts"), ( - ({"model": "gemini/gemini-3.8-flash", "api_key": "sk-gemini-test"}, _THINKING_BODY, None, 1), - ({"model": "bedrock/us.amazon.nova-2-lite-v1:0", **_AWS_CREDENTIALS}, _THINKING_BODY, None, 1), - (_ANTHROPIC_DEPLOYMENT, _THINKING_BODY, None, 2), - ({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1), - ({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, _CALLER_BETA, 2), - (_AZURE_AI_DEPLOYMENT, _THINKING_BODY, None, 2), - ({**_AZURE_AI_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1), - (_VERTEX_DEPLOYMENT, _THINKING_BODY, None, 2), - ({**_VERTEX_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 1), - ({**_BEDROCK_CLAUDE_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, 2), + ( + {"model": "gemini/gemini-3.8-flash", "api_key": "sk-gemini-test"}, + _THINKING_BODY, + None, + ("generativelanguage.googleapis.com",), + ), + ( + {"model": "bedrock/us.amazon.nova-2-lite-v1:0", **_AWS_CREDENTIALS}, + _THINKING_BODY, + None, + (_BEDROCK_RUNTIME_HOST,), + ), + (_ANTHROPIC_DEPLOYMENT, _THINKING_BODY, None, (_ANTHROPIC_HOST, _ANTHROPIC_HOST)), + ({**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_ANTHROPIC_HOST,)), + ( + {**_ANTHROPIC_DEPLOYMENT, **_DROPS_THE_EFFORT}, + _EFFORT_BODY, + _CALLER_BETA, + (_ANTHROPIC_HOST, _ANTHROPIC_HOST), + ), + (_AZURE_AI_DEPLOYMENT, _THINKING_BODY, None, (_AZURE_AI_HOST, _AZURE_AI_HOST)), + ({**_AZURE_AI_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_AZURE_AI_HOST,)), + (_VERTEX_DEPLOYMENT, _THINKING_BODY, None, (_VERTEX_HOST, _VERTEX_HOST)), + ({**_VERTEX_DEPLOYMENT, **_DROPS_THE_EFFORT}, _EFFORT_BODY, None, (_VERTEX_HOST,)), + ( + {**_BEDROCK_CLAUDE_DEPLOYMENT, **_DROPS_THE_EFFORT}, + _EFFORT_BODY, + None, + (_BEDROCK_RUNTIME_HOST, _BEDROCK_MANTLE_HOST), + ), ), ids=( "counter-without-claude-options", @@ -1632,7 +1655,7 @@ async def test_anthropic_count_tokens_counts_again_without_the_options_only_when deployment_params: Mapping[str, JsonValue], body: Mapping[str, JsonValue], caller_beta: str | None, - provider_calls: int, + counted_hosts: tuple[str, ...], ) -> None: from litellm.proxy import proxy_server @@ -1666,7 +1689,7 @@ async def test_anthropic_count_tokens_counts_again_without_the_options_only_when assert isinstance(response, JSONResponse) assert response.status_code == 400 - assert route.call_count == provider_calls + assert tuple(call.request.url.host for call in route.calls) == counted_hosts async def _count_on_a_bedrock_claude_deployment(