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

fix(bedrock): refuse Mantle-only token counts on a pinned deployment with a clear 400

Browse source 
A token count model on bedrock_mantle used to fall back to the pinned bedrock-runtime endpoint with a model id only Mantle hosts, which failed every time. It now answers a 400 before any call, naming BEDROCK_MANTLE_API_BASE, and the proxy skips its region hint and framed local estimate for it.
This commit is contained in:
mateo-berri 2026-10-29 11:25:03 -07:00
parent 3520696cac
commit a279510e35
4 changed files with 114 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -19,6 +19,7 @@ from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.mantle_handler import (
NO_MANTLE_ENDPOINT_MESSAGE,
BedrockMantleCountTokensHandler,
deployment_mantle_endpoint,
)
@ -82,6 +83,19 @@ async def _has_mantle_endpoint(litellm_params: Mapping[str, object]) -> bool:
return await run_aws_signing(deployment_mantle_endpoint, litellm_params) is not None
def _no_mantle_endpoint_refusal(request_model: str, model_used: str) -> TokenCountResponse:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_used,
tokenizer_type=MANTLE_TOKENIZER_TYPE,
error=True,
error_message=NO_MANTLE_ENDPOINT_MESSAGE,
status_code=400,
rejected_before_provider_call=True,
)
def _request_data(
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
) -> Mapping[str, object]:
@ -159,8 +173,9 @@ class BedrockTokenCounter(BaseTokenCounter):
counted_model: Final = model_to_use.removeprefix("bedrock_mantle/")
resolved_model: Final = get_bedrock_base_model(counted_model)
litellm_params: Final = (deployment or {}).get("litellm_params", {})
counts_on_mantle: Final = use_mantle and await _has_mantle_endpoint(litellm_params)
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if counts_on_mantle else RUNTIME_TOKENIZER_TYPE
if use_mantle and not await _has_mantle_endpoint(litellm_params):
return _no_mantle_endpoint_refusal(request_model, model_to_use)
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE
translated: Final = (
await anthropic_count_fields(
counted_model,
@ -187,7 +202,7 @@ class BedrockTokenCounter(BaseTokenCounter):
await self._count_with(
self._mantle_handler, primary_tokenizer, mantle_request, litellm_params, resolved_model
)
if counts_on_mantle
if use_mantle
else await self._count_with(
self._runtime_handler, primary_tokenizer, runtime_request, litellm_params, resolved_model
)

5
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -19,8 +19,9 @@ from litellm.secret_managers.main import get_secret_str
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
NO_MANTLE_ENDPOINT_MESSAGE: Final = (
"This deployment pins its own bedrock-runtime endpoint, so its prompt is not counted on the public "
"bedrock-mantle host. Set BEDROCK_MANTLE_API_BASE to the bedrock-mantle endpoint it may use."
"This model counts tokens on bedrock-mantle, but the deployment pins its own bedrock-runtime endpoint "
"through api_base, aws_bedrock_runtime_endpoint, or AWS_BEDROCK_RUNTIME_ENDPOINT, so its prompt is not sent "
"to the public bedrock-mantle host. Set BEDROCK_MANTLE_API_BASE to the bedrock-mantle endpoint it may use."
)

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

@ -8,6 +8,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
from litellm.llms.bedrock.count_tokens.mantle_handler import NO_MANTLE_ENDPOINT_MESSAGE
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
@ -502,20 +503,23 @@ async def test_pinned_runtime_endpoint_retries_on_the_mantle_host_the_operator_n
@pytest.mark.asyncio
@pytest.mark.parametrize(
("runtime", "expected"),
("endpoint", "environment"),
(
((200, {"inputTokens": 193}), (False, 193, None)),
((400, UNSUPPORTED), (True, 0, 400)),
({"api_base": PRIVATE_RUNTIME}, {}),
({"aws_bedrock_runtime_endpoint": PRIVATE_RUNTIME}, {}),
({}, {"AWS_BEDROCK_RUNTIME_ENDPOINT": PRIVATE_RUNTIME}),
),
ids=("runtime_counts", "runtime_rejects"),
ids=("api_base", "aws_bedrock_runtime_endpoint", "AWS_BEDROCK_RUNTIME_ENDPOINT"),
)
async def test_mantle_counting_model_on_a_pinned_runtime_deployment_counts_on_that_endpoint_only(
monkeypatch: pytest.MonkeyPatch,
runtime: tuple[int, Mapping[str, object]],
expected: tuple[bool, int, int | None],
async def test_mantle_counting_model_on_a_pinned_runtime_deployment_is_refused_before_any_call(
monkeypatch: pytest.MonkeyPatch, endpoint: Mapping[str, str], environment: Mapping[str, str]
) -> None:
vault: Final = host_in_loop_probing_vault(monkeypatch, "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT")
bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 271}), runtime_host=PRIVATE_RUNTIME_HOST)
for name, value in environment.items():
monkeypatch.setenv(name, value)
vault: Final = host_in_loop_probing_vault(monkeypatch, "BEDROCK_MANTLE_API_BASE")
bedrock: Final = _Bedrock(
runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271}), runtime_host=PRIVATE_RUNTIME_HOST
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="bedrock_mantle/anthropic.claude-opus-5-5",
@ -525,16 +529,22 @@ async def test_mantle_counting_model_on_a_pinned_runtime_deployment_counts_on_th
"litellm_params": {
**DEPLOYMENT["litellm_params"],
"model": "bedrock/us.anthropic.claude-opus-5-5",
"api_base": PRIVATE_RUNTIME,
**endpoint,
}
},
request_model="claude-opus-5-5",
)
assert result is not None
assert (result.error, result.total_tokens, result.status_code) == expected
assert result.tokenizer_type == "bedrock_api"
assert bedrock.posted_urls() == (f"{PRIVATE_RUNTIME}/model/anthropic.claude-opus-5-5/count-tokens",)
assert (result.error, result.status_code, result.total_tokens, result.tokenizer_type) == (
True,
400,
0,
"bedrock_mantle_api",
)
assert result.error_message == NO_MANTLE_ENDPOINT_MESSAGE
assert result.rejected_before_provider_call
assert bedrock.requests == ()
assert vault.read.is_set()
assert not vault.read_on_loop.is_set()

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

@ -1762,6 +1762,75 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the
)
class _RejectingBedrockPeer:
def __init__(self) -> None:
self.requests: tuple[httpx.Request, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.requests = (*self.requests, request)
return httpx.Response(400, json={"message": "The provided model identifier is invalid."})
_MANTLE_COUNT_MODEL: Final = MappingProxyType(
{
"token_count_framing": 10,
"token_count_model": {
"model": "anthropic.claude-opus-5-5",
"litellm_provider": "bedrock_mantle",
"region": "us-east-1",
},
}
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("configured", "default"),
((_MANTLE_COUNT_MODEL, {}), ({}, _MANTLE_COUNT_MODEL)),
ids=("configured_token_count_model", "cost_map_token_count_model"),
)
async def test_mantle_token_count_model_on_a_pinned_runtime_deployment_answers_a_400_naming_the_fix(
configured: Mapping[str, object], default: Mapping[str, object], monkeypatch: pytest.MonkeyPatch
) -> None:
for name in ("AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT"):
monkeypatch.delenv(name, raising=False)
peer: Final = _RejectingBedrockPeer()
counter: Final = BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)))
with pytest.raises(ProxyException) as raised:
await _provider_token_count(
count_with=partial(
_try_provider_token_count,
counter,
"bedrock",
messages=_COUNT_REQUEST.messages,
contents=_COUNT_REQUEST.contents,
request_model=_COUNT_REQUEST.model,
deployment={
"litellm_params": {
"model": "bedrock/us.anthropic.claude-opus-5-5",
"api_base": "https://vpce-0a1b2c3d.bedrock-runtime.us-east-1.vpce.amazonaws.com",
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
}
},
tools=_COUNT_REQUEST.tools,
system=_COUNT_REQUEST.system,
),
model_to_use="us.anthropic.claude-opus-5-5",
request_model=_COUNT_REQUEST.model,
settings=resolve_token_count_settings(
configured=configured, default=default, deployment_region=lambda: "us-east-1"
),
local_estimate_allowed=True,
)
assert raised.value.code == "400"
assert "BEDROCK_MANTLE_API_BASE" in raised.value.message
assert "If that region is not enabled" not in raised.value.message
assert peer.requests == ()
@pytest.mark.parametrize(
("count_model", "region"),
[