751 lines
31 KiB
Python
751 lines
31 KiB
Python
"""
|
|
Tests for AWS Bedrock GovCloud model support
|
|
"""
|
|
|
|
import os
|
|
|
|
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" # Load from local file
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
|
|
# Import modules that need to be reloaded
|
|
import importlib
|
|
import litellm.litellm_core_utils.get_model_cost_map
|
|
import litellm
|
|
|
|
# Reload modules to pick up environment variable
|
|
importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
|
|
importlib.reload(litellm)
|
|
|
|
from litellm import completion
|
|
from litellm.llms.bedrock.common_utils import (
|
|
BedrockModelInfo,
|
|
AmazonBedrockGlobalConfig,
|
|
)
|
|
|
|
|
|
class TestBedrockGovCloudSupport:
|
|
"""Test suite for GovCloud model support in Bedrock"""
|
|
|
|
def test_govcloud_regions_in_config(self):
|
|
"""Test that GovCloud regions are included in the configuration"""
|
|
config = AmazonBedrockGlobalConfig()
|
|
us_regions = config.get_us_regions()
|
|
|
|
assert "us-gov-east-1" in us_regions
|
|
assert "us-gov-west-1" in us_regions
|
|
|
|
all_regions = config.get_all_regions()
|
|
assert "us-gov-east-1" in all_regions
|
|
assert "us-gov-west-1" in all_regions
|
|
|
|
def test_govcloud_models_in_model_cost(self):
|
|
"""Test that GovCloud models are present in model cost configuration"""
|
|
from litellm import model_cost
|
|
|
|
# Test Claude models in GovCloud
|
|
assert (
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
in model_cost
|
|
)
|
|
assert (
|
|
"bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
in model_cost
|
|
)
|
|
assert (
|
|
"bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
|
|
)
|
|
assert (
|
|
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost
|
|
)
|
|
assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
|
|
assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost
|
|
|
|
# Test Llama models in GovCloud
|
|
assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost
|
|
assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost
|
|
assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost
|
|
assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost
|
|
|
|
# Test Titan models in GovCloud
|
|
assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost
|
|
assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost
|
|
|
|
def test_govcloud_model_routing(self):
|
|
"""Test that GovCloud models are routed correctly"""
|
|
# Test Claude model routing
|
|
route = BedrockModelInfo.get_bedrock_route(
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
)
|
|
assert route == "converse"
|
|
|
|
route = BedrockModelInfo.get_bedrock_route(
|
|
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0"
|
|
)
|
|
assert route == "converse"
|
|
|
|
# Test Llama model routing
|
|
route = BedrockModelInfo.get_bedrock_route(
|
|
"bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0"
|
|
)
|
|
assert route == "converse"
|
|
|
|
route = BedrockModelInfo.get_bedrock_route(
|
|
"bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0"
|
|
)
|
|
assert route == "converse"
|
|
|
|
# Test Titan model routing (should use invoke)
|
|
route = BedrockModelInfo.get_bedrock_route(
|
|
"bedrock/us-gov-east-1/amazon.titan-text-lite-v1"
|
|
)
|
|
assert route == "invoke"
|
|
|
|
def test_base_model_extraction(self):
|
|
"""Test that base model names are correctly extracted from GovCloud models"""
|
|
# Test GovCloud model extraction
|
|
base_model = BedrockModelInfo.get_base_model(
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
)
|
|
assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
|
|
base_model = BedrockModelInfo.get_base_model(
|
|
"bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"
|
|
)
|
|
assert base_model == "meta.llama3-8b-instruct-v1:0"
|
|
|
|
@patch("litellm.llms.bedrock.common_utils.init_bedrock_client")
|
|
def test_govcloud_client_initialization(self, mock_init_client):
|
|
"""Test that Bedrock client can be initialized with GovCloud regions"""
|
|
mock_client = Mock()
|
|
mock_init_client.return_value = mock_client
|
|
|
|
# Test that init_bedrock_client accepts GovCloud regions
|
|
from litellm.llms.bedrock.common_utils import init_bedrock_client
|
|
|
|
# This should not raise an error
|
|
client = init_bedrock_client(
|
|
region_name="us-gov-east-1",
|
|
aws_access_key_id=None,
|
|
aws_secret_access_key=None,
|
|
aws_region_name="us-gov-east-1",
|
|
aws_bedrock_runtime_endpoint=None,
|
|
aws_session_name=None,
|
|
aws_profile_name=None,
|
|
aws_role_name=None,
|
|
aws_web_identity_token=None,
|
|
extra_headers=None,
|
|
timeout=None,
|
|
)
|
|
|
|
assert mock_init_client.called
|
|
|
|
def test_govcloud_model_in_bedrock_models_list(self):
|
|
"""Test that GovCloud models are NOT included in bedrock_models list (they are pricing-only)"""
|
|
# Regional models including GovCloud should be excluded from bedrock_models list
|
|
# They are only in model_cost for pricing purposes
|
|
assert not any("us-gov-east-1" in model for model in litellm.bedrock_models)
|
|
assert not any("us-gov-west-1" in model for model in litellm.bedrock_models)
|
|
|
|
def test_govcloud_model_cost_properties(self):
|
|
"""Test that GovCloud models have proper cost configuration"""
|
|
from litellm import model_cost
|
|
|
|
# Check a specific GovCloud model has all required properties
|
|
govcloud_model = model_cost[
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
]
|
|
|
|
assert "max_tokens" in govcloud_model
|
|
assert "max_input_tokens" in govcloud_model
|
|
assert "max_output_tokens" in govcloud_model
|
|
assert "input_cost_per_token" in govcloud_model
|
|
assert "output_cost_per_token" in govcloud_model
|
|
assert govcloud_model["litellm_provider"] == "bedrock"
|
|
assert govcloud_model["mode"] == "chat"
|
|
|
|
def test_govcloud_model_pricing_verification(self):
|
|
"""Test that GovCloud models have correct pricing that differs from base models"""
|
|
from litellm import model_cost
|
|
|
|
# Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id
|
|
base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
gov_east_model = (
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
)
|
|
gov_west_model = (
|
|
"bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
)
|
|
|
|
# Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok)
|
|
base_pricing = model_cost[base_model]
|
|
assert base_pricing["input_cost_per_token"] == 1.1e-06
|
|
assert base_pricing["output_cost_per_token"] == 5.5e-06
|
|
|
|
# Verify GovCloud models have different (higher) pricing
|
|
gov_east_pricing = model_cost[gov_east_model]
|
|
gov_west_pricing = model_cost[gov_west_model]
|
|
|
|
# GovCloud models should have ~20% higher pricing than base models
|
|
assert gov_east_pricing["input_cost_per_token"] == 1.2e-06
|
|
assert gov_east_pricing["output_cost_per_token"] == 6e-06
|
|
assert gov_west_pricing["input_cost_per_token"] == 1.2e-06
|
|
assert gov_west_pricing["output_cost_per_token"] == 6e-06
|
|
|
|
# Verify the pricing difference is approximately 20%
|
|
assert (
|
|
abs(
|
|
gov_east_pricing["input_cost_per_token"]
|
|
/ base_pricing["input_cost_per_token"]
|
|
- 1.2
|
|
)
|
|
< 0.15
|
|
)
|
|
assert (
|
|
abs(
|
|
gov_east_pricing["output_cost_per_token"]
|
|
/ base_pricing["output_cost_per_token"]
|
|
- 1.2
|
|
)
|
|
< 0.15
|
|
)
|
|
assert (
|
|
abs(
|
|
gov_west_pricing["input_cost_per_token"]
|
|
/ base_pricing["input_cost_per_token"]
|
|
- 1.2
|
|
)
|
|
< 0.15
|
|
)
|
|
assert (
|
|
abs(
|
|
gov_west_pricing["output_cost_per_token"]
|
|
/ base_pricing["output_cost_per_token"]
|
|
- 1.2
|
|
)
|
|
< 0.15
|
|
)
|
|
|
|
# Test Claude 3 Haiku pricing
|
|
base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0"
|
|
gov_east_haiku_model = (
|
|
"bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0"
|
|
)
|
|
gov_west_haiku_model = (
|
|
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0"
|
|
)
|
|
|
|
# Verify base Haiku model pricing
|
|
base_haiku_pricing = model_cost[base_haiku_model]
|
|
assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07 # 0.00000025
|
|
assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06 # 0.00000125
|
|
|
|
# Verify GovCloud Haiku models have different (higher) pricing
|
|
gov_east_haiku_pricing = model_cost[gov_east_haiku_model]
|
|
gov_west_haiku_pricing = model_cost[gov_west_haiku_model]
|
|
|
|
# GovCloud Haiku models should have 20% higher pricing than base models
|
|
assert (
|
|
gov_east_haiku_pricing["input_cost_per_token"] == 3e-07
|
|
) # 0.0000003 (20% higher)
|
|
assert (
|
|
gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06
|
|
) # 0.0000015 (20% higher)
|
|
assert (
|
|
gov_west_haiku_pricing["input_cost_per_token"] == 3e-07
|
|
) # 0.0000003 (20% higher)
|
|
assert (
|
|
gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06
|
|
) # 0.0000015 (20% higher)
|
|
|
|
# Verify the pricing difference is exactly 20%
|
|
assert (
|
|
gov_east_haiku_pricing["input_cost_per_token"]
|
|
== base_haiku_pricing["input_cost_per_token"] * 1.2
|
|
)
|
|
assert (
|
|
gov_east_haiku_pricing["output_cost_per_token"]
|
|
== base_haiku_pricing["output_cost_per_token"] * 1.2
|
|
)
|
|
assert (
|
|
gov_west_haiku_pricing["input_cost_per_token"]
|
|
== base_haiku_pricing["input_cost_per_token"] * 1.2
|
|
)
|
|
assert (
|
|
gov_west_haiku_pricing["output_cost_per_token"]
|
|
== base_haiku_pricing["output_cost_per_token"] * 1.2
|
|
)
|
|
|
|
@patch("litellm.completion")
|
|
def test_govcloud_completion_cost_calculation(self, mock_completion):
|
|
"""Test that completion requests use correct pricing for GovCloud models"""
|
|
from litellm import completion_cost, Choices, Message, ModelResponse
|
|
from litellm.utils import Usage
|
|
|
|
# Mock completion response for base model
|
|
# Use us.* inference profile ID to match us.* pricing ($1.10/$5.50 per MTok)
|
|
base_model_response = ModelResponse(
|
|
id="test-base",
|
|
choices=[
|
|
Choices(
|
|
finish_reason="stop",
|
|
index=0,
|
|
message=Message(content="Hello", role="assistant"),
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
object="chat.completion",
|
|
system_fingerprint=None,
|
|
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
)
|
|
base_model_response._hidden_params = {
|
|
"custom_llm_provider": "bedrock",
|
|
"region_name": "us-east-1",
|
|
}
|
|
|
|
# Mock completion response for gov model
|
|
# GovCloud responses use base anthropic.* model ID; pricing is looked up
|
|
# via bedrock/us-gov-east-1/anthropic.* entries in model_cost
|
|
gov_model_response = ModelResponse(
|
|
id="test-gov",
|
|
choices=[
|
|
Choices(
|
|
finish_reason="stop",
|
|
index=0,
|
|
message=Message(content="Hello", role="assistant"),
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
object="chat.completion",
|
|
system_fingerprint=None,
|
|
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
)
|
|
gov_model_response._hidden_params = {
|
|
"custom_llm_provider": "bedrock",
|
|
"region_name": "us-gov-east-1",
|
|
}
|
|
|
|
# Mock completion response for gov-west model
|
|
gov_west_model_response = ModelResponse(
|
|
id="test-gov-west",
|
|
choices=[
|
|
Choices(
|
|
finish_reason="stop",
|
|
index=0,
|
|
message=Message(content="Hello", role="assistant"),
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
object="chat.completion",
|
|
system_fingerprint=None,
|
|
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
)
|
|
gov_west_model_response._hidden_params = {
|
|
"custom_llm_provider": "bedrock",
|
|
"region_name": "us-gov-west-1",
|
|
}
|
|
|
|
# Test messages
|
|
messages = [{"role": "user", "content": "Hello, how are you?"}]
|
|
|
|
# Calculate costs using the standard Bedrock format with region parameter
|
|
# Base model uses us.* inference profile — no region_name needed since
|
|
# the response model already contains the us.* prefix for pricing lookup.
|
|
base_cost = completion_cost(
|
|
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
completion_response=base_model_response,
|
|
messages=messages,
|
|
)
|
|
|
|
# GovCloud models use region_name to look up bedrock/us-gov-*/anthropic.* pricing
|
|
gov_east_cost = completion_cost(
|
|
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
completion_response=gov_model_response,
|
|
messages=messages,
|
|
region_name="us-gov-east-1",
|
|
)
|
|
|
|
gov_west_cost = completion_cost(
|
|
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
completion_response=gov_west_model_response,
|
|
messages=messages,
|
|
region_name="us-gov-west-1",
|
|
)
|
|
|
|
# Expected costs based on pricing:
|
|
# Base model (us.*): 10 * 1.1e-06 + 5 * 5.5e-06 = 1.1e-05 + 2.75e-05 = 3.85e-05
|
|
# Gov models: 10 * 1.2e-06 + 5 * 6e-06 = 1.2e-05 + 3e-05 = 4.2e-05
|
|
expected_base_cost = 10 * 1.1e-06 + 5 * 5.5e-06
|
|
expected_gov_cost = 10 * 1.2e-06 + 5 * 6e-06
|
|
|
|
# Verify costs are calculated correctly
|
|
assert (
|
|
abs(base_cost - expected_base_cost) < 1e-10
|
|
), f"Base cost mismatch: got {base_cost}, expected {expected_base_cost}"
|
|
assert (
|
|
abs(gov_east_cost - expected_gov_cost) < 1e-10
|
|
), f"Gov East cost mismatch: got {gov_east_cost}, expected {expected_gov_cost}"
|
|
assert (
|
|
abs(gov_west_cost - expected_gov_cost) < 1e-10
|
|
), f"Gov West cost mismatch: got {gov_west_cost}, expected {expected_gov_cost}"
|
|
|
|
# Verify GovCloud costs are approximately 20% higher than base cost
|
|
assert (
|
|
abs(gov_east_cost / base_cost - 1.2) < 0.15
|
|
), f"Gov East cost should be ~20% higher than base: got {gov_east_cost}, base {base_cost}"
|
|
assert (
|
|
abs(gov_west_cost / base_cost - 1.2) < 0.15
|
|
), f"Gov West cost should be ~20% higher than base: got {gov_west_cost}, base {base_cost}"
|
|
|
|
# Test with different token counts
|
|
large_response = ModelResponse(
|
|
id="test-large",
|
|
choices=[
|
|
Choices(
|
|
finish_reason="stop",
|
|
index=0,
|
|
message=Message(content="A longer response", role="assistant"),
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
object="chat.completion",
|
|
system_fingerprint=None,
|
|
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
|
|
)
|
|
large_response._hidden_params = {
|
|
"custom_llm_provider": "bedrock",
|
|
"region_name": "us-east-1",
|
|
}
|
|
|
|
large_base_cost = completion_cost(
|
|
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
completion_response=large_response,
|
|
messages=messages,
|
|
)
|
|
|
|
# Create large response for gov model
|
|
large_gov_response = ModelResponse(
|
|
id="test-large-gov",
|
|
choices=[
|
|
Choices(
|
|
finish_reason="stop",
|
|
index=0,
|
|
message=Message(content="A longer response", role="assistant"),
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
object="chat.completion",
|
|
system_fingerprint=None,
|
|
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
|
|
)
|
|
large_gov_response._hidden_params = {
|
|
"custom_llm_provider": "bedrock",
|
|
"region_name": "us-gov-east-1",
|
|
}
|
|
|
|
large_gov_cost = completion_cost(
|
|
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
completion_response=large_gov_response,
|
|
messages=messages,
|
|
region_name="us-gov-east-1",
|
|
)
|
|
|
|
# Expected costs for larger response:
|
|
# Base model (us.*): 100 * 1.1e-06 + 50 * 5.5e-06 = 1.1e-04 + 2.75e-04 = 3.85e-04
|
|
# Gov model: 100 * 1.2e-06 + 50 * 6e-06 = 1.2e-04 + 3e-04 = 4.2e-04
|
|
expected_large_base_cost = 100 * 1.1e-06 + 50 * 5.5e-06
|
|
expected_large_gov_cost = 100 * 1.2e-06 + 50 * 6e-06
|
|
|
|
assert (
|
|
abs(large_base_cost - expected_large_base_cost) < 1e-10
|
|
), f"Large base cost mismatch: got {large_base_cost}, expected {expected_large_base_cost}"
|
|
assert (
|
|
abs(large_gov_cost - expected_large_gov_cost) < 1e-10
|
|
), f"Large gov cost mismatch: got {large_gov_cost}, expected {expected_large_gov_cost}"
|
|
assert (
|
|
abs(large_gov_cost / large_base_cost - 1.2) < 0.15
|
|
), f"Large gov cost should be ~20% higher than base: got {large_gov_cost}, base {large_base_cost}"
|
|
|
|
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
|
|
def test_govcloud_completion_with_cost_tracking(self, mock_post):
|
|
"""Test that completion requests with cost tracking use correct pricing for GovCloud models"""
|
|
from litellm import completion
|
|
from unittest.mock import Mock
|
|
import json
|
|
|
|
# Mock the HTTP client's post method to return responses
|
|
def mock_post_side_effect(url, headers=None, data=None, **kwargs):
|
|
# Extract region from the URL to determine which response to return
|
|
region = "us-east-1" # default
|
|
if "us-gov-east-1" in url:
|
|
region = "us-gov-east-1"
|
|
elif "us-gov-west-1" in url:
|
|
region = "us-gov-west-1"
|
|
|
|
# Create mock response based on region
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.headers = {}
|
|
|
|
# Create a realistic Bedrock converse response structure
|
|
bedrock_response = {
|
|
"output": {
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": f"Hello from {region}"}],
|
|
}
|
|
},
|
|
"usage": {"inputTokens": 15, "outputTokens": 8, "totalTokens": 23},
|
|
"stopReason": "end_turn",
|
|
}
|
|
|
|
mock_response.json.return_value = bedrock_response
|
|
mock_response.text = json.dumps(bedrock_response)
|
|
mock_response.raise_for_status = Mock() # Don't raise exceptions
|
|
|
|
return mock_response
|
|
|
|
mock_post.side_effect = mock_post_side_effect
|
|
|
|
# Test base model completion
|
|
base_result = completion(
|
|
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
# Test gov-east model completion
|
|
# GovCloud users specify the base anthropic.* model ID with the gov region
|
|
gov_east_result = completion(
|
|
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
aws_region_name="us-gov-east-1",
|
|
)
|
|
|
|
# Test gov-west model completion
|
|
gov_west_result = completion(
|
|
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
aws_region_name="us-gov-west-1",
|
|
)
|
|
|
|
# Verify the mock was called correctly
|
|
assert mock_post.call_count == 3
|
|
|
|
# Verify usage information is present
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
assert isinstance(base_result, ModelResponse)
|
|
assert isinstance(gov_east_result, ModelResponse)
|
|
assert isinstance(gov_west_result, ModelResponse)
|
|
|
|
base_result_typed: ModelResponse = base_result
|
|
gov_east_result_typed: ModelResponse = gov_east_result
|
|
gov_west_result_typed: ModelResponse = gov_west_result
|
|
|
|
# Verify usage information is present
|
|
assert (
|
|
hasattr(base_result_typed, "usage")
|
|
and base_result_typed.usage.prompt_tokens == 15
|
|
)
|
|
assert (
|
|
hasattr(base_result_typed, "usage")
|
|
and base_result_typed.usage.completion_tokens == 8
|
|
)
|
|
assert (
|
|
hasattr(gov_east_result_typed, "usage")
|
|
and gov_east_result_typed.usage.prompt_tokens == 15
|
|
)
|
|
assert (
|
|
hasattr(gov_east_result_typed, "usage")
|
|
and gov_east_result_typed.usage.completion_tokens == 8
|
|
)
|
|
assert (
|
|
hasattr(gov_west_result_typed, "usage")
|
|
and gov_west_result_typed.usage.prompt_tokens == 15
|
|
)
|
|
assert (
|
|
hasattr(gov_west_result_typed, "usage")
|
|
and gov_west_result_typed.usage.completion_tokens == 8
|
|
)
|
|
|
|
# Verify cost calculation uses correct pricing for each region
|
|
# Get costs directly from the completion response _hidden_params
|
|
base_cost = base_result_typed._hidden_params.get("response_cost", 0.0)
|
|
gov_east_cost = gov_east_result_typed._hidden_params.get("response_cost", 0.0)
|
|
gov_west_cost = gov_west_result_typed._hidden_params.get("response_cost", 0.0)
|
|
|
|
print(f"Base cost: {base_cost}")
|
|
print(f"Gov East cost: {gov_east_cost}")
|
|
print(f"Gov West cost: {gov_west_cost}")
|
|
|
|
# Expected costs based on pricing:
|
|
# Base model (us.*): 15 * 1.1e-06 + 8 * 5.5e-06 = 1.65e-05 + 4.4e-05 = 6.05e-05
|
|
# Gov models: 15 * 1.2e-06 + 8 * 6e-06 = 1.8e-05 + 4.8e-05 = 6.6e-05
|
|
expected_base_cost = 15 * 1.1e-06 + 8 * 5.5e-06
|
|
expected_gov_cost = 15 * 1.2e-06 + 8 * 6e-06
|
|
|
|
# Verify costs are calculated correctly
|
|
assert (
|
|
abs(base_cost - expected_base_cost) < 1e-10
|
|
), f"Base cost mismatch: got {base_cost}, expected {expected_base_cost}"
|
|
assert (
|
|
abs(gov_east_cost - expected_gov_cost) < 1e-10
|
|
), f"Gov East cost mismatch: got {gov_east_cost}, expected {expected_gov_cost}"
|
|
assert (
|
|
abs(gov_west_cost - expected_gov_cost) < 1e-10
|
|
), f"Gov West cost mismatch: got {gov_west_cost}, expected {expected_gov_cost}"
|
|
|
|
# Verify GovCloud costs are approximately 20% higher than base cost
|
|
assert (
|
|
abs(gov_east_cost / base_cost - 1.2) < 0.15
|
|
), f"Gov East cost should be ~20% higher than base: got {gov_east_cost}, base {base_cost}"
|
|
assert (
|
|
abs(gov_west_cost / base_cost - 1.2) < 0.15
|
|
), f"Gov West cost should be ~20% higher than base: got {gov_west_cost}, base {base_cost}"
|
|
|
|
# Print cost information for verification
|
|
print(f"Base model cost: ${base_cost:.6f}")
|
|
print(f"GovCloud East cost: ${gov_east_cost:.6f}")
|
|
print(f"GovCloud West cost: ${gov_west_cost:.6f}")
|
|
print(f"GovCloud cost increase: {((gov_east_cost / base_cost) - 1) * 100:.1f}%")
|
|
|
|
def test_govcloud_cost_per_token_with_region(self):
|
|
"""Test that cost_per_token function correctly uses region-based pricing for GovCloud models"""
|
|
from litellm import cost_per_token
|
|
from litellm.utils import Usage
|
|
|
|
# Test usage object
|
|
usage = Usage(prompt_tokens=20, completion_tokens=10, total_tokens=30)
|
|
|
|
# Commercial list pricing uses the us.* inference profile id; GovCloud keys use anthropic.* + region
|
|
haiku_us_id = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
haiku_anthropic_id = "anthropic.claude-haiku-4-5-20251001-v1:0"
|
|
# Test base model with standard region
|
|
base_prompt_cost, base_completion_cost = cost_per_token(
|
|
model=haiku_us_id,
|
|
prompt_tokens=20,
|
|
completion_tokens=10,
|
|
custom_llm_provider="bedrock",
|
|
region_name="us-east-1",
|
|
)
|
|
|
|
# Test gov models with gov regions
|
|
gov_east_prompt_cost, gov_east_completion_cost = cost_per_token(
|
|
model=haiku_anthropic_id,
|
|
prompt_tokens=20,
|
|
completion_tokens=10,
|
|
custom_llm_provider="bedrock",
|
|
region_name="us-gov-east-1",
|
|
)
|
|
|
|
gov_west_prompt_cost, gov_west_completion_cost = cost_per_token(
|
|
model=haiku_anthropic_id,
|
|
prompt_tokens=20,
|
|
completion_tokens=10,
|
|
custom_llm_provider="bedrock",
|
|
region_name="us-gov-west-1",
|
|
)
|
|
|
|
# Expected costs:
|
|
# Base model (us.*): 20 * 1.1e-06 + 10 * 5.5e-06 = 2.2e-05 + 5.5e-05 = 7.7e-05
|
|
# Gov models: 20 * 1.2e-06 + 10 * 6e-06 = 2.4e-05 + 6e-05 = 8.4e-05
|
|
expected_base_prompt_cost = 20 * 1.1e-06
|
|
expected_base_completion_cost = 10 * 5.5e-06
|
|
expected_gov_prompt_cost = 20 * 1.2e-06
|
|
expected_gov_completion_cost = 10 * 6e-06
|
|
|
|
# Verify costs are calculated correctly
|
|
assert (
|
|
abs(base_prompt_cost - expected_base_prompt_cost) < 1e-10
|
|
), f"Base prompt cost mismatch: got {base_prompt_cost}, expected {expected_base_prompt_cost}"
|
|
assert (
|
|
abs(base_completion_cost - expected_base_completion_cost) < 1e-10
|
|
), f"Base completion cost mismatch: got {base_completion_cost}, expected {expected_base_completion_cost}"
|
|
|
|
assert (
|
|
abs(gov_east_prompt_cost - expected_gov_prompt_cost) < 1e-10
|
|
), f"Gov East prompt cost mismatch: got {gov_east_prompt_cost}, expected {expected_gov_prompt_cost}"
|
|
assert (
|
|
abs(gov_east_completion_cost - expected_gov_completion_cost) < 1e-10
|
|
), f"Gov East completion cost mismatch: got {gov_east_completion_cost}, expected {expected_gov_completion_cost}"
|
|
|
|
assert (
|
|
abs(gov_west_prompt_cost - expected_gov_prompt_cost) < 1e-10
|
|
), f"Gov West prompt cost mismatch: got {gov_west_prompt_cost}, expected {expected_gov_prompt_cost}"
|
|
assert (
|
|
abs(gov_west_completion_cost - expected_gov_completion_cost) < 1e-10
|
|
), f"Gov West completion cost mismatch: got {gov_west_completion_cost}, expected {expected_gov_completion_cost}"
|
|
|
|
# Verify GovCloud costs are approximately 20% higher than base costs
|
|
# (uses 1e-8 tolerance because GovCloud prices are independently rounded, not exact * 1.2)
|
|
assert (
|
|
abs(gov_east_prompt_cost / base_prompt_cost - 1.2) < 0.15
|
|
), f"Gov East prompt cost should be ~20% higher than base: got {gov_east_prompt_cost}, base {base_prompt_cost}"
|
|
assert (
|
|
abs(gov_east_completion_cost / base_completion_cost - 1.2) < 0.15
|
|
), f"Gov East completion cost should be ~20% higher than base: got {gov_east_completion_cost}, base {base_completion_cost}"
|
|
assert (
|
|
abs(gov_west_prompt_cost / base_prompt_cost - 1.2) < 0.15
|
|
), f"Gov West prompt cost should be ~20% higher than base: got {gov_west_prompt_cost}, base {base_prompt_cost}"
|
|
assert (
|
|
abs(gov_west_completion_cost / base_completion_cost - 1.2) < 0.15
|
|
), f"Gov West completion cost should be ~20% higher than base: got {gov_west_completion_cost}, base {base_completion_cost}"
|
|
|
|
# Test total costs
|
|
base_total_cost = base_prompt_cost + base_completion_cost
|
|
gov_east_total_cost = gov_east_prompt_cost + gov_east_completion_cost
|
|
gov_west_total_cost = gov_west_prompt_cost + gov_west_completion_cost
|
|
|
|
expected_base_total = expected_base_prompt_cost + expected_base_completion_cost
|
|
expected_gov_total = expected_gov_prompt_cost + expected_gov_completion_cost
|
|
|
|
assert (
|
|
abs(base_total_cost - expected_base_total) < 1e-10
|
|
), f"Base total cost mismatch: got {base_total_cost}, expected {expected_base_total}"
|
|
assert (
|
|
abs(gov_east_total_cost - expected_gov_total) < 1e-10
|
|
), f"Gov East total cost mismatch: got {gov_east_total_cost}, expected {expected_gov_total}"
|
|
assert (
|
|
abs(gov_west_total_cost - expected_gov_total) < 1e-10
|
|
), f"Gov West total cost mismatch: got {gov_west_total_cost}, expected {expected_gov_total}"
|
|
assert (
|
|
abs(gov_east_total_cost / base_total_cost - 1.2) < 0.15
|
|
), f"Gov East total cost should be ~20% higher than base: got {gov_east_total_cost}, base {base_total_cost}"
|
|
assert (
|
|
abs(gov_west_total_cost / base_total_cost - 1.2) < 0.15
|
|
), f"Gov West total cost should be ~20% higher than base: got {gov_west_total_cost}, base {base_total_cost}"
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[
|
|
"bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0",
|
|
"bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0",
|
|
"bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0",
|
|
],
|
|
)
|
|
def test_govcloud_converse_models(self, model_name):
|
|
"""Test that GovCloud Claude and Llama models support Converse API"""
|
|
route = BedrockModelInfo.get_bedrock_route(model_name)
|
|
assert route == "converse"
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_name",
|
|
[
|
|
"bedrock/us-gov-east-1/amazon.titan-text-lite-v1",
|
|
"bedrock/us-gov-west-1/amazon.titan-text-express-v1",
|
|
"bedrock/us-gov-east-1/amazon.titan-text-premier-v1:0",
|
|
],
|
|
)
|
|
def test_govcloud_invoke_models(self, model_name):
|
|
"""Test that GovCloud Titan models use Invoke API"""
|
|
route = BedrockModelInfo.get_bedrock_route(model_name)
|
|
assert route == "invoke"
|