From d402d5c618829cfdbc2aa0d7d249606ede4b9754 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 5 Jan 2026 16:26:27 +0530 Subject: [PATCH] fix estimate_cost --- .../cost_tracking_settings.py | 87 +++++++++++-------- .../test_cost_estimate_endpoint.py | 47 +++++----- 2 files changed, 71 insertions(+), 63 deletions(-) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index a94b3800314..daa456a233f 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -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, diff --git a/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py b/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py index 440dbb93525..b31423a582f 100644 --- a/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py +++ b/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py @@ -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: