scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(count_tokens): send a Converse Bedrock deployment's count only the betas Converse inference sends

Browse source 
A bedrock/converse/ Claude deployment's count filtered the caller's anthropic-beta with the bedrock_mantle row on Mantle and the bedrock row (plus the model's kept betas) on Bedrock Runtime, while Converse inference filters it with the bedrock_converse row. So the count still sent betas that inference drops, such as interleaved-thinking-2025-05-14, per-turn-control-2026-07-01 and web-search-2025-03-05 on Mantle, and per-turn-control-2026-07-01 and mid-conversation-tool-changes-2026-07-01 on Runtime. When the deployment's inference runs on Converse, the count now runs the caller's betas through the same bedrock_converse filter before it builds its options, the same way it already drops disable_parallel_tool_use for Converse. Betas the count adds for fields it carries are unchanged, and invoke deployments, bedrock_mantle deployments and counts with no deployment are untouched

The new Converse test derives the expected betas from what Converse inference sends for the same caller betas, limited by the count API's own row, on both Mantle and Runtime. The proxy retry table now asserts which hosts got the count, in order, instead of only the call count. The caller-beta tests derive which betas a row keeps from the beta config through the manager's filter instead of pinning it, and compare betas as sorted tuples so a beta sent twice fails
This commit is contained in:
mateo-berri 2026-11-07 04:27:56 -08:00
parent 3025f7f445
commit f09855a540
3 changed files with 136 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -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):

103
tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
View file

@ -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)

51
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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(