scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(tests): make cost tracking and max_tokens tests independent of model_prices JSON

Browse source 
- Gemini video cost tests: pass explicit ModelInfo with pricing data to
  video_generation_cost() instead of relying on get_model_info() lookup,
  which fails when gemini/veo-3.0-generate-preview is not in the JSON.

- Anthropic max_tokens tests: mock get_max_tokens() to return a known
  value instead of depending on model_prices_and_context_window.json
  having a specific max_output_tokens for claude-3-5-sonnet-20241022.

Co-authored-by: yuneng-jiang <yuneng-jiang@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-03-12 04:28:19 +00:00
parent 68ccf6d025
commit 73981d477c
2 changed files with 51 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -1824,16 +1824,19 @@ def test_get_max_tokens_for_model_claude_3():
def test_get_max_tokens_for_model_claude_35():
"""
Test that get_max_tokens_for_model returns correct value for Claude 3.5 models.
Claude 3.5 models have max_output_tokens of 8192.
Test that get_max_tokens_for_model returns the value from get_max_tokens for
Claude 3.5 models.
Fixes: https://github.com/BerriAI/litellm/issues/8835
"""
config = AnthropicConfig()
# Claude 3.5 Sonnet should return 8192
max_tokens = config.get_max_tokens_for_model("claude-3-5-sonnet-20241022")
assert max_tokens == 8192
with patch(
"litellm.llms.anthropic.chat.transformation.get_max_tokens",
return_value=8192,
):
max_tokens = config.get_max_tokens_for_model("claude-3-5-sonnet-20241022")
assert max_tokens == 8192
def test_get_max_tokens_for_model_claude_37():
@ -1911,16 +1914,19 @@ def test_transform_request_uses_dynamic_max_tokens():
messages = [{"role": "user", "content": "Hello"}]
# Claude 3.5 model should get 8192 as default max_tokens
result = config.transform_request(
model="claude-3-5-sonnet-20241022",
messages=messages,
optional_params={}, # No max_tokens provided
litellm_params={},
headers={}
)
with patch(
"litellm.llms.anthropic.chat.transformation.get_max_tokens",
return_value=8192,
):
result = config.transform_request(
model="claude-3-5-sonnet-20241022",
messages=messages,
optional_params={},
litellm_params={},
headers={}
)
assert result["max_tokens"] == 8192
assert result["max_tokens"] == 8192
def test_transform_request_respects_user_max_tokens():

57
tests/test_litellm/llms/gemini/videos/test_gemini_video_transformation.py
View file

@ -593,63 +593,69 @@ class TestGeminiVideoIntegration:
class TestGeminiVideoCostTracking:
"""Test cost tracking for Gemini video generation."""
def test_cost_calculation_with_duration(self):
"""Test that cost is calculated correctly using duration from usage."""
# Test VEO 2.0 ($0.35/second)
from litellm.types.utils import ModelInfo
# Test with explicit model_info to avoid dependency on the cost map JSON
veo2_info = ModelInfo(output_cost_per_second=0.35)
cost_veo2 = video_generation_cost(
model="gemini/veo-2.0-generate-001",
duration_seconds=5.0,
custom_llm_provider="gemini"
custom_llm_provider="gemini",
model_info=veo2_info,
)
expected_veo2 = 0.35 * 5.0 # $1.75
assert abs(cost_veo2 - expected_veo2) < 0.001, f"Expected ${expected_veo2}, got ${cost_veo2}"
# Test VEO 3.0 ($0.75/second)
veo3_info = ModelInfo(output_cost_per_second=0.75)
cost_veo3 = video_generation_cost(
model="gemini/veo-3.0-generate-preview",
duration_seconds=8.0,
custom_llm_provider="gemini"
custom_llm_provider="gemini",
model_info=veo3_info,
)
expected_veo3 = 0.75 * 8.0 # $6.00
assert abs(cost_veo3 - expected_veo3) < 0.001, f"Expected ${expected_veo3}, got ${cost_veo3}"
# Test VEO 3.1 Standard ($0.40/second)
veo31_info = ModelInfo(output_cost_per_second=0.40)
cost_veo31 = video_generation_cost(
model="gemini/veo-3.1-generate-preview",
duration_seconds=10.0,
custom_llm_provider="gemini"
custom_llm_provider="gemini",
model_info=veo31_info,
)
expected_veo31 = 0.40 * 10.0 # $4.00
assert abs(cost_veo31 - expected_veo31) < 0.001, f"Expected ${expected_veo31}, got ${cost_veo31}"
# Test VEO 3.1 Fast ($0.15/second)
veo31_fast_info = ModelInfo(output_cost_per_second=0.15)
cost_veo31_fast = video_generation_cost(
model="gemini/veo-3.1-fast-generate-preview",
duration_seconds=6.0,
custom_llm_provider="gemini"
custom_llm_provider="gemini",
model_info=veo31_fast_info,
)
expected_veo31_fast = 0.15 * 6.0 # $0.90
assert abs(cost_veo31_fast - expected_veo31_fast) < 0.001, f"Expected ${expected_veo31_fast}, got ${cost_veo31_fast}"
def test_cost_calculation_end_to_end(self):
"""Test complete cost tracking flow: request -> response -> cost calculation."""
from litellm.types.utils import ModelInfo
config = GeminiVideoConfig()
mock_logging_obj = Mock()
# Create request with duration
request_data = {
"instances": [{"prompt": "A beautiful sunset"}],
"parameters": {"durationSeconds": 5}
}
# Mock response
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = {
"name": "operations/generate_test123",
}
# Transform response
video_obj = config.transform_video_create_response(
model="gemini/veo-3.0-generate-preview",
raw_response=mock_response,
@ -657,20 +663,19 @@ class TestGeminiVideoCostTracking:
custom_llm_provider="gemini",
request_data=request_data
)
# Verify usage has duration
assert video_obj.usage is not None
assert "duration_seconds" in video_obj.usage
duration = video_obj.usage["duration_seconds"]
# Calculate cost using the duration from usage
veo3_info = ModelInfo(output_cost_per_second=0.75)
cost = video_generation_cost(
model="gemini/veo-3.0-generate-preview",
duration_seconds=duration,
custom_llm_provider="gemini"
custom_llm_provider="gemini",
model_info=veo3_info,
)
# Verify cost calculation (VEO 3.0 is $0.75/second)
expected_cost = 0.75 * 5.0 # $3.75
assert abs(cost - expected_cost) < 0.001, f"Expected ${expected_cost}, got ${cost}"