fix(accounting): preserve image and live adjustment pricing parity

This commit is contained in:
Joshua Valluru 2026-10-31 15:32:09 -07:00
parent f1b6559887
commit 8b73a8afe7
6 changed files with 203 additions and 10 deletions

View file

@ -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)
)

View file

@ -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(

View file

@ -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),

View file

@ -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

View file

@ -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)

View file

@ -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)