diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 1e8c28abd46..437d6705eca 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -10,7 +10,9 @@ from typing import Any, Final from pydantic import JsonValue from typing_extensions import assert_never +import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.count_tokens.transformation import ( NO_CLAUDE_COUNT_OPTIONS, ChatCountTranslationError, @@ -60,6 +62,31 @@ class CountTokensFailure: CountTokensOutcome = CountedTokens | CountTokensFailure +def _deployment_litellm_params(deployment: Mapping[str, object] | None) -> Mapping[str, object]: + return as_str_mapping((deployment or {}).get("litellm_params")) or MappingProxyType({}) + + +def _deployment_route(litellm_params: Mapping[str, object]) -> tuple[str, str] | None: + deployment_model: Final = litellm_params.get("model") + if not isinstance(deployment_model, str): + return None + provider_hint: Final = litellm_params.get("custom_llm_provider") + model, provider, _, _ = litellm.get_llm_provider( + model=deployment_model, custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None + ) + return model, provider + + +def _claude_options_provider(route: tuple[str, str] | None, use_mantle: bool) -> LlmProviders: + if not use_mantle: + return LlmProviders.BEDROCK + if route is None: + return LlmProviders.BEDROCK_MANTLE + model, provider = route + runs_on_invoke: Final = provider == LlmProviders.BEDROCK.value and not _runs_on_converse(model) + return LlmProviders.BEDROCK if runs_on_invoke else LlmProviders.BEDROCK_MANTLE + + def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) @@ -175,11 +202,12 @@ class BedrockTokenCounter(BaseTokenCounter): counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") 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 or {}).get("litellm_params", {}) + litellm_params: Final = _deployment_litellm_params(deployment) + route: Final = _deployment_route(litellm_params) if is_mantle_claude_model(resolved_model) else None claude_options: Final = ( claude_count_options( counted_model, - LlmProviders.BEDROCK_MANTLE if use_mantle else LlmProviders.BEDROCK, + _claude_options_provider(route, use_mantle), messages, tools, _options_converse_sends(request_options, deployment_inference_model(deployment)), @@ -215,7 +243,7 @@ class BedrockTokenCounter(BaseTokenCounter): else await self._count_with( self._runtime_handler.handle_count_tokens_request( request_data=runtime_request, - litellm_params=litellm_params, + litellm_params=dict(litellm_params), resolved_model=resolved_model, client=self._client, input_type="invokeModel" if invoke_fields else None, @@ -232,7 +260,8 @@ class BedrockTokenCounter(BaseTokenCounter): ) outcome: Final = ( fallback - if fallback is not None and _mantle_answered_for_runtime(fallback, forwarded_options=bool(invoke_fields)) + if fallback is not None + and _mantle_answered_for_runtime(fallback, forwarded_options=claude_options.carries_request_options) else primary ) match outcome: 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 b729aa2a4e4..796ce478735 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,4 +1,6 @@ import base64 +import functools +import itertools import json from collections.abc import Callable, Mapping from typing import Final @@ -13,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_EFFORT_BETA_HEADER, + ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER, ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER, ) @@ -756,3 +759,171 @@ async def test_claude_count_refuses_options_bedrock_messages_inference_refuses_w assert result is not None assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) assert bedrock.posted_hosts() == () + + +_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" + + +async def _inference_betas(model: str, caller_beta: str, model_info: Mapping[str, JsonValue]) -> frozenset[JsonValue]: + inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY)) + await litellm.anthropic_messages( + model=model, + messages=MESSAGES, + max_tokens=4096, + extra_headers={"anthropic-beta": caller_beta}, + model_info=dict(model_info), + client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), + **_THINKING, + **DEPLOYMENT["litellm_params"], + ) + 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))) + + +async def _counted_betas( + model_to_use: str, caller_beta: str, deployment: Mapping[str, JsonValue] | None +) -> tuple[str, frozenset[JsonValue]]: + 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=None if deployment is None else dict(deployment), + request_model="claude", + request_options=_THINKING, + anthropic_beta=caller_beta, + ) + 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))) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("caller_beta", "model_info", "deployment_route", "kept"), + ( + (_INTERLEAVED_THINKING, {}, {"model": _INVOKE_DEPLOYMENT_MODEL}, False), + (_PER_TURN_CONTROL, {"supports_per_turn_control": True}, {"model": _INVOKE_DEPLOYMENT_MODEL}, True), + (_PER_TURN_CONTROL, {"supports_per_turn_control": False}, {"model": _INVOKE_DEPLOYMENT_MODEL}, False), + ( + _INTERLEAVED_THINKING, + {}, + {"model": _INVOKE_DEPLOYMENT_MODEL.removeprefix("bedrock/"), "custom_llm_provider": "bedrock"}, + False, + ), + ), + ids=( + "interleaved-thinking", + "deployment-supports-per-turn-control", + "deployment-lacks-per-turn-control", + "provider-set-by-custom-llm-provider", + ), +) +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 +) -> None: + sent: Final = await _inference_betas(_INVOKE_DEPLOYMENT_MODEL, caller_beta, model_info) + deployment: Final[Mapping[str, JsonValue]] = { + "litellm_params": {**DEPLOYMENT["litellm_params"], **deployment_route}, + "model_info": dict(model_info), + } + host, counted = await _counted_betas(_MANTLE_MODEL, caller_beta, deployment) + + assert host == MANTLE_HOST + assert counted == sent + assert (caller_beta in counted) is kept + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("model_to_use", "inference_model", "count_host", "kept"), + ( + ("us.anthropic.claude-opus-5-5", _INVOKE_DEPLOYMENT_MODEL, RUNTIME_HOST, False), + (_MANTLE_MODEL, _MANTLE_MODEL, MANTLE_HOST, True), + ), + 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 +) -> None: + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + sent: Final = await _inference_betas(inference_model, _INTERLEAVED_THINKING, {}) + host, counted = await _counted_betas(model_to_use, _INTERLEAVED_THINKING, None) + + assert host == count_host + assert counted == sent + assert (_INTERLEAVED_THINKING in counted) is kept + + +_CONVERSE_DEPLOYMENT_MODEL: Final = "bedrock/converse/us.anthropic.claude-opus-5-5" +_NATIVE_OUTPUT_OPT_IN: Final[Mapping[str, JsonValue]] = {"supports_bedrock_native_structured_output": True} + + +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, ...]: + content: Final = message.get("content") if isinstance(message, dict) else None + return (content,) if isinstance(content, str) else _texts(message) + + +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))) + + +@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: + 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)) + await litellm.anthropic_messages( + model=_CONVERSE_DEPLOYMENT_MODEL, + messages=MESSAGES, + max_tokens=4096, + extra_headers={"anthropic-beta": caller_beta}, + model_info=dict(_NATIVE_OUTPUT_OPT_IN), + client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), + **request_options, + **DEPLOYMENT["litellm_params"], + ) + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + result: Final = await _counter(bedrock).count_tokens( + model_to_use=_MANTLE_MODEL, + messages=MESSAGES, + contents=None, + deployment={ + "litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL}, + "model_info": dict(_NATIVE_OUTPUT_OPT_IN), + }, + request_model="claude", + request_options=request_options, + anthropic_beta=caller_beta, + ) + + assert result is not None and result.error is False + assert (inference.posted_hosts(), bedrock.posted_hosts()) == ((RUNTIME_HOST,), (MANTLE_HOST,)) + sent: Final = _JSON_BODY.validate_json(inference.requests[0].content) + counted: Final = _JSON_BODY.validate_json(bedrock.requests[0].content) + sent_schema: Final = _at(sent, "outputConfig", "textFormat", "structure", "jsonSchema", "schema") + assert isinstance(sent_schema, str) + assert _at(counted, "output_config", "format") == { + "type": _at(sent, "outputConfig", "textFormat", "type"), + "schema": json.loads(sent_schema), + } + assert _message_texts(counted) == _message_texts(sent) + assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != () diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 1d19524c665..6d360dc776c 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1569,50 +1569,97 @@ async def test_anthropic_count_tokens_does_not_count_again_after_a_rate_limit( assert len(requests) == 1 +_BEDROCK_RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" +_BEDROCK_MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" +_AWS_CREDENTIALS: Final[Mapping[str, str]] = { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", +} +_BEDROCK_CLAUDE_DEPLOYMENT: Final[Mapping[str, str]] = { + "model": "bedrock/us.anthropic.claude-count-probe-v1:0", + **_AWS_CREDENTIALS, +} +_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_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { + "model": "anthropic/claude-haiku-5-5", + "api_key": "sk-ant-test", +} +_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_key": "sk-azure-test", +} +_VERTEX_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { + "model": "vertex_ai/claude-haiku-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", +} + + @pytest.mark.asyncio @pytest.mark.parametrize( - ("deployment_model", "deployment_params", "count_host", "provider_calls"), + ("deployment_params", "body", "caller_beta", "provider_calls"), ( - ("gemini/gemini-3.8-flash", {}, "generativelanguage.googleapis.com", 1), - ("anthropic/claude-haiku-5-5", {}, "api.anthropic.com", 2), - ( - "azure_ai/claude-haiku-5-5", - {"api_base": "https://example-resource.services.ai.azure.com/anthropic"}, - "example-resource.services.ai.azure.com", - 2, - ), + ({"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), ), ids=( "counter-without-claude-options", - "claude-counter-forwarding-the-options", - "azure-ai-claude-counter-forwarding-the-options", + "bedrock-counter-without-claude-options", + "anthropic-forwarding-the-options", + "anthropic-forwarding-nothing", + "anthropic-forwarding-only-the-caller-beta", + "azure-ai-forwarding-the-options", + "azure-ai-forwarding-nothing", + "vertex-forwarding-the-options", + "vertex-forwarding-nothing", + "bedrock-forwarding-nothing-counts-on-runtime-then-mantle-once", ), ) async def test_anthropic_count_tokens_counts_again_without_the_options_only_when_the_counter_forwarded_them( monkeypatch: pytest.MonkeyPatch, - deployment_model: str, - deployment_params: dict[str, JsonValue], - count_host: str, + deployment_params: Mapping[str, JsonValue], + body: Mapping[str, JsonValue], + caller_beta: str | None, provider_calls: int, ) -> None: from litellm.proxy import proxy_server - request: Final = _json_request( - { - "model": "claude-agent", - "messages": [{"role": "user", "content": "Weather in SF?"}], - "thinking": {"type": "enabled", "budget_tokens": 2000}, - }, - "python-httpx/0.28.1", + for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL", "AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr( + "google.auth.default", lambda *_, **__: (MagicMock(token="fake-token", valid=True, expiry=None), "test-project") ) - monkeypatch.setattr(proxy_server, "llm_router", _claude_router(deployment_model, deployment_params)) + request: Final = _json_request( + {"model": "claude-agent", "messages": [{"role": "user", "content": "Weather in SF?"}], **body}, + "python-httpx/0.28.1", + {} if caller_beta is None else {"anthropic-beta": caller_beta}, + ) + router: Final = Router(model_list=[{"model_name": "claude-agent", "litellm_params": dict(deployment_params)}]) + monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig()) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) with respx.mock: - route: Final = respx.route(host=count_host).mock( + route: Final = respx.route().mock( return_value=httpx.Response( - 400, json={"type": "error", "error": {"type": "invalid_request_error", "message": "rejected"}} + 400, + json={ + "type": "error", + "message": "rejected", + "error": {"type": "invalid_request_error", "message": "rejected"}, + }, ) ) response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth()) @@ -1622,16 +1669,6 @@ async def test_anthropic_count_tokens_counts_again_without_the_options_only_when assert route.call_count == provider_calls -_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]] = { - "model": "bedrock/us.anthropic.claude-count-probe-v1:0", - "aws_access_key_id": "AKIATESTACCESSKEY", - "aws_secret_access_key": "test-secret", - "aws_region_name": "us-east-1", -} - - async def _count_on_a_bedrock_claude_deployment( monkeypatch: pytest.MonkeyPatch, body: Mapping[str, JsonValue], @@ -1642,7 +1679,9 @@ async def _count_on_a_bedrock_claude_deployment( monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) - router: Final = Router(model_list=[{"model_name": "claude-agent", "litellm_params": _BEDROCK_CLAUDE_DEPLOYMENT}]) + router: Final = Router( + model_list=[{"model_name": "claude-agent", "litellm_params": dict(_BEDROCK_CLAUDE_DEPLOYMENT)}] + ) request: Final = _json_request( {"model": "claude-agent", "messages": [{"role": "user", "content": "Weather in SF?"}], **body}, "python-httpx/0.28.1",