mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
329 lines
13 KiB
Python
329 lines
13 KiB
Python
"""
|
|
Integration tests for Perplexity cost calculation and transformation.
|
|
|
|
Tests the end-to-end functionality of Perplexity cost calculation
|
|
including integration with the main LiteLLM cost calculator.
|
|
"""
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
|
|
import pytest
|
|
|
|
# Add the project root to Python path
|
|
|
|
import litellm
|
|
from litellm import ModelResponse
|
|
from litellm.cost_calculator import completion_cost, cost_per_token
|
|
from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig
|
|
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
|
|
from litellm.utils import get_model_info
|
|
|
|
|
|
class TestPerplexityIntegration:
|
|
"""Integration test suite for Perplexity functionality."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup_model_cost_map(self):
|
|
"""Set up the model cost map for testing."""
|
|
# Ensure we use local model cost map for consistent testing
|
|
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
|
|
|
# Load the model cost map
|
|
try:
|
|
with open("model_prices_and_context_window.json", "r") as f:
|
|
model_cost_map = json.load(f)
|
|
litellm.model_cost = model_cost_map
|
|
except FileNotFoundError:
|
|
# Fallback to ensure we have the Perplexity model configuration
|
|
litellm.model_cost = {
|
|
"perplexity/sonar-deep-research": {
|
|
"max_tokens": 128000,
|
|
"max_input_tokens": 128000,
|
|
"input_cost_per_token": 2e-06,
|
|
"output_cost_per_token": 8e-06,
|
|
"output_cost_per_reasoning_token": 3e-06,
|
|
"citation_cost_per_token": 2e-06,
|
|
"search_context_cost_per_query": {
|
|
"search_context_size_low": 0.005,
|
|
"search_context_size_medium": 0.005,
|
|
"search_context_size_high": 0.005,
|
|
},
|
|
"litellm_provider": "perplexity",
|
|
"mode": "chat",
|
|
"supports_reasoning": True,
|
|
"supports_web_search": True,
|
|
}
|
|
}
|
|
|
|
def test_end_to_end_cost_calculation_with_transformation(self):
|
|
"""Test end-to-end cost calculation with response transformation."""
|
|
# Create a Perplexity API response that includes citations and search queries
|
|
config = PerplexityChatConfig()
|
|
|
|
# Create a ModelResponse with basic usage (before transformation)
|
|
model_response = ModelResponse()
|
|
model_response.model = "sonar-deep-research"
|
|
model_response.usage = Usage(
|
|
prompt_tokens=100,
|
|
completion_tokens=50,
|
|
total_tokens=150,
|
|
reasoning_tokens=10,
|
|
)
|
|
|
|
# Simulate raw response from Perplexity API
|
|
raw_response_dict = {
|
|
"choices": [{"message": {"content": "Test response with citations"}}],
|
|
"usage": {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"total_tokens": 150,
|
|
"num_search_queries": 2,
|
|
},
|
|
"citations": [
|
|
"This is the first citation with important information about the topic",
|
|
"Another citation providing additional context for the response",
|
|
],
|
|
}
|
|
|
|
# Apply transformation to extract Perplexity-specific fields
|
|
config._enhance_usage_with_perplexity_fields(model_response, raw_response_dict)
|
|
|
|
# Now calculate the cost with the enhanced usage
|
|
total_cost = completion_cost(
|
|
completion_response=model_response, custom_llm_provider="perplexity"
|
|
)
|
|
|
|
# Calculate expected cost
|
|
citation_chars = sum(
|
|
len(citation) for citation in raw_response_dict["citations"]
|
|
)
|
|
citation_tokens = citation_chars // 4
|
|
|
|
expected_prompt_cost = (100 * 2e-6) + (citation_tokens * 2e-6)
|
|
expected_completion_cost = (
|
|
((50 - 10) * 8e-6) + (10 * 3e-6) + (2 * 0.005)
|
|
) # Output (text) + reasoning + search
|
|
expected_total = expected_prompt_cost + expected_completion_cost
|
|
|
|
assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
|
|
|
|
def test_cost_calculation_without_custom_fields(self):
|
|
"""Test that cost calculation works normally when custom fields are absent."""
|
|
# Create a standard response without Perplexity-specific fields
|
|
model_response = ModelResponse()
|
|
model_response.model = "sonar-deep-research"
|
|
model_response.usage = Usage(
|
|
prompt_tokens=100, completion_tokens=50, total_tokens=150
|
|
)
|
|
|
|
# Calculate cost without custom fields
|
|
total_cost = completion_cost(
|
|
completion_response=model_response, custom_llm_provider="perplexity"
|
|
)
|
|
|
|
# Should only include basic input/output costs
|
|
expected_cost = (100 * 2e-6) + (50 * 8e-6)
|
|
|
|
assert math.isclose(total_cost, expected_cost, rel_tol=1e-6)
|
|
|
|
def test_main_cost_calculator_integration(self):
|
|
"""Test integration with the main LiteLLM cost calculator."""
|
|
# Create usage with all Perplexity fields
|
|
usage = Usage(
|
|
prompt_tokens=200,
|
|
completion_tokens=100,
|
|
total_tokens=300,
|
|
reasoning_tokens=25,
|
|
prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=3),
|
|
)
|
|
usage.citation_tokens = 40
|
|
|
|
# Test main cost calculator
|
|
prompt_cost, completion_cost_val = cost_per_token(
|
|
model="sonar-deep-research",
|
|
custom_llm_provider="perplexity",
|
|
usage_object=usage,
|
|
)
|
|
|
|
expected_prompt_cost = (200 * 2e-6) + (40 * 2e-6)
|
|
expected_completion_cost = (
|
|
((100 - 25) * 8e-6) + (25 * 3e-6) + (3 * 0.005)
|
|
) # Output (text) + reasoning + search
|
|
|
|
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
|
|
assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)
|
|
|
|
def test_model_info_includes_custom_fields(self):
|
|
"""Test that get_model_info returns the custom Perplexity cost fields."""
|
|
model_info = get_model_info(
|
|
model="sonar-deep-research", custom_llm_provider="perplexity"
|
|
)
|
|
|
|
# Verify custom fields are included
|
|
required_fields = [
|
|
"citation_cost_per_token",
|
|
"search_context_cost_per_query",
|
|
"input_cost_per_token",
|
|
"output_cost_per_token",
|
|
"output_cost_per_reasoning_token",
|
|
]
|
|
|
|
for field in required_fields:
|
|
assert field in model_info, f"Missing field: {field}"
|
|
assert model_info[field] is not None, f"Null value for field: {field}"
|
|
|
|
def test_various_citation_sizes(self):
|
|
"""Test cost calculation with various citation sizes."""
|
|
config = PerplexityChatConfig()
|
|
|
|
test_cases = [
|
|
# (citations, expected_approximate_tokens)
|
|
(["Short"], 1),
|
|
(["This is a medium-length citation with some content"], 12),
|
|
(
|
|
["Very short", "Another citation", "Third one with more text content"],
|
|
15,
|
|
),
|
|
([""], 0), # Empty citation
|
|
]
|
|
|
|
for citations, expected_approx_tokens in test_cases:
|
|
model_response = ModelResponse()
|
|
model_response.model = "sonar-deep-research"
|
|
model_response.usage = Usage(
|
|
prompt_tokens=100, completion_tokens=50, total_tokens=150
|
|
)
|
|
|
|
raw_response_dict = {
|
|
"usage": {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"total_tokens": 150,
|
|
},
|
|
"citations": citations,
|
|
}
|
|
|
|
config._enhance_usage_with_perplexity_fields(
|
|
model_response, raw_response_dict
|
|
)
|
|
|
|
citation_tokens = getattr(model_response.usage, "citation_tokens", 0)
|
|
|
|
# Allow for reasonable variance in token estimation
|
|
if expected_approx_tokens == 0:
|
|
assert citation_tokens == 0
|
|
else:
|
|
assert abs(citation_tokens - expected_approx_tokens) <= 5
|
|
|
|
def test_cost_calculation_with_zero_values(self):
|
|
"""Test cost calculation handles zero values for custom fields correctly."""
|
|
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
|
|
|
|
# Set custom fields to zero
|
|
usage.citation_tokens = 0
|
|
usage.prompt_tokens_details = PromptTokensDetailsWrapper(web_search_requests=0)
|
|
|
|
# Should not add any extra cost
|
|
prompt_cost, completion_cost_val = cost_per_token(
|
|
model="sonar-deep-research",
|
|
custom_llm_provider="perplexity",
|
|
usage_object=usage,
|
|
)
|
|
|
|
expected_prompt_cost = 100 * 2e-6
|
|
expected_completion_cost = 50 * 8e-6
|
|
|
|
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
|
|
assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)
|
|
|
|
def test_high_volume_cost_calculation(self):
|
|
"""Test cost calculation with high token and query counts."""
|
|
usage = Usage(
|
|
prompt_tokens=50000,
|
|
completion_tokens=25000,
|
|
total_tokens=75000,
|
|
reasoning_tokens=10000,
|
|
)
|
|
|
|
usage.citation_tokens = 5000
|
|
usage.prompt_tokens_details = PromptTokensDetailsWrapper(
|
|
web_search_requests=100
|
|
)
|
|
|
|
total_cost = completion_cost(
|
|
completion_response=ModelResponse(usage=usage, model="sonar-deep-research"),
|
|
custom_llm_provider="perplexity",
|
|
)
|
|
|
|
expected_prompt_cost = (50000 * 2e-6) + (5000 * 2e-6)
|
|
expected_completion_cost = (
|
|
((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 * 0.005)
|
|
) # $0.65
|
|
expected_total = expected_prompt_cost + expected_completion_cost # $0.76
|
|
|
|
assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
|
|
assert total_cost > 0.25
|
|
|
|
def test_transformation_preserves_existing_usage_fields(self):
|
|
"""Test that transformation doesn't overwrite existing standard usage fields."""
|
|
config = PerplexityChatConfig()
|
|
|
|
model_response = ModelResponse()
|
|
model_response.usage = Usage(
|
|
prompt_tokens=100,
|
|
completion_tokens=50,
|
|
total_tokens=150,
|
|
reasoning_tokens=20,
|
|
)
|
|
|
|
# Store original values
|
|
original_prompt_tokens = model_response.usage.prompt_tokens
|
|
original_completion_tokens = model_response.usage.completion_tokens
|
|
original_total_tokens = model_response.usage.total_tokens
|
|
|
|
raw_response_dict = {
|
|
"usage": {
|
|
"prompt_tokens": 999, # Different from original
|
|
"completion_tokens": 999, # Different from original
|
|
"total_tokens": 999, # Different from original
|
|
"num_search_queries": 3,
|
|
},
|
|
"citations": ["Some citation"],
|
|
}
|
|
|
|
config._enhance_usage_with_perplexity_fields(model_response, raw_response_dict)
|
|
|
|
# Original usage fields should be preserved
|
|
assert model_response.usage.prompt_tokens == original_prompt_tokens
|
|
assert model_response.usage.completion_tokens == original_completion_tokens
|
|
assert model_response.usage.total_tokens == original_total_tokens
|
|
|
|
# But custom fields should be added
|
|
assert hasattr(model_response.usage, "prompt_tokens_details")
|
|
assert hasattr(model_response.usage, "citation_tokens")
|
|
assert model_response.usage.prompt_tokens_details.web_search_requests == 3
|
|
|
|
@pytest.mark.parametrize(
|
|
"provider_name", ["perplexity", "PERPLEXITY", "Perplexity"]
|
|
)
|
|
def test_case_insensitive_provider_matching(self, provider_name):
|
|
"""Test that cost calculation works with different case variations of provider name."""
|
|
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
|
|
usage.citation_tokens = 10
|
|
usage.prompt_tokens_details = PromptTokensDetailsWrapper(web_search_requests=1)
|
|
|
|
# Should work regardless of case
|
|
prompt_cost, completion_cost_val = cost_per_token(
|
|
model="sonar-deep-research",
|
|
custom_llm_provider=provider_name.lower(), # Normalize to lowercase
|
|
usage_object=usage,
|
|
)
|
|
|
|
# Should calculate costs correctly
|
|
expected_prompt_cost = (100 * 2e-6) + (10 * 2e-6)
|
|
expected_completion_cost = (50 * 8e-6) + (1 * 0.005)
|
|
|
|
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
|
|
assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)
|