mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(count_tokens)!: audit #47441 and keep the framed estimate for Bedrock count models from the cost map (#48805)
* fix(count_tokens): keep the framed estimate when a region-less token_count_model cannot count (audit #47441) * fix(count_tokens): let framing stand in for a shipped token_count_model, keep a configured one's error The framed local estimate now replaces a 400/404 from any count model the cost map ships, Mantle-pinned rows included, and an admin-set token_count_model keeps its error as the docs say. * fix(count_tokens): warn with the region when framing stands in for a pinned count model Test a Mantle row #47441 pinned in its own region, and keep the region hint check on an admin-set count model. * fix(count_tokens): name the deployment in the framing stand-in warning * fix(count_tokens): let framing stand in when a cost map count model's region denies the account A count model the cost map pins to a region can answer 401 or 403 when the account has not enabled that region or lacks Mantle count access there. With framing set, the estimate now stands in, as it does for 400 and 404. An admin set count model and a count model with no pinned region keep the error. The warning now names the model group, the count model and its region, and says the proxy falls back to the estimate. * test(count_tokens): cover framed 401 and 403 from a shipped count model region on the wire --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
rj-cairn
GitHub
parent
9b1d69e32c
commit
ff5bf4826d
5 changed files with 334 additions and 48 deletions
|
|
@ -13,6 +13,7 @@ from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields
|
|||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
|
||||
REGION_DENIED_STATUSES: Final = frozenset({401, 403})
|
||||
TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0
|
||||
TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0
|
||||
TOKENIZER_LOAD_MAX_RUNNING: Final = 4
|
||||
|
|
@ -47,6 +48,7 @@ class TokenCountModel:
|
|||
model: str
|
||||
region: str | None
|
||||
litellm_provider: TokenCountProvider | None = None
|
||||
configured: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -66,6 +68,10 @@ class ProviderCountFailure:
|
|||
def cannot_count_model(self) -> bool:
|
||||
return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES
|
||||
|
||||
@property
|
||||
def denied_in_region(self) -> bool:
|
||||
return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES
|
||||
|
||||
def with_hint(self, hint: str) -> "ProviderCountFailure":
|
||||
return replace(self, message=_with_hint(self.message, hint))
|
||||
|
||||
|
|
@ -111,7 +117,7 @@ def _count_provider(value: object) -> TokenCountProvider | None:
|
|||
return None
|
||||
|
||||
|
||||
def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
|
||||
def _count_model(value: Mapping[str, object], configured: bool) -> TokenCountModel | None:
|
||||
model: Final = _non_blank(value.get("model"))
|
||||
return (
|
||||
None
|
||||
|
|
@ -120,6 +126,7 @@ def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
|
|||
model=model,
|
||||
region=_non_blank(value.get("region")),
|
||||
litellm_provider=_count_provider(value.get("litellm_provider")),
|
||||
configured=configured,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -146,12 +153,12 @@ def _resolved_count_model(
|
|||
) -> TokenCountModel | None:
|
||||
configured_mapping: Final = as_str_mapping(configured)
|
||||
if configured_mapping is not None:
|
||||
return _count_model(configured_mapping)
|
||||
return _count_model(configured_mapping, configured=True)
|
||||
default_mapping: Final = as_str_mapping(default)
|
||||
return (
|
||||
None
|
||||
if default_mapping is None
|
||||
else _within_deployment_geography(_count_model(default_mapping), deployment_region)
|
||||
else _within_deployment_geography(_count_model(default_mapping, configured=False), deployment_region)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -14853,6 +14853,7 @@ async def _try_native_responses_count(
|
|||
async def _provider_token_count(
|
||||
count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]],
|
||||
model_to_use: str,
|
||||
request_model: str,
|
||||
settings: TokenCountSettings,
|
||||
local_estimate_allowed: bool,
|
||||
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
|
||||
|
|
@ -14875,17 +14876,28 @@ async def _provider_token_count(
|
|||
if count_model is None:
|
||||
return counted.model_copy(update={"count_source": "provider"})
|
||||
return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model})
|
||||
failure: Final = (
|
||||
counted
|
||||
if count_model is None or count_model.region is None
|
||||
else counted.with_hint(
|
||||
f"This model counts tokens with {count_model.model} in {count_model.region}. "
|
||||
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
|
||||
)
|
||||
pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model
|
||||
region_hint: Final = (
|
||||
None
|
||||
if pinned_count_model is None
|
||||
else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. "
|
||||
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
|
||||
)
|
||||
if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed:
|
||||
return failure
|
||||
raise failure.to_proxy_exception()
|
||||
failure: Final = counted if region_hint is None else counted.with_hint(region_hint)
|
||||
framing_can_stand_in: Final = (count_model is None or not count_model.configured) and (
|
||||
failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region)
|
||||
)
|
||||
if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed):
|
||||
raise failure.to_proxy_exception()
|
||||
if pinned_count_model is not None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Token count for %s answered %s from %s in %s, so it falls back to the framed local estimate",
|
||||
request_model,
|
||||
failure.status_code,
|
||||
pinned_count_model.model,
|
||||
pinned_count_model.region,
|
||||
)
|
||||
return failure
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -15455,6 +15467,7 @@ async def _count_request_tokens(
|
|||
chat_params=MappingProxyType({}) if chat_count is None else chat_count.params,
|
||||
),
|
||||
model_to_use=model_to_use,
|
||||
request_model=request.model,
|
||||
settings=settings,
|
||||
local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -24,7 +24,10 @@ _SONNET: Final = "us.anthropic.claude-sonnet-5-5"
|
|||
_SONNET_BASE: Final = "anthropic.claude-sonnet-5-5"
|
||||
_OPUS: Final = "us.anthropic.claude-opus-5-5"
|
||||
_OPUS_BASE: Final = "anthropic.claude-opus-5-5"
|
||||
_OPUS_4_7: Final = "us.anthropic.claude-opus-4-7"
|
||||
_NOVA: Final = "amazon.nova-lite-v1:0"
|
||||
_RUNTIME_ROUTED: Final = "us.anthropic.claude-sonnet-4-6"
|
||||
_CONFIGURED_COUNT_REGION: Final = "us-east-2"
|
||||
_KIMI: Final = "us.moonshotai.kimi-k3"
|
||||
_KIMI_BASE: Final = "moonshotai.kimi-k3"
|
||||
_PROFILE_ID: Final = "lit9563count"
|
||||
|
|
@ -168,6 +171,7 @@ def _mantle(request: Request) -> Reply:
|
|||
return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT})
|
||||
error_type: Final = {
|
||||
401: "authentication_error",
|
||||
403: "permission_error",
|
||||
404: "not_found_error",
|
||||
429: "rate_limit_error",
|
||||
}[status]
|
||||
|
|
@ -206,9 +210,9 @@ def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str,
|
|||
return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests)
|
||||
|
||||
|
||||
def _mantle_body(prompt: str) -> dict[str, JsonValue]:
|
||||
def _mantle_body(prompt: str, model: str = _OPUS_BASE) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": _OPUS_BASE,
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"system": _SYSTEM,
|
||||
"tools": _SDK_TOOLS,
|
||||
|
|
@ -219,13 +223,13 @@ def _url(port: int) -> str:
|
|||
return f"http://127.0.0.1:{port}"
|
||||
|
||||
|
||||
def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]:
|
||||
def _litellm_params(model: str, api_base: str, region: str = _REGION) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": f"bedrock/{model}",
|
||||
"api_base": api_base,
|
||||
"aws_access_key_id": _ACCESS_KEY,
|
||||
"aws_secret_access_key": _SECRET_KEY,
|
||||
"aws_region_name": _REGION,
|
||||
"aws_region_name": region,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -259,9 +263,13 @@ def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -
|
|||
|
||||
|
||||
def _deployment(
|
||||
scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None
|
||||
scenario: Scenario,
|
||||
runtime_port: int,
|
||||
model: str = _SONNET,
|
||||
model_info: Mapping[str, JsonValue] | None = None,
|
||||
region: str = _REGION,
|
||||
) -> str:
|
||||
return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port)))
|
||||
return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port), region))
|
||||
|
||||
|
||||
def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount:
|
||||
|
|
@ -382,6 +390,171 @@ def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_th
|
|||
assert _mantle_models(mantle) == ()
|
||||
|
||||
|
||||
def test_framing_is_added_to_the_local_estimate_when_a_shipped_runtime_count_model_cannot_count(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model")
|
||||
assert _shipped_token_count_model(_RUNTIME_ROUTED).get("region") is None
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
framed: Final = _count(
|
||||
gateway,
|
||||
_deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": _FRAMING}),
|
||||
400,
|
||||
count_model_status=404,
|
||||
)
|
||||
unframed: Final = _count(
|
||||
gateway,
|
||||
_deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": 0}),
|
||||
400,
|
||||
count_model_status=404,
|
||||
)
|
||||
assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed)
|
||||
assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed)
|
||||
assert _runtime_models(runtime) == (count_model, count_model)
|
||||
assert _mantle_models(mantle) == (count_model, count_model)
|
||||
|
||||
|
||||
_OPUS_4_7_PINNED: Final = (
|
||||
_OPUS_4_7,
|
||||
_REGION,
|
||||
_shipped_token_count_model_field(_OPUS_4_7, "region"),
|
||||
"anthropic.claude-opus-4-7",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "deployment_region", "count_region", "count_model", "count_model_status"),
|
||||
[
|
||||
(_SONNET, _REGION, _COUNT_REGION, _OPUS_BASE, 404),
|
||||
(_SONNET_BASE, "ca-central-1", "ca-central-1", _OPUS_BASE, 404),
|
||||
(*_OPUS_4_7_PINNED, 404),
|
||||
(*_OPUS_4_7_PINNED, 401),
|
||||
(*_OPUS_4_7_PINNED, 403),
|
||||
],
|
||||
ids=[
|
||||
"sonnet_5_5_in_its_pinned_region",
|
||||
"sonnet_5_5_outside_the_pinned_geography",
|
||||
"opus_4_7_in_its_pinned_region",
|
||||
"opus_4_7_pinned_region_denies_the_key",
|
||||
"opus_4_7_pinned_region_denies_the_account",
|
||||
],
|
||||
)
|
||||
def test_framing_is_added_to_the_local_estimate_when_a_shipped_mantle_count_model_cannot_count(
|
||||
counting_proxy: OwnedProxy,
|
||||
runtime_port: int,
|
||||
mantle_port: int,
|
||||
model: str,
|
||||
deployment_region: str,
|
||||
count_region: str,
|
||||
count_model: str,
|
||||
count_model_status: int,
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
assert _shipped_token_count_model_field(model, "litellm_provider") == "bedrock_mantle"
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
framed: Final = _count(
|
||||
gateway,
|
||||
_deployment(scenario, runtime_port, model, {"token_count_framing": _FRAMING}, deployment_region),
|
||||
200,
|
||||
count_model_status=count_model_status,
|
||||
)
|
||||
unframed: Final = _count(
|
||||
gateway,
|
||||
_deployment(scenario, runtime_port, model, {"token_count_framing": 0}, deployment_region),
|
||||
200,
|
||||
count_model_status=count_model_status,
|
||||
)
|
||||
assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed)
|
||||
assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed)
|
||||
assert _runtime_counts(runtime) == ()
|
||||
assert (
|
||||
_mantle_count_requests(mantle)
|
||||
== (((_ACCESS_KEY, count_region), _mantle_body(_prompt(200, count_model_status), count_model)),) * 2
|
||||
)
|
||||
|
||||
|
||||
def test_framing_keeps_the_403_when_the_shipped_count_model_runs_in_the_deployment_region(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
with pytest.raises(anthropic.PermissionDeniedError, match=re.escape(f"scripted mantle 403 for {_OPUS_BASE}")):
|
||||
_count(
|
||||
gateway,
|
||||
_deployment(scenario, runtime_port, _SONNET_BASE, {"token_count_framing": _FRAMING}, "ca-central-1"),
|
||||
200,
|
||||
count_model_status=403,
|
||||
)
|
||||
assert _runtime_counts(runtime) == ()
|
||||
assert _mantle_count_requests(mantle) == (
|
||||
((_ACCESS_KEY, "ca-central-1"), _mantle_body(_prompt(200, 403))),
|
||||
)
|
||||
|
||||
|
||||
def test_configured_token_count_model_error_stands_with_framing(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model")
|
||||
model_info: Final[dict[str, JsonValue]] = {
|
||||
"token_count_framing": _FRAMING,
|
||||
"token_count_model": {"model": count_model},
|
||||
}
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
with pytest.raises(anthropic.NotFoundError, match=re.escape(f"scripted mantle 404 for {count_model}")):
|
||||
_count(
|
||||
gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, model_info), 400, count_model_status=404
|
||||
)
|
||||
assert _runtime_models(runtime) == (count_model,)
|
||||
assert _mantle_models(mantle) == (count_model,)
|
||||
|
||||
|
||||
def test_configured_token_count_model_provider_counts_on_mantle_without_trying_the_runtime(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model")
|
||||
assert _shipped_token_count_model_field(_RUNTIME_ROUTED, "litellm_provider") == "bedrock"
|
||||
configured: Final[dict[str, JsonValue]] = {
|
||||
"model": count_model,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"region": _CONFIGURED_COUNT_REGION,
|
||||
}
|
||||
with (
|
||||
wire_server(_runtime, port=runtime_port) as runtime,
|
||||
wire_server(_mantle, port=mantle_port) as mantle,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
counted: Final = _count(
|
||||
gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_model": configured}), 400
|
||||
)
|
||||
assert (counted.input_tokens, counted.model_extra) == (
|
||||
_COUNT_MODEL_COUNT,
|
||||
{"litellm_count_source": "token_count_model", "litellm_count_model": count_model},
|
||||
), counted
|
||||
assert _runtime_counts(runtime) == ()
|
||||
assert _mantle_count_requests(mantle) == (
|
||||
((_ACCESS_KEY, _CONFIGURED_COUNT_REGION), {**_mantle_body(_prompt(400)), "model": count_model}),
|
||||
)
|
||||
|
||||
|
||||
def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region(
|
||||
counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -32,22 +32,31 @@ def _in_us_west_2() -> str:
|
|||
({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
|
||||
TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
|
||||
TokenCountSettings(
|
||||
framing=5,
|
||||
count_model=TokenCountModel(model="custom-model", region="ap-southeast-4", configured=True),
|
||||
),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model", "litellm_provider": "bedrock_mantle"}},
|
||||
TokenCountSettings(
|
||||
framing=97,
|
||||
count_model=TokenCountModel(model="custom-model", region=None, litellm_provider="bedrock_mantle"),
|
||||
count_model=TokenCountModel(
|
||||
model="custom-model", region=None, litellm_provider="bedrock_mantle", configured=True
|
||||
),
|
||||
),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model"}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
TokenCountSettings(
|
||||
framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True)
|
||||
),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model", "region": " "}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
TokenCountSettings(
|
||||
framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True)
|
||||
),
|
||||
),
|
||||
({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
|
|
@ -80,26 +89,26 @@ def test_config_overrides_the_cost_map_default_key_by_key(
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_region", "configured", "expected_region"),
|
||||
("deployment_region", "configured", "expected"),
|
||||
[
|
||||
("us-west-2", {}, "us-east-1"),
|
||||
("eu-central-1", {}, None),
|
||||
("ca-central-1", {}, None),
|
||||
("us-gov-west-1", {}, None),
|
||||
("us-west-2", {}, _OPUS_US_EAST_1),
|
||||
("eu-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)),
|
||||
("ca-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)),
|
||||
("us-gov-west-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)),
|
||||
(
|
||||
"eu-central-1",
|
||||
{"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}},
|
||||
"us-east-1",
|
||||
TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1", configured=True),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography(
|
||||
deployment_region: str, configured: dict[str, object], expected_region: str | None
|
||||
deployment_region: str, configured: dict[str, object], expected: TokenCountModel
|
||||
) -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
|
||||
assert settings.count_model == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/"))
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import threading
|
|||
import traceback
|
||||
import uuid
|
||||
from collections.abc import Callable, Coroutine, Mapping
|
||||
from dataclasses import replace
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
|
|
@ -209,7 +210,6 @@ async def test_gpt_token_counting():
|
|||
assert response.request_model == "gpt-4"
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_counter_leaves_the_routers_deployment_without_the_credentials_it_resolves(tmp_path, monkeypatch):
|
||||
"""token_counter reads the router's own deployment, so loading the named credential into it would leave the
|
||||
|
|
@ -1685,6 +1685,7 @@ async def _ladder(
|
|||
system=_COUNT_REQUEST.system,
|
||||
),
|
||||
model_to_use="us.anthropic.claude-sonnet-5-5",
|
||||
request_model=_COUNT_REQUEST.model,
|
||||
settings=settings,
|
||||
local_estimate_allowed=local_estimate_allowed,
|
||||
)
|
||||
|
|
@ -1748,6 +1749,7 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the
|
|||
system=_COUNT_REQUEST.system,
|
||||
),
|
||||
model_to_use="us.anthropic.claude-sonnet-5-5",
|
||||
request_model=_COUNT_REQUEST.model,
|
||||
settings=settings,
|
||||
local_estimate_allowed=False,
|
||||
)
|
||||
|
|
@ -1786,15 +1788,22 @@ async def test_token_count_model_counts_the_same_request_in_its_own_region_witho
|
|||
assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"count_model",
|
||||
[None, TokenCountModel(model="us.anthropic.claude-sonnet-5-5", region=None)],
|
||||
ids=["no_token_count_model", "shipped_token_count_model_in_the_deployment_region"],
|
||||
)
|
||||
@pytest.mark.parametrize("status_code", [401, 403, 429, 500])
|
||||
@pytest.mark.asyncio
|
||||
async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int):
|
||||
async def test_errors_other_than_cannot_count_raise_even_with_framing(
|
||||
count_model: TokenCountModel | None, status_code: int
|
||||
):
|
||||
counter: Final = _ScriptedCounter(
|
||||
{"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)}
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await _ladder(counter, TokenCountSettings(framing=10, count_model=None))
|
||||
await _ladder(counter, TokenCountSettings(framing=10, count_model=count_model))
|
||||
|
||||
assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet")
|
||||
|
||||
|
|
@ -1840,28 +1849,97 @@ async def test_uncountable_model_without_framing_raises():
|
|||
assert raised.value.code == "404"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count_model", "counted_as"),
|
||||
[
|
||||
(None, "us.anthropic.claude-sonnet-5-5"),
|
||||
(
|
||||
TokenCountModel(model="anthropic.claude-sonnet-5-5", region=None, litellm_provider="bedrock"),
|
||||
"bedrock/anthropic.claude-sonnet-5-5",
|
||||
),
|
||||
(_OPUS_IN_US_EAST_1, "us.anthropic.claude-opus-5-5"),
|
||||
],
|
||||
ids=[
|
||||
"no_token_count_model",
|
||||
"shipped_token_count_model_in_the_deployment_region",
|
||||
"shipped_token_count_model_pinned_to_a_region",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate():
|
||||
counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)})
|
||||
async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(
|
||||
count_model: TokenCountModel | None, counted_as: str
|
||||
):
|
||||
counter: Final = _ScriptedCounter({counted_as: _provider_reply("sonnet", status_code=400)})
|
||||
|
||||
result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None))
|
||||
result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=count_model))
|
||||
|
||||
assert isinstance(result, ProviderCountFailure)
|
||||
assert len(counter.calls) == 1
|
||||
assert [call[0] for call in counter.calls] == [counted_as]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status_code", [400, 404])
|
||||
@pytest.mark.parametrize("status_code", [401, 403])
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_count_model_error_stands_even_with_framing(status_code: int):
|
||||
async def test_framing_stands_in_when_a_shipped_token_count_model_region_denies_the_account(status_code: int):
|
||||
counter: Final = _ScriptedCounter(
|
||||
{"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)}
|
||||
)
|
||||
|
||||
result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=_OPUS_IN_US_EAST_1))
|
||||
|
||||
assert isinstance(result, ProviderCountFailure)
|
||||
assert result.status_code == status_code
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count_model", "region_hints"),
|
||||
[
|
||||
(replace(_OPUS_IN_US_EAST_1, configured=True), 1),
|
||||
(replace(_OPUS_IN_US_EAST_1, region=None, configured=True), 0),
|
||||
],
|
||||
ids=["pinned_to_a_region", "in_the_deployment_region"],
|
||||
)
|
||||
@pytest.mark.parametrize("status_code", [400, 401, 403, 404])
|
||||
@pytest.mark.asyncio
|
||||
async def test_configured_token_count_model_error_stands_even_with_framing(
|
||||
count_model: TokenCountModel, region_hints: int, status_code: int
|
||||
):
|
||||
counter: Final = _ScriptedCounter(
|
||||
{"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)}
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1))
|
||||
await _ladder(counter, TokenCountSettings(framing=78, count_model=count_model))
|
||||
|
||||
assert raised.value.code == str(status_code)
|
||||
assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message
|
||||
assert raised.value.message.startswith(f"scripted {status_code} for opus")
|
||||
assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count_model", "region_hints"),
|
||||
[(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())],
|
||||
ids=["pinned_to_a_region", "in_the_deployment_region"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_framing_standing_in_for_a_pinned_token_count_model_warns_with_its_model_and_region(
|
||||
count_model: TokenCountModel, region_hints: frozenset[bool], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)})
|
||||
verbose_proxy_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name):
|
||||
result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=count_model))
|
||||
finally:
|
||||
verbose_proxy_logger.removeHandler(caplog.handler)
|
||||
|
||||
assert isinstance(result, ProviderCountFailure)
|
||||
stand_in_warnings: Final = [
|
||||
r.getMessage() for r in caplog.records if "falls back to the framed local estimate" in r.getMessage()
|
||||
]
|
||||
assert {
|
||||
f"for {_COUNT_REQUEST.model} answered 404 from us.anthropic.claude-opus-5-5 in us-east-1" in message
|
||||
for message in stand_in_warnings
|
||||
} == region_hints
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2837,7 +2915,9 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp
|
|||
async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography(
|
||||
monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str
|
||||
):
|
||||
counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
|
||||
counter: Final = _ScriptedCounter(
|
||||
{"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.proxy_server,
|
||||
"llm_router",
|
||||
|
|
@ -2888,7 +2968,9 @@ class _VaultWaitingForLoopProgress:
|
|||
async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
|
||||
counter: Final = _ScriptedCounter(
|
||||
{"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}
|
||||
)
|
||||
|
||||
def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]:
|
||||
return counter, model, "bedrock"
|
||||
|
|
@ -3653,7 +3735,11 @@ async def test_responses_count_sends_openai_the_provider_file_id_inference_sends
|
|||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-files",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"},
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://count.test/v1",
|
||||
},
|
||||
"model_info": {"id": "deployment-a"},
|
||||
}
|
||||
]
|
||||
|
|
@ -3833,9 +3919,7 @@ async def test_responses_count_on_claude_rejects_forced_tool_choice_without_drop
|
|||
monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False})
|
||||
|
||||
with pytest.raises(ProxyException) as rejected:
|
||||
await _claude_count(
|
||||
monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"}
|
||||
)
|
||||
await _claude_count(monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"})
|
||||
|
||||
assert (rejected.value.code, "forced tool use" in rejected.value.message) == ("400", True)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue