mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
test(bedrock): cover SDK count_tokens cross-region prefix path
This commit is contained in:
parent
ba0caf1ad8
commit
61d8d8e1f6
1 changed files with 33 additions and 0 deletions
|
|
@ -21,6 +21,7 @@ sys.path.insert(
|
|||
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
|
||||
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
||||
from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest
|
||||
|
||||
|
||||
|
|
@ -112,6 +113,38 @@ class TestBedrockTokenCounter(BaseTokenCounterTest):
|
|||
), f"Token counting should not error: {result.error_message}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sdk_token_counter_preserves_cross_region_prefix():
|
||||
"""
|
||||
Regression test for #32683 on the SDK token_counter() path: the resolved
|
||||
model passed to the Bedrock CountTokens handler must keep the cross-region
|
||||
inference-profile prefix instead of collapsing to the bare foundation model.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
counter = BedrockTokenCounter()
|
||||
|
||||
with patch.object(
|
||||
BedrockCountTokensHandler,
|
||||
"handle_count_tokens_request",
|
||||
new=AsyncMock(return_value={"input_tokens": 7}),
|
||||
) as mock_handler:
|
||||
result = await counter.count_tokens(
|
||||
model_to_use="bedrock/global.anthropic.claude-opus-4-8",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"aws_region_name": "eu-central-1"}},
|
||||
request_model="global.anthropic.claude-opus-4-8",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.total_tokens == 7
|
||||
assert (
|
||||
mock_handler.call_args.kwargs["resolved_model"]
|
||||
== "global.anthropic.claude-opus-4-8"
|
||||
)
|
||||
|
||||
|
||||
class TestBedrockCountTokensEndpoint:
|
||||
"""Unit tests for custom endpoint URL resolution in BedrockCountTokensConfig."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue