From 43dcc2d80a808a5789015c0d174eaf012fbbed6f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:36:13 +0000 Subject: [PATCH 1/4] feat(bedrock): route token counts from model metadata (#47441) * feat(bedrock): route token counts from model metadata * fix(bedrock): keep counting metadata compatible with older proxies (cherry picked from commit 7b82b2a3c38cd3b641b581b08970aeaed3f03f0c) --- .../count_tokens/bedrock_token_counter.py | 18 +++-- .../common_utils/token_count_settings.py | 28 +++++++- litellm/proxy/proxy_server.py | 9 ++- litellm/types/token_count.py | 2 + model_prices_and_context_window.schema.json | 9 ++- .../test_bedrock_token_counter.py | 47 +++++++++++++ .../common_utils/test_token_count_settings.py | 35 +++++++--- tests/unit/proxy/test_proxy_token_counter.py | 69 +++++++++++++++++-- 8 files changed, 189 insertions(+), 28 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 97f28b9fbc5..068cf4f2e67 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -59,8 +59,7 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: class BedrockTokenCounter(BaseTokenCounter): - """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens - on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for""" + """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" def __init__( self, @@ -127,17 +126,22 @@ class BedrockTokenCounter(BaseTokenCounter): **({"tools": tools} if tools else {}), **({"system": system} if system else {}), } - resolved_model: Final = get_bedrock_base_model(model_to_use) + use_mantle: Final = model_to_use.startswith("bedrock_mantle/") + resolved_model: Final = get_bedrock_base_model(model_to_use.removeprefix("bedrock_mantle/")) - runtime: Final = await self._count_with( - self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + primary: Final = await self._count_with( + self._mantle_handler if use_mantle else self._runtime_handler, + MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE, + request_data, + litellm_params, + resolved_model, ) outcome: Final = ( await self._count_with( self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model ) - if _runtime_rejected_claude_model(runtime, resolved_model) - else runtime + if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model) + else primary ) match outcome: case CountedTokens(): diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 6f805616820..22b35bd1001 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -9,7 +9,7 @@ import litellm from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException -from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields, TokenCountTokenizer from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) @@ -46,6 +46,7 @@ def _with_hint(message: str, hint: str) -> str: class TokenCountModel: model: str region: str | None + litellm_provider: TokenCountProvider | None = None @dataclass(frozen=True, slots=True) @@ -100,9 +101,27 @@ def _non_blank(value: object) -> str | None: return value if isinstance(value, str) and value.strip() else None +def _count_provider(value: object) -> TokenCountProvider | None: + match value: + case "bedrock": + return "bedrock" + case "bedrock_mantle": + return "bedrock_mantle" + case _: + return None + + def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: model: Final = _non_blank(value.get("model")) - return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + return ( + None + if model is None + else TokenCountModel( + model=model, + region=_non_blank(value.get("region")), + litellm_provider=_count_provider(value.get("litellm_provider")), + ) + ) def _geography(region: str) -> str: @@ -112,7 +131,10 @@ def _geography(region: str) -> str: def _within_deployment_geography( count_model: TokenCountModel | None, deployment_region: Callable[[], str] ) -> TokenCountModel | None: - if count_model is None or count_model.region is None or count_model.model.startswith("global."): + if count_model is None or count_model.region is None: + return count_model + model: Final = count_model.model.removeprefix("bedrock/").removeprefix("bedrock_mantle/") + if model.startswith("global."): return count_model if _geography(deployment_region()) == _geography(count_model.region): return count_model diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cf291d65a5b..3fb7043aa2c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14534,9 +14534,14 @@ async def _provider_token_count( Returns None when the count should be estimated locally, or the provider's failure when the framed local estimate stands in for it and that failure is the answer if the estimate cannot read the request.""" count_model: Final = settings.count_model - counted: Final = await ( - count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + counting_model: Final = ( + model_to_use + if count_model is None + else count_model.model + if count_model.litellm_provider is None + else f"{count_model.litellm_provider}/{count_model.model.removeprefix('bedrock/').removeprefix('bedrock_mantle/')}" ) + counted: Final = await count_with(counting_model, None if count_model is None else count_model.region) if counted is None: return None if isinstance(counted, TokenCountResponse): diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py index 3b21401066c..90ccefea177 100644 --- a/litellm/types/token_count.py +++ b/litellm/types/token_count.py @@ -4,11 +4,13 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] +TokenCountProvider: TypeAlias = Literal["bedrock", "bedrock_mantle"] class TokenCountModelInfo(TypedDict): model: ReadOnly[str] region: ReadOnly[NotRequired[str]] + litellm_provider: ReadOnly[NotRequired[TokenCountProvider]] class TokenCountSourceFields(TypedDict, total=False): diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 45724d0d520..7a8b2c0cc7d 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1253,9 +1253,16 @@ "region": { "type": "string", "pattern": "\\S" + }, + "litellm_provider": { + "type": "string", + "enum": [ + "bedrock", + "bedrock_mantle" + ] } }, - "description": "Model that counts tokens for this one, and the AWS region its count API runs in." + "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in." }, "token_count_tokenizer": { "type": "string", diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 8a21027872a..dd0d6429f8a 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -174,3 +174,50 @@ async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> N "system": "You are a terse assistant.", "tools": tools, } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "host", "tokenizer", "expected_tokens"), + ( + ("bedrock_mantle/anthropic.claude-opus-5-5", MANTLE_HOST, "bedrock_mantle_api", 271), + ("bedrock/anthropic.claude-3-haiku-20240307-v1:0", RUNTIME_HOST, "bedrock_api", 193), + ), +) +async def test_qualified_counting_model_uses_its_endpoint_without_a_failed_probe( + model: str, host: str, tokenizer: str, expected_tokens: int +) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model, + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="served-model", + ) + + assert result is not None + assert result.total_tokens == expected_tokens + assert result.tokenizer_type == tokenizer + assert bedrock.posted_hosts() == (host,) + if host == MANTLE_HOST: + assert _JSON_BODY.validate_json(bedrock.requests[0].content)["model"] == model.split("/", 1)[1] + + +@pytest.mark.asyncio +async def test_explicit_mantle_failure_does_not_probe_runtime() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(404, UNSUPPORTED)) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + ) + + assert result is not None + assert result.error is True + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert bedrock.posted_hosts() == (MANTLE_HOST,) diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 8bd5ced1b58..d7314f1a6d3 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -4,6 +4,7 @@ from typing import Final import pytest +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy.common_utils.token_count_settings import ( ProviderCountFailure, TokenCountModel, @@ -33,6 +34,13 @@ def _in_us_west_2() -> str: {"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")), ), + ( + {"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"), + ), + ), ( {"token_count_model": {"model": "custom-model"}}, TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), @@ -94,13 +102,15 @@ def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geogr assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) -def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: +@pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/")) +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography(prefix: str) -> None: + model: Final = f"{prefix}global.anthropic.claude-opus-5-5" settings: Final = resolve_token_count_settings( configured={}, - default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + default={"token_count_model": {"model": model, "region": "us-east-1"}}, deployment_region=lambda: "eu-west-1", ) - assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + assert settings.count_model == TokenCountModel(model=model, region="us-east-1") def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: @@ -157,21 +167,24 @@ def test_count_source_fields_name_the_counting_model_only_when_token_count_model assert count_source_fields(_counted(source)) == expected -def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None: +def test_every_shipped_token_count_model_resolves_to_a_mapped_aws_model() -> None: cost_map: Final = json.loads(_COST_MAP.read_text()) + aws_models: Final = frozenset( + get_bedrock_base_model(name) + for name, entry in cost_map.items() + if isinstance(entry, dict) and entry.get("litellm_provider") in ("bedrock", "bedrock_converse") + ) count_models: Final = { name: entry["token_count_model"] for name, entry in cost_map.items() if isinstance(entry, dict) and "token_count_model" in entry } assert count_models - assert { - name: count_model - for name, count_model in count_models.items() - if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"] - or count_model["model"] == name - or not count_model.get("region") - } == {} + for name, count_model in count_models.items(): + model: Final = count_model["model"] + assert get_bedrock_base_model(model) in aws_models + assert cost_map[name]["litellm_provider"] in ("bedrock", "bedrock_converse") + assert "region" not in count_model or count_model["region"].strip() _UPSTREAM_404: Final = { diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 720739a730b..d17fad7e9b3 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -65,6 +65,7 @@ from litellm.proxy.common_utils.token_count_settings import ( TokenCountModel, TokenCountSettings, cost_map_entry, + resolve_token_count_settings, ) from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request from litellm.proxy.proxy_server import ( @@ -1368,6 +1369,65 @@ async def test_model_without_token_count_model_counts_itself_in_the_deployment_r assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider", "expected_host", "expected_tokens"), + ( + ("bedrock", "bedrock-runtime.us-east-1.amazonaws.com", 193), + ("bedrock_mantle", "bedrock-mantle.us-east-1.api.aws", 271), + ), +) +async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the_public_model( + provider: str, expected_host: str, expected_tokens: int, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + model: Final = "anthropic.claude-opus-5-5" + + def peer(request: httpx.Request) -> httpx.Response: + assert request.url.host == expected_host + if provider == "bedrock_mantle": + assert json.loads(request.content)["model"] == model + return httpx.Response(200, json={"input_tokens": expected_tokens}) + assert request.url.path == f"/model/{model}/count-tokens" + return httpx.Response(200, json={"inputTokens": expected_tokens}) + + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": model, "litellm_provider": provider, "region": "us-east-1"}}, + deployment_region=lambda: "us-west-2", + ) + counter: Final = BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) + result: Final = 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": { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + } + }, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=False, + ) + + assert isinstance(result, TokenCountResponse) + assert (result.total_tokens, result.model_used, result.count_source) == ( + expected_tokens, + model, + "token_count_model", + ) + + @pytest.mark.parametrize( ("count_model", "region"), [ @@ -2390,7 +2450,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is @pytest.mark.parametrize( ("model_info", "counted_with"), [ - ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ({}, ("bedrock_mantle/us.anthropic.claude-opus-5-5", "us-east-1", 447)), ( {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), @@ -2406,6 +2466,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp { "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + "bedrock_mantle/us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), } ) monkeypatch.setattr( @@ -2444,7 +2505,7 @@ 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({"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", @@ -2468,7 +2529,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True ) - assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + assert (counter.calls[0][0], counter.regions[0]) == ("bedrock_mantle/anthropic.claude-opus-5-5", counted_in) class _SlowVault: @@ -2488,7 +2549,7 @@ async def _longest_stall(done: asyncio.Event) -> float: 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({"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" From c8218796c7f991d114c28d1d3be0bc92ed741ce1 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:42:43 +0000 Subject: [PATCH 2/4] test(integration): follow #47441 metadata-routed counting in count_tokens wire tests (#47525) * test(integration): follow #47441 metadata-routed Mantle counting in count_tokens wire tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): build the runtime count body without mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit f5d6e89ad8341e4df5400a0bd0e99509b8ae62c8) --- .../test_bedrock_mantle_count_tokens_wire.py | 205 ++++++++++++------ .../providers/test_count_tokens_model_wire.py | 81 ++++++- 2 files changed, 207 insertions(+), 79 deletions(-) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index ab0c35d0ca4..4fe29fa49d8 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -70,6 +70,18 @@ _SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[JsonValue]) +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + + +def _count_model_source(model: str) -> dict[str, JsonValue]: + cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + token_count_model: Final = cost_map[model].get("token_count_model") + assert isinstance(token_count_model, dict), model + count_model: Final = token_count_model.get("model") + assert isinstance(count_model, str), token_count_model + return {"litellm_count_source": "token_count_model", "litellm_count_model": count_model} + _SDK_MESSAGES: Final[list[MessageParam]] = [{"role": "user", "content": "Count this message"}] _MESSAGES: Final = _JSON_LIST.validate_json(json.dumps(_SDK_MESSAGES)) @@ -94,6 +106,32 @@ _GEMINI_BODY: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "part _GEMINI_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this"}] +def _runtime_count_body(prompt: str, include_system_and_tools: bool = False) -> dict[str, JsonValue]: + tool: Final = _SDK_TOOLS[0] + converse: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": [{"text": prompt}]}], + **( + { + "system": [{"text": _SYSTEM}], + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": tool["name"], + "description": tool["description"], + "inputSchema": {"json": tool["input_schema"]}, + } + } + ] + }, + } + if include_system_and_tools + else {} + ), + } + return {"input": {"converse": converse}} + + def _mantle_body(**fields: JsonValue) -> dict[str, JsonValue]: return {"model": _OPUS_BASE, "messages": _MESSAGES, **fields} @@ -223,7 +261,11 @@ def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue { **config, "model_list": [ - {"model_name": name, "litellm_params": _litellm_params(model, runtime_url)} + { + "model_name": name, + "litellm_params": _litellm_params(model, runtime_url), + "model_info": {"token_count_model": {}}, + } for name, model in _OWNED_MODELS.items() ], "litellm_settings": {**litellm_settings, **settings}, @@ -233,7 +275,11 @@ def _owned_config(path: Path, runtime_url: str, settings: Mapping[str, JsonValue return path -def _deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: +def _fallback_deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: + return scenario.model(model_info={"token_count_model": {}}, api_key=None, **_litellm_params(model, api_base)) + + +def _metadata_deployment(scenario: Scenario, api_base: str, model: str = _OPUS) -> str: return scenario.model(model_info=None, api_key=None, **_litellm_params(model, api_base)) @@ -290,6 +336,14 @@ def _runtime_count_targets(wire: Wire) -> tuple[str, ...]: return tuple(target_of(request) for request in counts) +def _runtime_count_requests(wire: Wire) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + counts: Final = tuple(request for request in wire.drain() if _COUNT_TARGET.match(target_of(request))) + for request in counts: + assert request.method == "POST", request.method + _assert_sigv4(request) + return tuple((target_of(request), _JSON_OBJECT.validate_json(request.body)) for request in counts) + + def _clients(stack: ExitStack, base_url: str, count: int, timeout: float = 30) -> tuple[httpx.Client, ...]: return tuple( stack.enter_context(httpx.Client(base_url=base_url, timeout=timeout, trust_env=False)) for _ in range(count) @@ -376,7 +430,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) @@ -392,7 +446,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -408,7 +462,7 @@ def test_utils_token_counter_call_endpoint_reports_the_mantle_tokenizer( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) @@ -427,14 +481,14 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _metadata_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, "promptTokensDetails": [], - "litellm_count_source": "provider", + **_count_model_source(_OPUS), }, response.text - assert len(_runtime_count_targets(runtime)) == 1 + assert _runtime_count_requests(runtime) == () assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) @@ -447,14 +501,16 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _SONNET) + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) assert _payload(response) == { "totalTokens": _RUNTIME_COUNT, "promptTokensDetails": [], - "litellm_count_source": "provider", + **_count_model_source(_SONNET), }, response.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) + assert _runtime_count_requests(runtime) == ( + (f"/model/{_SONNET}/count-tokens", _runtime_count_body("Count this")), + ) assert mantle.drain() == () @@ -466,7 +522,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -481,7 +537,7 @@ def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: Owned wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) assert response.status_code == 404, response.text assert _REJECTION in response.text @@ -496,7 +552,7 @@ def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, ma wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _NOVA) + model: Final = _fallback_deployment(scenario, runtime.url, _NOVA) body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]} response: Final = _count(gateway, body) assert response.status_code == 404, response.text @@ -504,9 +560,7 @@ def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, ma assert mantle.drain() == () -def test_google_count_tokens_forwards_system_and_tools_to_mantle( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: +def test_google_count_tokens_forwards_system_and_tools_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway system_text: Final = "Use the provided tool" tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}} @@ -526,7 +580,7 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, @@ -575,7 +629,7 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) assert _payload(response) == { "totalTokens": _MANTLE_COUNT, @@ -607,7 +661,7 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) @@ -627,7 +681,7 @@ def test_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: OwnedP wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = _anthropic_client(gateway).messages.count_tokens( model=model, messages=_SDK_MESSAGES, system=_SYSTEM, tools=_SDK_TOOLS ) @@ -643,7 +697,7 @@ def test_async_anthropic_sdk_count_tokens_counts_through_mantle(counting_proxy: wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_SDK_MESSAGES, system=_SYSTEM) ) @@ -659,17 +713,36 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _SONNET) + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) assert _payload(_count(gateway, _full(model))) == { "input_tokens": _RUNTIME_COUNT, - "litellm_count_source": "provider", + **_count_model_source(_SONNET), } detailed: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) payload: Final = _payload(detailed) - assert (payload["total_tokens"], payload["tokenizer_type"]) == (_RUNTIME_COUNT, "bedrock_api"), detailed.text - assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) * 2 + assert payload == { + "total_tokens": _RUNTIME_COUNT, + "request_model": model, + "model_used": _count_model_source(_SONNET)["litellm_count_model"], + "tokenizer_type": "bedrock_api", + "count_source": "token_count_model", + "original_response": {"input_tokens": _RUNTIME_COUNT}, + "error": False, + "error_message": None, + "status_code": None, + }, detailed.text + assert _runtime_count_requests(runtime) == ( + ( + f"/model/{_SONNET}/count-tokens", + _runtime_count_body("Count this message", include_system_and_tools=True), + ), + ( + f"/model/{_SONNET}/count-tokens", + _runtime_count_body("Count this message", include_system_and_tools=True), + ), + ) assert mantle.drain() == () @@ -682,7 +755,7 @@ def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_ma wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url, _NOVA) + model: Final = _fallback_deployment(scenario, runtime.url, _NOVA) response: Final = _count(gateway, _bare(model)) assert response.status_code == 400, response.text assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text @@ -697,7 +770,7 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = { "model": model, "messages": [{"role": "user", "content": "status=403 Count this message"}], @@ -712,7 +785,7 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario: - model: Final = _deployment(scenario, _closed_port_url()) + model: Final = _fallback_deployment(scenario, _closed_port_url()) response: Final = _count(gateway, _full(model)) assert 500 <= response.status_code < 600, response.text assert "input_tokens" not in response.text, response.text @@ -728,7 +801,7 @@ def test_bedrock_passthrough_count_tokens_still_answers_the_runtime_rejection( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", "/bedrock/v1/messages/count_tokens", _full(model)) assert response.status_code == 400, response.text assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text @@ -745,7 +818,7 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", @@ -770,7 +843,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, count_request(model, case)) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -786,7 +859,7 @@ def test_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = _anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) assert counted.input_tokens == _MANTLE_COUNT, counted assert len(_runtime_count_targets(runtime)) == 1 @@ -802,7 +875,7 @@ def test_async_anthropic_sdk_count_tokens_lifts_the_leading_system_through_mantl wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( _async_anthropic_client(gateway).messages.count_tokens(model=model, messages=_LIFT_SDK_MESSAGES) ) @@ -820,7 +893,7 @@ def test_utils_token_counter_call_endpoint_counts_a_leading_system_through_mantl wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/utils/token_counter", count_request(model, STRING_CASE), params={"call_endpoint": "true"} ) @@ -841,7 +914,7 @@ def test_openai_sdk_input_tokens_lifts_instructions_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = openai_client(gateway).responses.input_tokens.count( model=model, input=USER_TEXT, instructions=INSTRUCTION ) @@ -859,7 +932,7 @@ def test_async_openai_sdk_input_tokens_lifts_instructions_through_mantle( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) counted: Final = asyncio.run( async_openai_client(gateway).responses.input_tokens.count( model=model, input=USER_TEXT, instructions=INSTRUCTION @@ -879,7 +952,7 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/responses/input_tokens", @@ -903,7 +976,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -919,7 +992,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = { "model": model, "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], @@ -941,7 +1014,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_s wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]} response: Final = _count(gateway, body) assert response.status_code == 400, response.text @@ -959,7 +1032,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_syste wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5}) assert response.status_code == 400, response.text assert "input_tokens" not in response.text, response.text @@ -976,7 +1049,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 @@ -992,7 +1065,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]} response: Final = _count(gateway, body) assert 500 <= response.status_code < 600, response.text @@ -1015,7 +1088,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) first: Final = json.dumps([USER]) last: Final = json.dumps([LEADING, USER]) response: Final = gateway.client.post( @@ -1064,7 +1137,7 @@ def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers( pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) runtime: Final = stack.enter_context(wire_server(_runtime)) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = count_request(model, STRING_CASE) def count(client: httpx.Client) -> tuple[int, JsonValue]: @@ -1092,7 +1165,7 @@ def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors( wire_server(_mantle(_rejecting(status)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) assert response.status_code == status, response.text assert "input_tokens" not in response.text, response.text @@ -1114,7 +1187,7 @@ def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( wire_server(_mantle(lambda _request: Reply(body=body)), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) assert 500 <= response.status_code < 600, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1124,16 +1197,14 @@ def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens( def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_runtime) as runtime, gateway.scenario() as scenario: - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, _full(model)) assert 500 <= response.status_code < 600, response.text assert "input_tokens" not in response.text, response.text assert len(_runtime_count_targets(runtime)) == 1 -def test_messages_count_tokens_errors_when_mantle_never_answers( - counting_proxy: OwnedProxy, mantle_port: int -) -> None: +def test_messages_count_tokens_errors_when_mantle_never_answers(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway held: Final[SimpleQueue[str]] = SimpleQueue() release: Final = threading.Event() @@ -1150,7 +1221,7 @@ def test_messages_count_tokens_errors_when_mantle_never_answers( mantle: Final = stack.enter_context(wire_server(_mantle(hold), port=mantle_port)) stack.callback(release.set) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = _full(model) future: Final = pool.submit(_counted_on, client, gateway.key, body) eventually(held.qsize, lambda size: size == 1, seconds=30) @@ -1172,7 +1243,7 @@ def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system( wire_server(_mantle(_strict), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": 5}) assert response.status_code == 400, response.text assert "input_tokens" not in response.text, response.text @@ -1194,7 +1265,7 @@ def test_messages_count_tokens_rejects_malformed_tools_without_calling_either_pe wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) refused: Final = _count(gateway, {**_bare(model), "tools": tools}) assert refused.status_code == 400, refused.text assert object_value(_JSON_OBJECT.validate_json(refused.content)["error"])["type"] == "invalid_request_error" @@ -1213,7 +1284,7 @@ def test_messages_count_tokens_without_messages_is_rejected( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {"model": model, "system": _SYSTEM, "tools": _TOOLS, **fields}) assert response.status_code == 400, response.text assert "messages parameter is required" in response.text, response.text @@ -1234,7 +1305,7 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_either_peer wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/messages/count_tokens", {**_full(model), **fields}, key="sk-not-a-key-this-proxy-issued" ) @@ -1253,7 +1324,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1269,7 +1340,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) fields: Final = f'"system": {json.dumps(_SYSTEM)}, "tools": {json.dumps(_TOOLS)}' response: Final = gateway.client.post( "/v1/messages/count_tokens", @@ -1292,7 +1363,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 @@ -1308,7 +1379,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 @@ -1364,7 +1435,7 @@ def test_mantle_outage_between_concurrent_waves_errors_then_recovers( pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) runtime: Final = stack.enter_context(wire_server(_runtime)) scenario: Final = stack.enter_context(gateway.scenario()) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) body: Final = _full(model) def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]: @@ -1398,7 +1469,7 @@ def test_slow_mantle_holds_concurrent_counts_without_stalling_the_proxy( scenario: Final = stack.enter_context(gateway.scenario()) pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) stack.callback(release.set) - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) futures: Final = tuple( pool.submit(_generated_then_counted, client, gateway.key, model, _full(model)) for client in clients ) @@ -1495,7 +1566,7 @@ def test_chat_completions_on_the_same_deployment_still_generate( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/chat/completions", @@ -1523,7 +1594,7 @@ def test_messages_endpoint_on_the_same_deployment_still_generates(counting_proxy wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request( "POST", "/v1/messages", @@ -1546,7 +1617,7 @@ def test_responses_endpoint_on_the_same_deployment_still_generates( wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": f"marker-{marker}"}) assert response.status_code == 200, response.text assert answer(marker) in response.text, response.text @@ -1562,6 +1633,6 @@ def test_utils_token_counter_local_mode_never_calls_either_peer(counting_proxy: wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario, ): - model: Final = _deployment(scenario, runtime.url) + model: Final = _fallback_deployment(scenario, runtime.url) assert _local_count(gateway, _full(model)) > 0 assert runtime.drain() == () and mantle.drain() == () diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py index 46e5396ef64..b6ace7ddef3 100644 --- a/tests/integration/providers/test_count_tokens_model_wire.py +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -31,7 +31,6 @@ _PROFILE_ID: Final = "lit9563count" _PROFILE_ARN: Final = f"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/{_PROFILE_ID}" _COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" _REGION: Final = "us-west-2" -_COUNT_REGION: Final = "us-east-1" _ACCESS_KEY: Final = "AKIAINTEGRATIONCOUNT" _SECRET_KEY: Final = "integration-count-secret" _PROVIDER_COUNT: Final = 1345 @@ -39,12 +38,37 @@ _COUNT_MODEL_COUNT: Final = 2718 _FRAMING: Final = 50 _UNSUPPORTED: Final = "The provided model doesn't support counting tokens." _COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") +_MANTLE_COUNT_TARGET: Final = "/anthropic/v1/messages/count_tokens" +_MANTLE_VERSION: Final = "2023-06-01" _SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") _SCRIPTED_COUNT_MODEL_STATUS: Final = re.compile(r"count_model=(\d{3})") _CREDENTIAL: Final = re.compile(r"Credential=([^/]+)/\d+/([^/]+)/") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + +def _shipped_token_count_model(model: str) -> dict[str, JsonValue]: + token_count_model: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model].get("token_count_model") + assert isinstance(token_count_model, dict), model + return token_count_model + + +def _shipped_token_count_model_field(model: str, field: str) -> str: + value: Final = _shipped_token_count_model(model).get(field) + assert isinstance(value, str), (model, field, value) + return value + + +def _count_model_source(model: str) -> dict[str, JsonValue]: + return { + "litellm_count_source": "token_count_model", + "litellm_count_model": _shipped_token_count_model_field(model, "model"), + } + + +_COUNT_REGION: Final = _shipped_token_count_model_field(_SONNET, "region") +_COUNT_MODEL_SOURCE: Final = _count_model_source(_SONNET) + _SYSTEM: Final = "You are a terse assistant that answers in one sentence" _WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", @@ -62,7 +86,6 @@ _SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ }, } ] -_COUNT_MODEL_SOURCE: Final = {"litellm_count_source": "token_count_model", "litellm_count_model": _OPUS} _OPENAI_PRIMARY: Final = "gpt-count-primary" _OPENAI_COUNT_MODEL: Final = "gpt-count-model" _OPENAI_COUNT_TARGET: Final = "/v1/responses/input_tokens" @@ -140,8 +163,17 @@ def _runtime(request: Request) -> Reply: def _mantle(request: Request) -> Reply: model: Final = _JSON_OBJECT.validate_json(request.body).get("model") + status: Final = _scripted(_SCRIPTED_COUNT_MODEL_STATUS, request, 200) + if status == 200: + return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT}) + error_type: Final = { + 401: "authentication_error", + 404: "not_found_error", + 429: "rate_limit_error", + }[status] return _json_reply( - 404, {"type": "error", "error": {"type": "not_found_error", "message": f"scripted mantle 404 for {model}"}} + status, + {"type": "error", "error": {"type": error_type, "message": f"scripted mantle {status} for {model}"}}, ) @@ -165,6 +197,24 @@ def _mantle_models(wire: Wire) -> tuple[JsonValue, ...]: return tuple(_JSON_OBJECT.validate_json(request.body).get("model") for request in wire.drain()) +def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str, JsonValue]], ...]: + requests: Final = tuple(wire.drain()) + for request in requests: + assert (request.method, request.target) == ("POST", _MANTLE_COUNT_TARGET), request.target + assert request.headers["anthropic-version"] == _MANTLE_VERSION, request.headers + assert request.headers["content-type"] == "application/json", request.headers + return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests) + + +def _mantle_body(prompt: str) -> dict[str, JsonValue]: + return { + "model": _OPUS_BASE, + "messages": [{"role": "user", "content": prompt}], + "system": _SYSTEM, + "tools": _SDK_TOOLS, + } + + def _url(port: int) -> str: return f"http://127.0.0.1:{port}" @@ -255,15 +305,15 @@ def test_shipped_token_count_model_counts_the_request_in_its_region_without_tryi ): counted: Final = _count(gateway, _deployment(scenario, runtime_port), 404) assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted - assert _runtime_counts(runtime) == ((_OPUS_BASE, (_ACCESS_KEY, _COUNT_REGION), _converse_body(_prompt(404))),) - assert _mantle_models(mantle) == () + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == (((_ACCESS_KEY, _COUNT_REGION), _mantle_body(_prompt(404))),) @pytest.mark.parametrize( ("count_model_status", "error", "message"), [ - (401, anthropic.AuthenticationError, f"scripted runtime 401 for {_OPUS_BASE}"), - (429, anthropic.RateLimitError, f"scripted runtime 429 for {_OPUS_BASE}"), + (401, anthropic.AuthenticationError, f"scripted mantle 401 for {_OPUS_BASE}"), + (429, anthropic.RateLimitError, f"scripted mantle 429 for {_OPUS_BASE}"), (404, anthropic.NotFoundError, f"scripted mantle 404 for {_OPUS_BASE}"), ], ) @@ -283,9 +333,16 @@ def test_token_count_model_failure_answers_its_error_and_names_the_region_to_ena ): with pytest.raises(error) as refused: _count(gateway, _deployment(scenario, runtime_port), 200, count_model_status=count_model_status) - assert message in str(refused.value) and f"{_OPUS} in {_COUNT_REGION}" in str(refused.value), refused.value - assert _runtime_models(runtime) == (_OPUS_BASE,) - assert _mantle_models(mantle) == ((_OPUS_BASE,) if count_model_status == 404 else ()) + assert message in str(refused.value) and ( + f"{_COUNT_MODEL_SOURCE['litellm_count_model']} in {_COUNT_REGION}" in str(refused.value) + ), refused.value + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ( + (_ACCESS_KEY, _COUNT_REGION), + _mantle_body(_prompt(200, count_model_status)), + ), + ) def test_framing_is_added_to_the_local_estimate_when_the_runtime_cannot_count( @@ -359,8 +416,8 @@ def test_inference_profile_arn_counts_with_the_token_count_model_of_its_base_mod ) counted: Final = _count(gateway, model, 404) assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted - assert _runtime_models(runtime) == (_OPUS_BASE,) - assert _mantle_models(mantle) == () + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == (((_ACCESS_KEY, _COUNT_REGION), _mantle_body(_prompt(404))),) @pytest.mark.parametrize( From 38482f7cda185e2b532c32d80b43029680411242 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 31 Oct 2026 15:47:06 -0700 Subject: [PATCH 3/4] test(integration): backport the cursor responses wire test fix from #47600 to rc/1.108.0 rc/1.108.0 carries the #43935 backport, which enforces gpt-5's supports_none_reasoning_effort and supports_xhigh_reasoning_effort flags, so the test now reads the levels it parametrizes from the cost map like main does. Only this file of #47600 is taken. Its other files follow changes rc does not have (cherry picked from commit e78f7b4602, tests/integration/providers/test_cursor_responses_wire.py only) --- .../providers/test_cursor_responses_wire.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/integration/providers/test_cursor_responses_wire.py b/tests/integration/providers/test_cursor_responses_wire.py index 6e8c877f7ca..d77d35f2ff8 100644 --- a/tests/integration/providers/test_cursor_responses_wire.py +++ b/tests/integration/providers/test_cursor_responses_wire.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import os import uuid +from importlib.resources import files from typing import Final import pytest @@ -16,8 +17,16 @@ _MODEL: Final = "gpt-5" _API_KEY: Final = "synthetic-cursor-key" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) -_THINKING_LEVELS: Final = ("none", "minimal", "low", "medium", "high", "xhigh", "max") -_CHAT_THINKING_LEVELS: Final = ("none", "minimal", "low", "medium", "high", "max") +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) +_MODEL_ROW: Final = _COST_MAP.validate_json( + (files("litellm") / "model_prices_and_context_window_backup.json").read_bytes() +)[_MODEL] +_THINKING_LEVELS: Final = tuple( + level + for level in ("none", "minimal", "low", "medium", "high", "xhigh", "max") + if _MODEL_ROW.get(f"supports_{level}_reasoning_effort") is not False +) +_CHAT_THINKING_LEVELS: Final = tuple(level for level in _THINKING_LEVELS if level != "xhigh") _RESPONSES_SUFFIXES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = tuple( (f"-thinking-{level}", {"reasoning": {"effort": level}}) for level in _THINKING_LEVELS ) + ( From 187d933de24aa65fbab0c2956731ae743e6d246f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 31 Oct 2026 15:47:07 -0700 Subject: [PATCH 4/4] test(integration): backport the stale-test fixes from #48665 that rc/1.108.0 needs rc/1.108.0 carries OpenTelemetry 1.45.1 (#48055), #47611 and the #47639 base_model strip, so the OTel traceresponse header, langtrace outage, guardrail concurrent-disconnect and Bedrock lookaround tests expected the old behavior. This takes #48665's hunks for those four files only, without the unhandled to non_object_body rename, which depends on #48438 that rc does not have (cherry picked from commit 21ab2cb87b, four test files only) --- .../observability/test_guardrail_effects.py | 26 +++++++++++++------ .../observability/test_langtrace_delivery.py | 7 ----- .../test_otel_trace_response_header.py | 12 ++++++--- ..._bedrock_converse_lookaround_regex_wire.py | 6 ++--- 4 files changed, 30 insertions(+), 21 deletions(-) diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index f267cff04f3..eb786a4f89b 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -2571,9 +2571,15 @@ def _scanned_while_upstream_is_held( rig.gate.set() -def _post_call_statuses(rig: _DisconnectRig, model: str, rows: int = 1) -> tuple[tuple[str, ...], ...]: +def _post_call_rows( + rig: _DisconnectRig, model: str, rows: int = 1, call_types: frozenset[str] | None = None +) -> tuple[tuple[str, tuple[str, ...]], ...]: found: Final = eventually( - lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda: tuple( + row + for row in read_rows('SELECT call_type, metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)) + if call_types is None or row["call_type"] in call_types + ), lambda values: len(values) == rows, seconds=70, ) @@ -2588,7 +2594,11 @@ def _post_call_statuses(rig: _DisconnectRig, model: str, rows: int = 1) -> tuple ) return tuple(str(entry["guardrail_status"]) for entry in post_call) - return tuple(statuses(row["metadata"]) for row in found) + return tuple((str(row["call_type"]), statuses(row["metadata"])) for row in found) + + +def _post_call_statuses(rig: _DisconnectRig, model: str, rows: int = 1) -> tuple[tuple[str, ...], ...]: + return tuple(statuses for _, statuses in _post_call_rows(rig, model, rows)) _ENDPOINT_CLIENTS: Final = ( @@ -2621,7 +2631,7 @@ def test_client_disconnect_mid_stream_still_scans_the_content_it_already_receive ( pytest.param(_chat_httpx, id="chat-httpx"), pytest.param(_chat_async_openai_sdk, id="chat-async-openai-sdk"), - pytest.param(_responses_httpx, id="responses-httpx", marks=_NO_DISCONNECT_ROW_LIT_8603), + pytest.param(_responses_httpx, id="responses-httpx"), pytest.param(_messages_anthropic_sdk, id="messages-anthropic-sdk", marks=_NO_DISCONNECT_ROW_LIT_8603), pytest.param( _gemini_httpx, @@ -2902,10 +2912,10 @@ def test_concurrent_disconnects_during_a_guardrail_outage_each_record_exactly_on seconds=20, ) assert scanned == (1,) * 30, scanned - chat_rows: Final = _post_call_statuses(rig, rig.model, rows=10) - assert sorted(chat_rows) == sorted( - ("guardrail_failed_to_respond",) if index % 3 == 0 else ("success",) for index in range(0, 30, 3) - ), chat_rows + rows: Final = _post_call_rows(rig, rig.model, rows=20, call_types=frozenset({"acompletion", "aresponses"})) + assert sorted(rows) == sorted( + (("acompletion", ("guardrail_failed_to_respond",)),) * 10 + (("aresponses", ("success",)),) * 10 + ), rows def test_disconnect_scans_keep_recording_after_a_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/integration/observability/test_langtrace_delivery.py b/tests/integration/observability/test_langtrace_delivery.py index f81780f9f72..de953788861 100644 --- a/tests/integration/observability/test_langtrace_delivery.py +++ b/tests/integration/observability/test_langtrace_delivery.py @@ -610,14 +610,7 @@ def test_langtrace_sink_outage_mid_burst_recovers_on_the_same_port(gateway: Gate with wire_server(_accepted, port=port) as sink: rig: Final = _Rig(owned.gateway, model, provider, _Sink(sink, key)) _assert_each_once(rig.sink, _burst(rig, 6)) - log: Final = owned.log - failures_before: Final = log.read_text().count("Exception while exporting Span batch") outage: Final = _burst(rig, 12) - eventually( - lambda: log.read_text().count("Exception while exporting Span batch"), - lambda value: value > failures_before, - seconds=20, - ) assert owned.gateway.request("GET", "/health/liveliness").status_code == 200 with wire_server(_accepted, port=port) as revived: recovered: Final = _Rig(owned.gateway, model, provider, _Sink(revived, key)) diff --git a/tests/integration/observability/test_otel_trace_response_header.py b/tests/integration/observability/test_otel_trace_response_header.py index e84aeab48cf..a9846d184f5 100644 --- a/tests/integration/observability/test_otel_trace_response_header.py +++ b/tests/integration/observability/test_otel_trace_response_header.py @@ -24,7 +24,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter MARKER: Final = re.compile(rb"traceresp-[0-9a-f]{32}") -TRACERESPONSE: Final = re.compile(r"00-([0-9a-f]{32})-([0-9a-f]{16})-01") +TRACERESPONSE: Final = re.compile(r"00-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})") +SAMPLED: Final = 0x01 SERVER: Final = 2 CALL_ID: Final = "litellm.call_id" ANSWER: Final = "trace ok" @@ -134,8 +135,13 @@ def _query(request: Request) -> Mapping[str, str]: return MappingProxyType({name: values[0] for name, values in parse_qs(urlsplit(request.target).query).items()}) +def _sampled(value: str) -> re.Match[str] | None: + match: Final = TRACERESPONSE.fullmatch(value) + return match if match is not None and int(match.group(3), 16) & SAMPLED else None + + def _continued_traceresponse(request: Request) -> Reply: - incoming: Final = TRACERESPONSE.fullmatch(request.headers.get("traceparent", "")) + incoming: Final = _sampled(request.headers.get("traceparent", "")) if incoming is None: return Reply(status=400, body=b'{"error":"no sampled traceparent reached the upstream"}') return Reply(body=b'{"case":"continued"}', headers={"traceresponse": f"00-{incoming.group(1)}-{_span_id()}-01"}) @@ -372,7 +378,7 @@ def control(factory: Callable[[bool], RigFactory]) -> Iterator[Rig]: def _named(headers: httpx.Headers) -> tuple[str, str]: values: Final = headers.get_list("traceresponse") assert len(values) == 1, values - match: Final = TRACERESPONSE.fullmatch(values[0]) + match: Final = _sampled(values[0]) assert match is not None, values return match.group(1), match.group(2) diff --git a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py index 428f0015766..823c9e775e5 100644 --- a/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py +++ b/tests/integration/providers/test_bedrock_converse_lookaround_regex_wire.py @@ -577,10 +577,10 @@ def test_a_deployment_flag_never_reaches_its_sibling_on_the_same_model( @pytest.mark.parametrize( ("model_id", "body_base_model", "expected"), - ((_NOVA, f"bedrock/{_KIMI}", _WIRE_LOOKAROUND_FREE), (_KIMI, f"bedrock/{_NOVA}", _WIRE_LOOKAROUND_FREE)), - ids=("client-base-model-can-loosen-an-unflagged-deployment", "client-base-model-cannot-restore-a-flagged-one"), + ((_NOVA, f"bedrock/{_KIMI}", _WIRE_AS_SENT), (_KIMI, f"bedrock/{_NOVA}", _WIRE_LOOKAROUND_FREE)), + ids=("client-base-model-cannot-flag-an-unflagged-deployment", "client-base-model-cannot-restore-a-flagged-one"), ) -def test_a_base_model_in_the_request_body_only_ever_loosens( +def test_a_base_model_in_the_request_body_leaves_the_deployment_decision( gateway: Gateway, model_id: str, body_base_model: str, expected: Mapping[str, JsonValue] ) -> None: with wire_server(_bedrock_peer) as wire, gateway.scenario() as scenario: