scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(count_tokens): keep the caller betas a Converse Bedrock deployment's count body needs

Browse source 
Filtering a Converse deployment's count betas through the bedrock_converse row
dropped betas the count body still needs, because the count keeps fields like
tools[].input_examples and a system message's output_config that Converse
inference strips. Mantle counts carrying those fields went from 200 to 400, so
the count goes back to filtering through the count API's own row
This commit is contained in:
mateo-berri 2026-11-07 05:52:57 -08:00
parent e9779083c1
commit 8b0a21b05e
2 changed files with 104 additions and 54 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,7 +12,6 @@ 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,
@ -138,12 +137,6 @@ 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"""
@ -210,20 +203,18 @@ 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)
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
route: Final = _deployment_route(litellm_params) if is_mantle_claude_model(resolved_model) else None
claude_options: Final = (
claude_count_options(
counted_model,
_claude_options_provider(route, use_mantle),
messages,
tools,
_options_converse_sends(request_options, inference_model),
_betas_converse_sends(anthropic_beta, inference_model),
_options_converse_sends(request_options, deployment_inference_model(deployment)),
anthropic_beta,
deployment,
)
if claude_model
if is_mantle_claude_model(resolved_model)
else NO_CLAUDE_COUNT_OPTIONS
)
if isinstance(claude_options, ChatCountTranslationError):

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

@ -903,7 +903,7 @@ def _message_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
@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:
async def test_mantle_count_for_a_converse_deployment_sends_converse_format_and_messages_with_mantle_betas() -> 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))
@ -942,58 +942,117 @@ async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas
"schema": json.loads(sent_schema),
}
assert _message_texts(counted) == _message_texts(sent)
assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != ()
assert _sent_betas(bedrock.requests[0]) == _row_sends(caller_beta.split(","), "bedrock_mantle") != ()
_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",
)
_TOOL_EXAMPLES_BETA: Final = "tool-examples-2025-10-29"
_TOOLS_WITH_INPUT_EXAMPLES: Final[list[dict[str, JsonValue]]] = [
{**_WEATHER_TOOL, "input_examples": [{"city": "Paris"}]}
]
_EFFORT_DIRECTIVE: Final[dict[str, JsonValue]] = {"effort": "low"}
_DIRECTIVE_ONLY_SYSTEM_TURN: Final[list[dict[str, JsonValue]]] = [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hello"},
{"role": "system", "content": [], "output_config": _EFFORT_DIRECTIVE},
{"role": "user", "content": "Go on."},
]
_SYSTEM_TURN_AFTER_USER: Final[list[dict[str, JsonValue]]] = [
{"role": "user", "content": "Hi"},
{"role": "system", "content": "Answer briefly", "output_config": _EFFORT_DIRECTIVE},
{"role": "assistant", "content": "Hello"},
{"role": "user", "content": "Go on."},
]
_PER_TURN_CONTROL_OPT_IN: Final[dict[str, JsonValue]] = {"supports_per_turn_control": True}
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)
def _carries(body: Mapping[str, JsonValue], key: str, field: str) -> bool:
entries: Final = body.get(key)
return isinstance(entries, list) and any(isinstance(entry, dict) and field in entry for entry in entries)
def _counted_body_and_betas(request: httpx.Request) -> tuple[dict[str, JsonValue], tuple[str, ...]]:
if request.url.host == MANTLE_HOST:
return _JSON_BODY.validate_json(request.content), _sent_betas(request)
body: Final = _invoke_body(request)
return body, _sorted_betas(_betas_besides_effort(body))
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize(
("model_to_use", "count_host", "count_row"),
("model_to_use", "count_host", "model_info", "caller_beta", "messages", "tools", "carried"),
(
(_MANTLE_MODEL, MANTLE_HOST, "bedrock_mantle"),
("us.anthropic.claude-opus-5-5", RUNTIME_HOST, "bedrock"),
(
_MANTLE_MODEL,
MANTLE_HOST,
{},
_TOOL_EXAMPLES_BETA,
MESSAGES,
_TOOLS_WITH_INPUT_EXAMPLES,
("tools", "input_examples"),
),
(
_MANTLE_MODEL,
MANTLE_HOST,
{},
_PER_TURN_CONTROL,
_DIRECTIVE_ONLY_SYSTEM_TURN,
None,
("messages", "output_config"),
),
(
_MANTLE_MODEL,
MANTLE_HOST,
{},
_PER_TURN_CONTROL,
_SYSTEM_TURN_AFTER_USER,
None,
("messages", "output_config"),
),
(
"us.anthropic.claude-opus-5-5",
RUNTIME_HOST,
_PER_TURN_CONTROL_OPT_IN,
_PER_TURN_CONTROL,
_DIRECTIVE_ONLY_SYSTEM_TURN,
None,
("messages", "output_config"),
),
),
ids=(
"mantle-tool-input-examples",
"mantle-directive-only-system-output-config",
"mantle-system-output-config-after-user",
"runtime-directive-only-system-output-config",
),
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
async def test_count_for_a_converse_deployment_keeps_the_caller_beta_its_body_needs(
model_to_use: str,
count_host: str,
model_info: dict[str, JsonValue],
caller_beta: str,
messages: list[dict[str, JsonValue]],
tools: list[dict[str, JsonValue]] | None,
carried: tuple[str, 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}},
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={
"litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL},
"model_info": model_info,
},
request_model="claude",
tools=tools,
request_options=_THINKING,
anthropic_beta=caller_beta,
)
assert host == count_host
assert counted == _row_sends(sent, count_row) != ()
assert counted != _row_sends(_CONVERSE_CALLER_BETAS, count_row)
assert result is not None and result.error is False
(request,) = bedrock.requests
body, betas = _counted_body_and_betas(request)
assert request.url.host == count_host
assert _carries(body, *carried)
assert caller_beta in betas