From 8b73a8afe789ebd51448bc35525e24fd53b52eca Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 15:32:09 -0700 Subject: [PATCH] fix(accounting): preserve image and live adjustment pricing parity --- .../llm_cost_calc/pricing_context.py | 12 ++- litellm/proxy/spend_tracking/cost_compiler.py | 15 ++- litellm/proxy/spend_tracking/usage_capture.py | 10 +- .../test_sagemaker_completion_handler.py | 15 +-- .../spend_tracking/test_cost_compiler.py | 98 +++++++++++++++++++ .../test_register_model_custom_pricing.py | 63 ++++++++++++ 6 files changed, 203 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py index bb575da9bff..9cc9255cfeb 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py +++ b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py @@ -66,11 +66,19 @@ def pricing_discounts() -> Mapping[str, float]: import litellm context: Final = _CURRENT.get() - return litellm.cost_discount_config if context is None else _DISCOUNTS.validate_json(context.discounts_json) + return ( + litellm.cost_discount_config + if context is None or context.allow_live_models + else _DISCOUNTS.validate_json(context.discounts_json) + ) def pricing_margins() -> Mapping[str, float | Mapping[str, float]]: import litellm context: Final = _CURRENT.get() - return litellm.cost_margin_config if context is None else _MARGINS.validate_json(context.margins_json) + return ( + litellm.cost_margin_config + if context is None or context.allow_live_models + else _MARGINS.validate_json(context.margins_json) + ) diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index 2c3bd8754d7..1afa83af83f 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -33,6 +33,7 @@ from litellm.types.utils import ( CallTypes, CostBreakdown, CostPerToken, + ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelInfo, ModelResponse, @@ -93,7 +94,15 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: realtime_usage: Final = event.call_type is CallTypes.arealtime and any( bool(item.get("usage")) for item in realtime_events ) - if captured.usage is None and not fixed_call and not realtime_usage and not speech_usage and not rerank_usage: + image_inputs: Final = _fields(captured.pricing.get("image_response")) + if ( + captured.usage is None + and not fixed_call + and not realtime_usage + and not speech_usage + and not rerank_usage + and not image_inputs + ): return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") pricing: Final = captured.pricing mapped: Final = _fields(pricing.get("model_map_information")) @@ -238,6 +247,10 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: } ) if tool_response.get("responses_api") + else ImageResponse.model_validate( + {"created": int(event.started_at.timestamp()), "model": model, **image_inputs} + ) + if image_inputs else token_response ) logger: Final = Logging( diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index edb7015f9a4..001d0713e56 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -31,7 +31,7 @@ from litellm.proxy.spend_tracking.usage_meter import ( UsageCaptureSettings, usage_delivery_targets, ) -from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams +from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams, ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -1159,6 +1159,14 @@ def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: from litellm.litellm_core_utils.litellm_logging import Logging return { + "image_response": { + "data": tuple(dict[str, JsonValue]() for _ in response.data or ()), + "usage": _response_fields(response.usage) if response.usage is not None else None, + "size": response.size, + "quality": response.quality, + } + if isinstance(response, ImageResponse) + else None, "tool_response": _tool_cost_inputs(response), "realtime_events": _realtime_cost_events(response), "request": _selected_pricing(kwargs), diff --git a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py index 5e40c2ebb60..496eb5fe25b 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py @@ -298,8 +298,13 @@ def test_installed_botocore_signs_the_completion_request(): from litellm.llms.sagemaker.completion.handler import SagemakerLLM request = SagemakerLLM()._prepare_request( - credentials=Credentials("test-key", "test-secret"), model="test-endpoint", data={"inputs": "ping"}, - messages=[], litellm_params={}, optional_params={}, aws_region_name="us-west-2", + credentials=Credentials("test-key", "test-secret"), + model="test-endpoint", + data={"inputs": "ping"}, + messages=[], + litellm_params={}, + optional_params={}, + aws_region_name="us-west-2", ) assert request.body == b'{"inputs": "ping"}' assert "/us-west-2/sagemaker/aws4_request" in request.headers["Authorization"] @@ -447,9 +452,7 @@ def test_completion_sagemaker_prompt_template_non_stream( }, ) ) - url: Final = ( - "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/deepseek_coder_6.7_instruct/invocations" - ) + url: Final = "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/deepseek_coder_6.7_instruct/invocations" route: Final = respx_mock.post(url).mock( return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"}) ) @@ -549,7 +552,7 @@ async def test_completion_sagemaker( assert route.called assert response.choices[0].message.content == "SageMaker reply" - assert litellm.completion_cost(completion_response=response) == pytest.approx(rate) + assert litellm.completion_cost(completion_response=response, custom_cost_per_second=rate) == pytest.approx(rate) @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index 0e9c8f48609..083ac2e0355 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -1575,3 +1575,101 @@ async def test_capture_replays_effective_ptu_spillover_prices( assert float(result.total_cost) == pytest.approx(source_cost) finally: await meter.stop(1) + + +@pytest.mark.asyncio +async def test_live_adjustment_reload_is_captured_and_replay_remains_frozen(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.utils import ModelResponse, Usage + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + monkeypatch.setattr(litellm, "cost_discount_config", {}) + monkeypatch.setattr(litellm, "cost_margin_config", {}) + logger: Final = _capture_logger() + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 0.02}}) + response: Final = ModelResponse(model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + amount: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + assert amount == pytest.approx(0.18 * 0.9 + 0.02) + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.close() + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.9}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 99.0}}) + replay: Final = price_usage(database.rows[-1].envelope) + assert isinstance(replay, PricingSnapshot) + assert float(replay.total_cost) == pytest.approx(amount) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_pricing", (True, False)) +async def test_image_capture_replays_the_existing_image_calculator( + monkeypatch: pytest.MonkeyPatch, token_pricing: bool +) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.utils import ImageResponse + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 0.2}}) + model: Final = "gpt-image-capture" if token_pricing else "dall-e-capture" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "openai", + "mode": "image_generation", + **( + {"input_cost_per_token": 0.01, "output_cost_per_image_token": 0.02} + if token_pricing + else {"output_cost_per_image": 0.25} + ), + }, + ) + logger: Final = _logger("aimage_generation") + logger.model = model + logger.update_environment_variables( + model=model, custom_llm_provider="openai", litellm_params={}, optional_params={} + ) + response: Final = ImageResponse.model_validate( + { + "created": 1, + "model": model, + "data": [{"b64_json": "private-image-content"}] * 2, + "size": "1024x1024", + "quality": "standard", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_tokens_details": {"text_tokens": 10, "image_tokens": 0}, + } + if token_pricing + else None, + } + ) + source: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + assert source == pytest.approx(0.2 if token_pricing else 0.5) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.close() + envelope: Final = database.rows[-1].envelope + assert "private-image-content" not in envelope + monkeypatch.delitem(litellm.model_cost, model) + replay: Final = price_usage(envelope) + assert isinstance(replay, PricingSnapshot) + assert float(replay.total_cost) == pytest.approx(source) + finally: + await meter.stop(1) diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 480c2c0f441..5d77492ee58 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -1807,3 +1807,66 @@ def test_changed_request_flat_rate_preserves_explicit_deployment_tier_prices() - ) finally: _restore_model_cost_entries(original) + + +@pytest.mark.parametrize("call_type", ("completion", "embedding")) +def test_sdk_followup_prices_change_only_after_explicit_registration(call_type: str) -> None: + model: Final = f"openai/followup-{call_type}-prices" + original: Final = _snapshot_model_cost_entries((model, model.removeprefix("openai/"))) + models_before: Final = frozenset(litellm.open_ai_chat_completion_models) + logger: Final = _request_logger(model, call_type) + try: + litellm.register_model( + { + model: { + "litellm_provider": "openai", + "mode": "chat" if call_type == "completion" else "embedding", + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + } + }, + persist_across_reloads=False, + ) + + def request(prices: dict[str, float], *, logging_obj: Logging | None = None) -> object: + if call_type == "completion": + return litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + mock_response="hello back", + litellm_logging_obj=logging_obj, + **prices, + ) + return litellm.embedding( + model=model, + input=["hello"], + api_key="fake-key", + mock_response=[0.1, 0.2], + litellm_logging_obj=logging_obj, + **prices, + ) + + first: Final = request({"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, logging_obj=logger) + second: Final = request({}) + assert logger.response_cost_calculator(first, ignore_stamped_cost=True) == pytest.approx( + 0.05 if call_type == "completion" else 0.01 + ) + assert litellm.completion_cost(completion_response=second, model=model, call_type=call_type) == pytest.approx( + 0.005 if call_type == "completion" else 0.001 + ) + assert litellm.model_cost[model]["input_cost_per_token"] == 0.0001 + litellm.register_model( + {model: {"input_cost_per_token": 0.0003, "output_cost_per_token": 0.0006}}, + persist_across_reloads=False, + ) + third: Final = request({}) + assert litellm.completion_cost(completion_response=third, model=model, call_type=call_type) == pytest.approx( + 0.015 if call_type == "completion" else 0.003 + ) + assert logger.response_cost_calculator(first, ignore_stamped_cost=True) == pytest.approx( + 0.05 if call_type == "completion" else 0.01 + ) + finally: + _restore_model_cost_entries(original) + litellm.open_ai_chat_completion_models.intersection_update(models_before)