Merge pull request #38315 from BerriAI/litellm_cost_estimate_pricing_edges

test(cost-estimate): pin the prices and period totals /cost/estimate returns

tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py

@@ -780,3 +780,132 @@ class TestBlockRequestsForModelsWithoutPricing:
 
         assert response.status_code == 500
         assert "error" in response.json()["detail"]
+
+
+AN_ALIAS = "onprem/alias"
+AN_UNDERLYING_MODEL = "vendor/model"
+A_MAPPED_MODEL = "openai/mapped-only-model"
+INPUT_TOKENS = 1000
+OUTPUT_TOKENS = 500
+
+
+def _router_pricing(**pricing: float) -> MagicMock:
+    mock_router = MagicMock()
+    mock_router.get_model_list.return_value = [
+        {
+            "model_name": AN_ALIAS,
+            "litellm_params": {
+                "model": AN_UNDERLYING_MODEL,
+                "custom_llm_provider": "openai",
+                **pricing,
+            },
+            "model_info": {},
+        }
+    ]
+    return mock_router
+
+
+async def _estimate(mock_router: MagicMock | None, model: str = AN_ALIAS, **overrides: int):
+    from litellm.proxy._types import CostEstimateRequest
+    from litellm.proxy.management_endpoints.cost_tracking_settings import estimate_cost
+
+    request = CostEstimateRequest(
+        model=model,
+        input_tokens=INPUT_TOKENS,
+        output_tokens=OUTPUT_TOKENS,
+        **overrides,
+    )
+    with patch(  # test-quality-ok: proxy_server module global is the endpoint's only injection point
+        "litellm.proxy.proxy_server.llm_router", mock_router
+    ):
+        return await estimate_cost(request=request, user_api_key_dict=MagicMock())
+
+
+class TestEstimateCostPartiallyPricedDeployments:
+    @pytest.mark.asyncio
+    async def test_a_deployment_that_prices_only_input_bills_output_at_zero(self):
+        response = await _estimate(_router_pricing(input_cost_per_token=0.000001))
+
+        assert response.input_cost_per_token == pytest.approx(0.000001)
+        assert response.output_cost_per_token == 0.0
+        assert response.cost_per_request == pytest.approx(0.001)
+
+    @pytest.mark.asyncio
+    async def test_a_deployment_that_prices_only_output_bills_input_at_zero(self):
+        response = await _estimate(_router_pricing(output_cost_per_token=0.000002))
+
+        assert response.input_cost_per_token == 0.0
+        assert response.output_cost_per_token == pytest.approx(0.000002)
+        assert response.cost_per_request == pytest.approx(0.001)
+
+    @pytest.mark.asyncio
+    async def test_a_model_priced_only_by_the_cost_map_reports_that_price_and_provider(self, monkeypatch):
+        monkeypatch.setitem(
+            litellm.model_cost,
+            A_MAPPED_MODEL,
+            {
+                "input_cost_per_token": 0.000005,
+                "output_cost_per_token": 0.000006,
+                "litellm_provider": "openai",
+                "mode": "chat",
+            },
+        )
+
+        response = await _estimate(None, model=A_MAPPED_MODEL)
+
+        assert response.input_cost_per_token == pytest.approx(0.000005)
+        assert response.output_cost_per_token == pytest.approx(0.000006)
+        assert response.provider == "openai"
+
+
+class TestEstimateCostPeriodTotals:
+    @pytest.mark.asyncio
+    async def test_zero_requests_a_day_reports_no_daily_cost_rather_than_zero(self):
+        response = await _estimate(
+            _router_pricing(input_cost_per_token=0.000001, output_cost_per_token=0.000002),
+            num_requests_per_day=0,
+        )
+
+        assert response.daily_cost is None
+        assert response.daily_input_cost is None
+        assert response.daily_output_cost is None
+
+    @pytest.mark.asyncio
+    async def test_daily_totals_scale_every_component_by_the_request_count(self):
+        response = await _estimate(
+            _router_pricing(input_cost_per_token=0.000001, output_cost_per_token=0.000002),
+            num_requests_per_day=100,
+        )
+
+        assert response.input_cost_per_request == pytest.approx(0.001)
+        assert response.output_cost_per_request == pytest.approx(0.001)
+        assert response.daily_input_cost == pytest.approx(0.1)
+        assert response.daily_output_cost == pytest.approx(0.1)
+        assert response.daily_cost == pytest.approx(0.2)
+
+    @pytest.mark.asyncio
+    async def test_a_month_and_a_day_are_totalled_from_their_own_request_counts(self):
+        response = await _estimate(
+            _router_pricing(input_cost_per_token=0.000001, output_cost_per_token=0.000002),
+            num_requests_per_day=100,
+            num_requests_per_month=3000,
+        )
+
+        assert response.daily_cost == pytest.approx(0.2)
+        assert response.monthly_cost == pytest.approx(6.0)
+        assert response.monthly_input_cost == pytest.approx(3.0)
+        assert response.monthly_output_cost == pytest.approx(3.0)
+
+    @pytest.mark.asyncio
+    async def test_a_configured_margin_is_totalled_per_period_like_the_other_components(self, monkeypatch):
+        monkeypatch.setattr(litellm, "cost_margin_config", {"openai": 0.10})
+
+        response = await _estimate(
+            _router_pricing(input_cost_per_token=0.000001, output_cost_per_token=0.000002),
+            num_requests_per_day=100,
+        )
+
+        assert response.margin_cost_per_request == pytest.approx(0.0002)
+        assert response.cost_per_request == pytest.approx(0.0022)
+        assert response.daily_margin_cost == pytest.approx(0.02)
+        assert response.daily_cost == pytest.approx(0.22)