mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix estimate_cost
This commit is contained in:
Ishaan Jaffer
parent
2e5e01fac3
commit
d402d5c618
2 changed files with 71 additions and 63 deletions
|
|
@ -372,7 +372,7 @@ async def estimate_cost(
|
|||
including any configured margins and discounts.
|
||||
|
||||
Parameters:
|
||||
- model: Model name from /model_group/info (e.g., "gpt-4", "claude-3-opus")
|
||||
- model: Model name (e.g., "gpt-4", "claude-3-opus")
|
||||
- input_tokens: Expected input tokens per request
|
||||
- output_tokens: Expected output tokens per request
|
||||
- num_requests: Number of requests (default: 1)
|
||||
|
|
@ -393,55 +393,66 @@ async def estimate_cost(
|
|||
}
|
||||
```
|
||||
"""
|
||||
from litellm.cost_calculator import _apply_cost_margin
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.utils import Usage
|
||||
from litellm.utils import ModelResponse
|
||||
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "Router not initialized. No models configured."},
|
||||
# Create a mock response with usage for completion_cost
|
||||
mock_response = ModelResponse(
|
||||
model=request.model,
|
||||
usage=Usage(
|
||||
prompt_tokens=request.input_tokens,
|
||||
completion_tokens=request.output_tokens,
|
||||
total_tokens=request.input_tokens + request.output_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
# Create a logging object to capture cost breakdown
|
||||
litellm_logging_obj = LiteLLMLoggingObj(
|
||||
model=request.model,
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=None,
|
||||
litellm_call_id="cost-estimate",
|
||||
function_id="cost-estimate",
|
||||
)
|
||||
|
||||
# Use completion_cost which handles all the logic including margins/discounts
|
||||
try:
|
||||
total_cost_per_request = completion_cost(
|
||||
completion_response=mock_response,
|
||||
model=request.model,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
# Get model group info from router to resolve pricing
|
||||
model_group_info = llm_router.get_model_group_info(model_group=request.model)
|
||||
|
||||
if model_group_info is None:
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={
|
||||
"error": f"Model '{request.model}' not found. Use /model_group/info to see available models."
|
||||
"error": f"Could not calculate cost for model '{request.model}': {str(e)}"
|
||||
},
|
||||
)
|
||||
|
||||
# Get the provider from the model group
|
||||
providers: List[str] = model_group_info.providers or []
|
||||
custom_llm_provider: Optional[str] = providers[0] if providers else None
|
||||
# Get cost breakdown from the logging object
|
||||
cost_breakdown = litellm_logging_obj.cost_breakdown
|
||||
|
||||
# Get cost per token from model group info
|
||||
input_cost_per_token = model_group_info.input_cost_per_token or 0.0
|
||||
output_cost_per_token = model_group_info.output_cost_per_token or 0.0
|
||||
input_cost = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0
|
||||
output_cost = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0
|
||||
margin_cost = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0
|
||||
|
||||
# Calculate base costs (before margin)
|
||||
input_cost = input_cost_per_token * request.input_tokens
|
||||
output_cost = output_cost_per_token * request.output_tokens
|
||||
base_cost = input_cost + output_cost
|
||||
|
||||
# Apply margin using the same function as completion_cost
|
||||
(
|
||||
cost_with_margin,
|
||||
_margin_percent,
|
||||
_margin_fixed_amount,
|
||||
margin_cost,
|
||||
) = _apply_cost_margin(
|
||||
base_cost=base_cost,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
cost_per_request = cost_with_margin
|
||||
# Get model info for per-token pricing display
|
||||
try:
|
||||
model_info = litellm.get_model_info(model=request.model)
|
||||
input_cost_per_token = model_info.get("input_cost_per_token")
|
||||
output_cost_per_token = model_info.get("output_cost_per_token")
|
||||
custom_llm_provider = model_info.get("litellm_provider")
|
||||
except Exception:
|
||||
input_cost_per_token = None
|
||||
output_cost_per_token = None
|
||||
custom_llm_provider = None
|
||||
|
||||
# Calculate totals based on number of requests
|
||||
total_cost = cost_per_request * request.num_requests
|
||||
total_cost = total_cost_per_request * request.num_requests
|
||||
total_input_cost = input_cost * request.num_requests
|
||||
total_output_cost = output_cost * request.num_requests
|
||||
total_margin_cost = margin_cost * request.num_requests
|
||||
|
|
@ -451,7 +462,7 @@ async def estimate_cost(
|
|||
input_tokens=request.input_tokens,
|
||||
output_tokens=request.output_tokens,
|
||||
num_requests=request.num_requests,
|
||||
cost_per_request=cost_per_request,
|
||||
cost_per_request=total_cost_per_request,
|
||||
input_cost_per_request=input_cost,
|
||||
output_cost_per_request=output_cost,
|
||||
margin_cost_per_request=margin_cost,
|
||||
|
|
|
|||
|
|
@ -14,18 +14,10 @@ class TestCostEstimateEndpoint:
|
|||
"""Tests for the cost estimation endpoint."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_estimate_cost_with_margin(self):
|
||||
async def test_estimate_cost_uses_completion_cost(self):
|
||||
"""
|
||||
Test that cost estimation returns correct breakdown including margin.
|
||||
Test that cost estimation uses completion_cost and returns breakdown.
|
||||
"""
|
||||
mock_model_group_info = MagicMock()
|
||||
mock_model_group_info.providers = ["openai"]
|
||||
mock_model_group_info.input_cost_per_token = 0.00003
|
||||
mock_model_group_info.output_cost_per_token = 0.00006
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_group_info.return_value = mock_model_group_info
|
||||
|
||||
request = CostEstimateRequest(
|
||||
model="gpt-4",
|
||||
input_tokens=1000,
|
||||
|
|
@ -33,40 +25,45 @@ class TestCostEstimateEndpoint:
|
|||
num_requests=10,
|
||||
)
|
||||
|
||||
# Base cost = 0.00003 * 1000 + 0.00006 * 500 = 0.03 + 0.03 = 0.06
|
||||
# With 10% margin = 0.06 + 0.006 = 0.066
|
||||
with patch("litellm.cost_calculator._apply_cost_margin") as mock_apply_margin:
|
||||
mock_apply_margin.return_value = (0.066, 0.10, 0.0, 0.006)
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.cost_tracking_settings.completion_cost"
|
||||
) as mock_completion_cost:
|
||||
mock_completion_cost.return_value = 0.066
|
||||
|
||||
with patch("litellm.get_model_info") as mock_get_model_info:
|
||||
mock_get_model_info.return_value = {
|
||||
"input_cost_per_token": 0.00003,
|
||||
"output_cost_per_token": 0.00006,
|
||||
"litellm_provider": "openai",
|
||||
}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
|
||||
response = await estimate_cost(
|
||||
request=request,
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
assert response.model == "gpt-4"
|
||||
assert response.input_cost_per_request == pytest.approx(0.03)
|
||||
assert response.output_cost_per_request == pytest.approx(0.03)
|
||||
assert response.margin_cost_per_request == pytest.approx(0.006)
|
||||
assert response.cost_per_request == pytest.approx(0.066)
|
||||
assert response.cost_per_request == 0.066
|
||||
assert response.total_cost == pytest.approx(0.66)
|
||||
assert response.total_margin_cost == pytest.approx(0.06)
|
||||
# Verify completion_cost was called
|
||||
mock_completion_cost.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_estimate_cost_model_not_found(self):
|
||||
"""
|
||||
Test that 404 is raised when model is not found.
|
||||
Test that 404 is raised when model cost calculation fails.
|
||||
"""
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_group_info.return_value = None
|
||||
|
||||
request = CostEstimateRequest(
|
||||
model="nonexistent-model",
|
||||
input_tokens=1000,
|
||||
output_tokens=500,
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.cost_tracking_settings.completion_cost"
|
||||
) as mock_completion_cost:
|
||||
mock_completion_cost.side_effect = Exception("Model not found in cost map")
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue