mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(count_tokens): send a Converse Bedrock deployment's count only the betas Converse inference sends
A bedrock/converse/ deployment, or a bedrock/anthropic/ deployment that runs on Converse, counted with the caller's anthropic-beta filtered only by the count API's own row. Converse inference first drops what the bedrock_converse row does not list, so the count sent betas like interleaved-thinking-2025-05-14 and token-efficient-tools-2025-02-19 on Mantle, and mid-conversation-tool-changes-2026-07-01 on Runtime, that inference never sends When the deployment's inference runs on Converse, the count now passes the caller's betas through the same bedrock_converse filter Converse inference uses before it builds its options. The count API's row still applies after that, so a beta Converse keeps but the count API drops is not sent. Betas the count adds for fields it carries, invoke deployments, bedrock_mantle deployments and counts with no deployment are unchanged A new test counts a Converse deployment on Runtime and on Mantle with caller betas the rows disagree on, and expects what Converse inference sent, run through the count API's own inference. The runtime test with no deployment is renamed to say what it checks, the format test's text helper is folded into _texts, and the proxy test's Bedrock deployment is a read-only Mapping
This commit is contained in:
parent
64c658486c
commit
d41a172eba
3 changed files with 95 additions and 15 deletions
|
|
@ -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 DeploymentBetaFilter
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits
|
||||
|
|
@ -46,6 +47,7 @@ from litellm.types.utils import (
|
|||
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
|
||||
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
|
||||
_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
|
||||
_CONVERSE_BETA_FILTER: Final = DeploymentBetaFilter(provider="bedrock_converse", api_base=None, unlisted="drop")
|
||||
|
||||
|
||||
class _CountTokensReply(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -167,6 +169,12 @@ def _options_converse_sends(
|
|||
}
|
||||
|
||||
|
||||
def _beta_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 _CONVERSE_BETA_FILTER.filtered({"anthropic-beta": anthropic_beta}, keep=frozenset()).get("anthropic-beta")
|
||||
|
||||
|
||||
def _request_data(
|
||||
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -284,14 +292,17 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
return translated.response(request_model, model_to_use, primary_tokenizer)
|
||||
sent_messages: Final = messages if translated is None else translated.messages
|
||||
sent_tools: Final = tools if translated is None else translated.tools
|
||||
inference_model: Final = (
|
||||
deployment_inference_model(deployment) if is_mantle_claude_model(resolved_model) else None
|
||||
)
|
||||
claude_options: Final = (
|
||||
claude_count_options(
|
||||
counted_model,
|
||||
_claude_options_provider(route, use_mantle),
|
||||
sent_messages,
|
||||
sent_tools,
|
||||
_options_converse_sends(request_options, deployment_inference_model(deployment)),
|
||||
anthropic_beta,
|
||||
_options_converse_sends(request_options, inference_model),
|
||||
_beta_converse_sends(anthropic_beta, inference_model),
|
||||
deployment,
|
||||
)
|
||||
if is_mantle_claude_model(resolved_model)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import base64
|
|||
import functools
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -769,7 +769,7 @@ _MANTLE_COUNT_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5"
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_beta_headers_config", "aws_credentials_in_the_environment")
|
||||
async def test_runtime_count_without_a_deployment_drops_the_caller_betas_invoke_inference_drops() -> None:
|
||||
async def test_runtime_count_without_a_deployment_sends_the_caller_betas_invoke_inference_sends() -> None:
|
||||
sent: Final = await _invoke_inference_betas(_INTERLEAVED_THINKING, {})
|
||||
bedrock: Final = await _count_with_caller_beta("us.anthropic.claude-opus-5-5", _INTERLEAVED_THINKING, None)
|
||||
|
||||
|
|
@ -866,15 +866,19 @@ 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, ...]:
|
||||
def _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)
|
||||
if isinstance(content, str):
|
||||
return (content,)
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(text for block in content if isinstance(block, dict) and isinstance(text := block.get("text"), str))
|
||||
|
||||
|
||||
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)))
|
||||
return tuple(itertools.chain.from_iterable(map(_texts, messages)))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -901,6 +905,78 @@ async def test_mantle_count_for_a_converse_deployment_sends_the_format_and_betas
|
|||
assert _sent_betas(bedrock.requests[0]) == _sent_betas(inference.requests[0]) != ()
|
||||
|
||||
|
||||
_CONVERSE_DEPLOYMENT_MODEL: Final = "bedrock/converse/us.anthropic.claude-opus-5-5"
|
||||
_CALLER_BETAS_THE_ROWS_SPLIT: Final = ",".join(
|
||||
(
|
||||
_INTERLEAVED_THINKING,
|
||||
"token-efficient-tools-2025-02-19",
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
|
||||
ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value,
|
||||
"context-1m-2025-08-07",
|
||||
)
|
||||
)
|
||||
_BETAS: Final = TypeAdapter(frozenset[str])
|
||||
|
||||
|
||||
async def _converse_inference_betas(caller_beta: str) -> frozenset[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"],
|
||||
)
|
||||
assert inference.posted_hosts() == (RUNTIME_HOST,)
|
||||
return frozenset(_sent_betas(inference.requests[0]))
|
||||
|
||||
|
||||
async def _invoke_api_inference_betas(caller_beta: str) -> frozenset[str]:
|
||||
return _BETAS.validate_python(await _invoke_inference_betas(caller_beta, {}))
|
||||
|
||||
|
||||
def _invoke_api_counted_betas(bedrock: _Bedrock) -> frozenset[str]:
|
||||
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
|
||||
return _BETAS.validate_python(_betas_besides_effort(_invoke_body(bedrock.requests[0])))
|
||||
|
||||
|
||||
async def _mantle_api_inference_betas(caller_beta: str) -> frozenset[str]:
|
||||
return await _mantle_inference_betas(caller_beta, {"model": _MANTLE_COUNT_MODEL})
|
||||
|
||||
|
||||
def _mantle_api_counted_betas(bedrock: _Bedrock) -> frozenset[str]:
|
||||
assert bedrock.posted_hosts() == (MANTLE_HOST,)
|
||||
return frozenset(_sent_betas(bedrock.requests[0]))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("local_beta_headers_config")
|
||||
@pytest.mark.parametrize(
|
||||
("model_to_use", "count_api_inference_betas", "counted_betas"),
|
||||
(
|
||||
("us.anthropic.claude-opus-5-5", _invoke_api_inference_betas, _invoke_api_counted_betas),
|
||||
(_MANTLE_COUNT_MODEL, _mantle_api_inference_betas, _mantle_api_counted_betas),
|
||||
),
|
||||
ids=("runtime", "mantle"),
|
||||
)
|
||||
async def test_count_for_a_converse_deployment_sends_only_the_caller_betas_converse_inference_sends(
|
||||
model_to_use: str,
|
||||
count_api_inference_betas: Callable[[str], Awaitable[frozenset[str]]],
|
||||
counted_betas: Callable[[_Bedrock], frozenset[str]],
|
||||
) -> None:
|
||||
converse_sent: Final = await _converse_inference_betas(_CALLER_BETAS_THE_ROWS_SPLIT)
|
||||
bedrock: Final = await _count_with_caller_beta(
|
||||
model_to_use,
|
||||
_CALLER_BETAS_THE_ROWS_SPLIT,
|
||||
{"litellm_params": {**DEPLOYMENT["litellm_params"], "model": _CONVERSE_DEPLOYMENT_MODEL}},
|
||||
)
|
||||
|
||||
assert await count_api_inference_betas(_CALLER_BETAS_THE_ROWS_SPLIT) != converse_sent
|
||||
assert counted_betas(bedrock) == await count_api_inference_betas(",".join(sorted(converse_sent)))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_on_a_stand_in_model_counts_the_messages_and_output_config_the_deployment_model_sends() -> None:
|
||||
inference, bedrock = await _messages_inference_and_count(
|
||||
|
|
@ -1055,13 +1131,6 @@ async def test_runtime_count_sends_the_tools_invoke_inference_sends() -> None:
|
|||
assert _invoke_body(bedrock.requests[0])["tools"] == sent_tools
|
||||
|
||||
|
||||
def _texts(message: JsonValue) -> tuple[str, ...]:
|
||||
content: Final = message.get("content") if isinstance(message, dict) else None
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(text for block in content if isinstance(block, dict) and isinstance(text := block.get("text"), str))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("native_output", "request_options", "schema_inlined"),
|
||||
|
|
|
|||
|
|
@ -1961,7 +1961,7 @@ async def test_anthropic_count_tokens_does_not_count_again_after_a_rate_limit(
|
|||
|
||||
_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]] = {
|
||||
_BEDROCK_CLAUDE_DEPLOYMENT: Final[Mapping[str, str]] = {
|
||||
"model": "bedrock/us.anthropic.claude-count-probe-v1:0",
|
||||
"aws_access_key_id": "AKIATESTACCESSKEY",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue