scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(count_tokens): filter a Bedrock invoke deployment's Mantle count betas like its inference, and test every counter's retry flag

Browse source 
A bedrock/ Claude deployment that counts on Mantle (the cost map's token_count_model for us.anthropic.claude-opus-5-5 and its siblings) built its count options with the Mantle messages config, so a caller beta that Bedrock invoke inference drops, such as interleaved-thinking-2025-05-14, still reached the Mantle count. When the deployment's /v1/messages runs on invoke, the count now builds its options and betas with the Bedrock invoke config, the same filter its inference applies. A bedrock/converse/ deployment, a bedrock_mantle deployment and a count with no deployment keep the Mantle config, the same split main uses

The Bedrock counter also reads carries_request_options for the Mantle fallback instead of recomputing the same flag from its invoke fields, and it validates the deployment's litellm_params into a typed mapping once per count

The proxy retry test now covers the Vertex AI partner, Anthropic and Azure AI counters forwarding options or nothing, a betas-only count, a Claude Bedrock count that forwards nothing, and a non-Claude Bedrock count, so each counter's retry flag is pinned in both directions. New Bedrock tests pin the Mantle count betas for invoke, Converse, Mantle and deployment-less counts against what each route's inference sends
This commit is contained in:
mateo-berri 2026-11-07 03:31:42 -08:00
parent af7e65b86e
commit 3025f7f445
3 changed files with 278 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

@ -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]) != ()

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

@ -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",