mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(bedrock): keep Mantle token counts on the deployment's own endpoint
The Mantle count endpoint ignored the deployment's api_base and aws_bedrock_runtime_endpoint, so a bedrock/mantle deployment counted on BEDROCK_MANTLE_API_BASE or the public host, and a runtime deployment pinned to a private endpoint retried its counts on the public Mantle host. A bedrock/mantle deployment now counts where its inference goes. Other deployments retry on BEDROCK_MANTLE_API_BASE when it is set, on the public host only when they use the default runtime endpoint, and otherwise keep the runtime's own answer with no Mantle retry.
This commit is contained in:
mateo-berri
parent
67ae3fb527
commit
3520696cac
5 changed files with 268 additions and 13 deletions
|
|
@ -1323,7 +1323,7 @@ MANTLE_MESSAGES_PATH: Final = "/anthropic/v1/messages"
|
|||
_MANTLE_OPENAI_BASE_SUFFIXES: Final = ("/openai/v1", "/v1")
|
||||
|
||||
|
||||
def _mantle_api_base_from_env() -> str | None:
|
||||
def mantle_api_base_from_env() -> str | None:
|
||||
env_base: Final = get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
if env_base is None:
|
||||
return None
|
||||
|
|
@ -1383,7 +1383,7 @@ def build_mantle_messages_url(
|
|||
shared with the OpenAI-surface ``bedrock_mantle/*`` routes, which need it to
|
||||
carry their ``/v1`` or ``/openai/v1`` base, so that suffix is dropped first.
|
||||
"""
|
||||
override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env()
|
||||
override: Final = api_base or aws_bedrock_runtime_endpoint or mantle_api_base_from_env()
|
||||
if override:
|
||||
return ensure_api_route(override, MANTLE_MESSAGES_PATH)
|
||||
return f"https://bedrock-mantle.{region}.api.aws{MANTLE_MESSAGES_PATH}"
|
||||
|
|
|
|||
|
|
@ -15,9 +15,13 @@ from litellm._logging import verbose_logger
|
|||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
|
||||
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
||||
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
|
||||
from litellm.llms.bedrock.count_tokens.mantle_handler import (
|
||||
BedrockMantleCountTokensHandler,
|
||||
deployment_mantle_endpoint,
|
||||
)
|
||||
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -74,6 +78,10 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
|
|||
)
|
||||
|
||||
|
||||
async def _has_mantle_endpoint(litellm_params: Mapping[str, object]) -> bool:
|
||||
return await run_aws_signing(deployment_mantle_endpoint, litellm_params) is not None
|
||||
|
||||
|
||||
def _request_data(
|
||||
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -150,8 +158,9 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
use_mantle: Final = model_to_use.startswith("bedrock_mantle/")
|
||||
counted_model: Final = model_to_use.removeprefix("bedrock_mantle/")
|
||||
resolved_model: Final = get_bedrock_base_model(counted_model)
|
||||
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE
|
||||
litellm_params: Final = (deployment or {}).get("litellm_params", {})
|
||||
counts_on_mantle: Final = use_mantle and await _has_mantle_endpoint(litellm_params)
|
||||
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if counts_on_mantle else RUNTIME_TOKENIZER_TYPE
|
||||
translated: Final = (
|
||||
await anthropic_count_fields(
|
||||
counted_model,
|
||||
|
|
@ -178,7 +187,7 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
await self._count_with(
|
||||
self._mantle_handler, primary_tokenizer, mantle_request, litellm_params, resolved_model
|
||||
)
|
||||
if use_mantle
|
||||
if counts_on_mantle
|
||||
else await self._count_with(
|
||||
self._runtime_handler, primary_tokenizer, runtime_request, litellm_params, resolved_model
|
||||
)
|
||||
|
|
@ -187,7 +196,9 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
await self._count_with(
|
||||
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
|
||||
)
|
||||
if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model)
|
||||
if not use_mantle
|
||||
and _runtime_rejected_claude_model(primary, resolved_model)
|
||||
and await _has_mantle_endpoint(litellm_params)
|
||||
else primary
|
||||
)
|
||||
match outcome:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -11,11 +12,16 @@ from litellm.litellm_core_utils.url_utils import join_api_base
|
|||
from litellm.llms.anthropic.common_utils import tools_without_custom_field
|
||||
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url, mantle_api_base_from_env
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
|
||||
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
|
||||
NO_MANTLE_ENDPOINT_MESSAGE: Final = (
|
||||
"This deployment pins its own bedrock-runtime endpoint, so its prompt is not counted on the public "
|
||||
"bedrock-mantle host. Set BEDROCK_MANTLE_API_BASE to the bedrock-mantle endpoint it may use."
|
||||
)
|
||||
|
||||
|
||||
class MantleCountTokensRequest(TypedDict):
|
||||
|
|
@ -30,6 +36,39 @@ _COUNT_OPTIONS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mappin
|
|||
_TOOLS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MantleEndpointSettings:
|
||||
api_base: str | None
|
||||
aws_bedrock_runtime_endpoint: str | None
|
||||
|
||||
|
||||
_OPERATOR_MANTLE_ENDPOINT: Final = MantleEndpointSettings(api_base=None, aws_bedrock_runtime_endpoint=None)
|
||||
|
||||
|
||||
def _endpoint_setting(litellm_params: Mapping[str, object], name: str) -> str | None:
|
||||
value: Final = litellm_params.get(name)
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _deploys_on_mantle(litellm_params: Mapping[str, object]) -> bool:
|
||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
model: Final = litellm_params.get("model")
|
||||
return isinstance(model, str) and BedrockModelInfo.get_bedrock_route(model) == "mantle"
|
||||
|
||||
|
||||
def deployment_mantle_endpoint(litellm_params: Mapping[str, object]) -> MantleEndpointSettings | None:
|
||||
api_base: Final = _endpoint_setting(litellm_params, "api_base")
|
||||
runtime_endpoint: Final = _endpoint_setting(litellm_params, "aws_bedrock_runtime_endpoint")
|
||||
if _deploys_on_mantle(litellm_params):
|
||||
return MantleEndpointSettings(api_base=api_base, aws_bedrock_runtime_endpoint=runtime_endpoint)
|
||||
if mantle_api_base_from_env() or not (
|
||||
api_base or runtime_endpoint or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT")
|
||||
):
|
||||
return _OPERATOR_MANTLE_ENDPOINT
|
||||
return None
|
||||
|
||||
|
||||
class BedrockMantleCountTokensHandler(BaseAWSLLM):
|
||||
"""Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint.
|
||||
|
||||
|
|
@ -42,9 +81,16 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
|
|||
super().__init__()
|
||||
self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig()
|
||||
|
||||
def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str:
|
||||
def get_mantle_count_tokens_endpoint(
|
||||
self, litellm_params: Mapping[str, object], aws_region_name: str
|
||||
) -> str | None:
|
||||
endpoint: Final = deployment_mantle_endpoint(litellm_params)
|
||||
if endpoint is None:
|
||||
return None
|
||||
messages_url: Final = build_mantle_messages_url(
|
||||
api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name
|
||||
api_base=endpoint.api_base,
|
||||
aws_bedrock_runtime_endpoint=endpoint.aws_bedrock_runtime_endpoint,
|
||||
region=aws_region_name,
|
||||
)
|
||||
return join_api_base(messages_url, MANTLE_COUNT_TOKENS_SUFFIX)
|
||||
|
||||
|
|
@ -60,7 +106,11 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
|
|||
aws_region_name: Final = await run_aws_signing(
|
||||
self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None
|
||||
)
|
||||
endpoint_url: Final = await run_aws_signing(self.get_mantle_count_tokens_endpoint, aws_region_name)
|
||||
endpoint_url: Final = await run_aws_signing(
|
||||
self.get_mantle_count_tokens_endpoint, litellm_params, aws_region_name
|
||||
)
|
||||
if endpoint_url is None:
|
||||
raise BedrockError(status_code=400, message=NO_MANTLE_ENDPOINT_MESSAGE)
|
||||
tools: Final = request.get("tools")
|
||||
body: Final = self._anthropic_config.transform_request_to_count_tokens(
|
||||
model=resolved_model,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|||
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
|
||||
|
||||
MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens"
|
||||
PRIVATE_MANTLE: Final = "https://vpce-0e1f2a3b.bedrock-mantle.us-east-1.vpce.api.aws"
|
||||
OPERATOR_MANTLE: Final = "https://vpce-9f8e7d6c.bedrock-mantle.us-east-1.vpce.api.aws"
|
||||
LITELLM_PARAMS: Final = {
|
||||
"aws_access_key_id": "AKIATESTACCESSKEY",
|
||||
"aws_secret_access_key": "test-secret",
|
||||
|
|
@ -47,6 +49,7 @@ def _client(endpoint: _MantleEndpoint) -> AsyncHTTPHandler:
|
|||
def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -141,7 +144,9 @@ async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None:
|
|||
@pytest.mark.asyncio
|
||||
async def test_region_and_endpoint_lookups_leave_the_event_loop_free(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_REGION_NAME", "BEDROCK_MANTLE_API_BASE")
|
||||
vault: Final = host_in_loop_probing_vault(
|
||||
monkeypatch, "AWS_REGION_NAME", "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT"
|
||||
)
|
||||
mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177})
|
||||
|
||||
result = await BedrockMantleCountTokensHandler().handle_count_tokens_request(
|
||||
|
|
@ -155,3 +160,52 @@ async def test_region_and_endpoint_lookups_leave_the_event_loop_free(monkeypatch
|
|||
assert not vault.read_on_loop.is_set()
|
||||
assert result == {"input_tokens": 2177}
|
||||
assert str(mantle.only_request().url) == MANTLE_COUNT_URL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("endpoint", "environment"),
|
||||
(
|
||||
({"api_base": PRIVATE_MANTLE}, {}),
|
||||
({"api_base": f"{PRIVATE_MANTLE}/anthropic/v1/messages"}, {}),
|
||||
({"aws_bedrock_runtime_endpoint": PRIVATE_MANTLE}, {}),
|
||||
({"api_base": PRIVATE_MANTLE}, {"BEDROCK_MANTLE_API_BASE": OPERATOR_MANTLE}),
|
||||
),
|
||||
ids=("api_base", "api_base_with_messages_path", "aws_bedrock_runtime_endpoint", "api_base_over_env"),
|
||||
)
|
||||
async def test_mantle_deployment_counts_on_the_endpoint_its_inference_uses(
|
||||
monkeypatch: pytest.MonkeyPatch, endpoint: Mapping[str, str], environment: Mapping[str, str]
|
||||
) -> None:
|
||||
for name, value in environment.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
mantle: Final = _MantleEndpoint(200, {"input_tokens": 3})
|
||||
|
||||
await BedrockMantleCountTokensHandler().handle_count_tokens_request(
|
||||
request_data=dict(REQUEST),
|
||||
litellm_params={**LITELLM_PARAMS, "model": "bedrock/mantle/anthropic.claude-opus-5-5", **endpoint},
|
||||
resolved_model="anthropic.claude-opus-5-5",
|
||||
client=_client(mantle),
|
||||
)
|
||||
|
||||
assert str(mantle.only_request().url) == f"{PRIVATE_MANTLE}/anthropic/v1/messages/count_tokens"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_deployment_with_its_own_endpoint_is_never_counted_on_public_mantle() -> None:
|
||||
mantle: Final = _MantleEndpoint(200, {"input_tokens": 3})
|
||||
|
||||
with pytest.raises(BedrockError) as raised:
|
||||
await BedrockMantleCountTokensHandler().handle_count_tokens_request(
|
||||
request_data=dict(REQUEST),
|
||||
litellm_params={
|
||||
**LITELLM_PARAMS,
|
||||
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"api_base": "https://vpce-0a1b2c3d.bedrock-runtime.us-east-1.vpce.amazonaws.com",
|
||||
},
|
||||
resolved_model="anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
client=_client(mantle),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert "BEDROCK_MANTLE_API_BASE" in raised.value.message
|
||||
assert mantle.requests == ()
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from pydantic import JsonValue, TypeAdapter
|
|||
import litellm
|
||||
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
|
||||
|
||||
RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com"
|
||||
MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws"
|
||||
|
|
@ -21,23 +22,37 @@ DEPLOYMENT: Final = {
|
|||
}
|
||||
}
|
||||
MESSAGES: Final = [{"role": "user", "content": "The quick brown fox jumps over the lazy dog."}]
|
||||
PRIVATE_RUNTIME_HOST: Final = "vpce-0a1b2c3d.bedrock-runtime.us-east-1.vpce.amazonaws.com"
|
||||
PRIVATE_RUNTIME: Final = f"https://{PRIVATE_RUNTIME_HOST}"
|
||||
PRIVATE_MANTLE: Final = "https://vpce-0e1f2a3b.bedrock-mantle.us-east-1.vpce.api.aws"
|
||||
SONNET_4_5: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
SONNET_4_5_RUNTIME_COUNT: Final = f"{PRIVATE_RUNTIME}/model/anthropic.claude-sonnet-4-5-20250929-v1%3A0/count-tokens"
|
||||
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class _Bedrock:
|
||||
def __init__(self, runtime: tuple[int, Mapping[str, object]], mantle: tuple[int, Mapping[str, object]]) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
runtime: tuple[int, Mapping[str, object]],
|
||||
mantle: tuple[int, Mapping[str, object]],
|
||||
runtime_host: str = RUNTIME_HOST,
|
||||
) -> None:
|
||||
self.runtime: Final = runtime
|
||||
self.mantle: Final = mantle
|
||||
self.runtime_host: Final = runtime_host
|
||||
self.requests: tuple[httpx.Request, ...] = ()
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.requests = (*self.requests, request)
|
||||
status_code, body = self.runtime if request.url.host == RUNTIME_HOST else self.mantle
|
||||
status_code, body = self.runtime if request.url.host == self.runtime_host else self.mantle
|
||||
return httpx.Response(status_code, json=dict(body), request=request)
|
||||
|
||||
def posted_hosts(self) -> tuple[str, ...]:
|
||||
return tuple(request.url.host for request in self.requests)
|
||||
|
||||
def posted_urls(self) -> tuple[str, ...]:
|
||||
return tuple(str(request.url) for request in self.requests)
|
||||
|
||||
|
||||
def _counter(bedrock: _Bedrock) -> BedrockTokenCounter:
|
||||
return BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock)))
|
||||
|
|
@ -47,6 +62,7 @@ def _counter(bedrock: _Bedrock) -> BedrockTokenCounter:
|
|||
def _sigv4_only(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -420,3 +436,127 @@ async def test_explicit_mantle_failure_does_not_probe_runtime() -> None:
|
|||
assert result.status_code == 404
|
||||
assert result.tokenizer_type == "bedrock_mantle_api"
|
||||
assert bedrock.posted_hosts() == (MANTLE_HOST,)
|
||||
|
||||
|
||||
def _pinned_runtime_deployment(endpoint: Mapping[str, str]) -> dict[str, object]:
|
||||
return {"litellm_params": {**DEPLOYMENT["litellm_params"], "model": f"bedrock/{SONNET_4_5}", **endpoint}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("endpoint", "environment"),
|
||||
(
|
||||
({"api_base": PRIVATE_RUNTIME}, {}),
|
||||
({"aws_bedrock_runtime_endpoint": PRIVATE_RUNTIME}, {}),
|
||||
({}, {"AWS_BEDROCK_RUNTIME_ENDPOINT": PRIVATE_RUNTIME}),
|
||||
),
|
||||
ids=("api_base", "aws_bedrock_runtime_endpoint", "AWS_BEDROCK_RUNTIME_ENDPOINT"),
|
||||
)
|
||||
async def test_pinned_runtime_endpoint_keeps_its_own_rejection_and_sends_nothing_to_mantle(
|
||||
monkeypatch: pytest.MonkeyPatch, endpoint: Mapping[str, str], environment: Mapping[str, str]
|
||||
) -> None:
|
||||
for name, value in environment.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
bedrock: Final = _Bedrock(
|
||||
runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}), runtime_host=PRIVATE_RUNTIME_HOST
|
||||
)
|
||||
|
||||
result: Final = await _counter(bedrock).count_tokens(
|
||||
model_to_use=SONNET_4_5,
|
||||
messages=MESSAGES,
|
||||
contents=None,
|
||||
deployment=_pinned_runtime_deployment(endpoint),
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert (result.error, result.status_code, result.tokenizer_type) == (True, 400, "bedrock_api")
|
||||
assert UNSUPPORTED["message"] in (result.error_message or "")
|
||||
assert bedrock.posted_urls() == (SONNET_4_5_RUNTIME_COUNT,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinned_runtime_endpoint_retries_on_the_mantle_host_the_operator_named(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_BASE", PRIVATE_MANTLE)
|
||||
bedrock: Final = _Bedrock(
|
||||
runtime=(400, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}), runtime_host=PRIVATE_RUNTIME_HOST
|
||||
)
|
||||
|
||||
result: Final = await _counter(bedrock).count_tokens(
|
||||
model_to_use=SONNET_4_5,
|
||||
messages=MESSAGES,
|
||||
contents=None,
|
||||
deployment=_pinned_runtime_deployment({"api_base": PRIVATE_RUNTIME}),
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert (result.error, result.total_tokens, result.tokenizer_type) == (False, 2177, "bedrock_mantle_api")
|
||||
assert bedrock.posted_urls() == (
|
||||
SONNET_4_5_RUNTIME_COUNT,
|
||||
f"{PRIVATE_MANTLE}/anthropic/v1/messages/count_tokens",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("runtime", "expected"),
|
||||
(
|
||||
((200, {"inputTokens": 193}), (False, 193, None)),
|
||||
((400, UNSUPPORTED), (True, 0, 400)),
|
||||
),
|
||||
ids=("runtime_counts", "runtime_rejects"),
|
||||
)
|
||||
async def test_mantle_counting_model_on_a_pinned_runtime_deployment_counts_on_that_endpoint_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
runtime: tuple[int, Mapping[str, object]],
|
||||
expected: tuple[bool, int, int | None],
|
||||
) -> None:
|
||||
vault: Final = host_in_loop_probing_vault(monkeypatch, "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT")
|
||||
bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 271}), runtime_host=PRIVATE_RUNTIME_HOST)
|
||||
|
||||
result: Final = await _counter(bedrock).count_tokens(
|
||||
model_to_use="bedrock_mantle/anthropic.claude-opus-5-5",
|
||||
messages=MESSAGES,
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
**DEPLOYMENT["litellm_params"],
|
||||
"model": "bedrock/us.anthropic.claude-opus-5-5",
|
||||
"api_base": PRIVATE_RUNTIME,
|
||||
}
|
||||
},
|
||||
request_model="claude-opus-5-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert (result.error, result.total_tokens, result.status_code) == expected
|
||||
assert result.tokenizer_type == "bedrock_api"
|
||||
assert bedrock.posted_urls() == (f"{PRIVATE_RUNTIME}/model/anthropic.claude-opus-5-5/count-tokens",)
|
||||
assert vault.read.is_set()
|
||||
assert not vault.read_on_loop.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mantle_deployment_counts_on_its_own_api_base() -> None:
|
||||
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271}))
|
||||
|
||||
result: Final = await _counter(bedrock).count_tokens(
|
||||
model_to_use="bedrock_mantle/anthropic.claude-opus-5-5",
|
||||
messages=MESSAGES,
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
**DEPLOYMENT["litellm_params"],
|
||||
"model": "bedrock/mantle/anthropic.claude-opus-5-5",
|
||||
"api_base": PRIVATE_MANTLE,
|
||||
}
|
||||
},
|
||||
request_model="claude-opus-5-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert (result.error, result.total_tokens, result.tokenizer_type) == (False, 271, "bedrock_mantle_api")
|
||||
assert bedrock.posted_urls() == (f"{PRIVATE_MANTLE}/anthropic/v1/messages/count_tokens",)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue