fix(proxy): log the provider usage on deferred /v1/messages calls and price cache writes without a creation rate

With a post-call guardrail the proxy defers async success logging, and every nested wrapper on a
/v1/messages call bridged to the Responses API overwrote the stored closure, so the spend log was
built from the outermost Anthropic-shaped reply under Responses semantics and recorded the prompt
tokens without the cache hit. The first wrapper to exit now keeps the slot, which is the innermost
provider response, the same one the non-deferred path logs.

The flat cost path also billed cache-creation tokens at 0 when the model had no
cache_creation_input_token_cost. It now falls back to the input rate, and the 1h rate to the
creation rate, matching the tiered path and the custom pricing helper.
This commit is contained in:
mateo-berri 2026-09-14 19:21:03 -07:00
parent d2859e18d7
commit 22b377fe2a
4 changed files with 263 additions and 55 deletions

View file

@ -532,6 +532,11 @@ def _get_token_base_cost(
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
An absent cache-creation rate always resolves to the resolved input rate, the way the
tiered table and custom deployment pricing already do, since a provider that publishes
no write price bills cache writes as ordinary input. An absent 1h write rate resolves
to the cache-creation rate. An explicit 0.0 stays a real price for both.
Returns:
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
"""
@ -554,10 +559,9 @@ def _get_token_base_cost(
output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
if output_image_cost is not None:
completion_base_cost = cast(float, output_image_cost)
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key))
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None)
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, "cache_creation_input_token_cost_above_1hr", default_value=None
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
@ -639,22 +643,10 @@ def _get_token_base_cost(
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
)
cache_creation_cost = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_tiered_key,
cache_creation_cost,
),
)
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost)
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_1hr_tiered_key,
cache_creation_cost_above_1hr,
),
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
@ -665,16 +657,19 @@ def _get_token_base_cost(
except Exception:
continue
input_rate_for_missing_cache_rates: Final = _off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if cache_read_cost is None:
cache_read_cost = (
_off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if missing_cache_read_uses_input
else 0.0
)
cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0
resolved_cache_creation_cost: Final = (
input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost
)
resolved_cache_creation_cost_above_1hr: Final = (
resolved_cache_creation_cost if cache_creation_cost_above_1hr is None else cache_creation_cost_above_1hr
)
return _apply_off_peak_to_base_costs(
model_info,
@ -682,8 +677,8 @@ def _get_token_base_cost(
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
resolved_cache_creation_cost,
resolved_cache_creation_cost_above_1hr,
cache_read_cost,
),
)

View file

@ -1208,30 +1208,13 @@ def _dispatch_success_logging(
is_litellm_internal_call: bool,
) -> None:
if not is_litellm_internal_call:
if getattr(logging_obj, "_defer_async_logging", False):
def _enqueue_deferred_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging
else:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
_schedule_async_success_logging(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
@ -1240,6 +1223,43 @@ def _dispatch_success_logging(
)
def _schedule_async_success_logging(
logging_obj: LiteLLMLoggingObject,
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
is_completion_with_fallbacks: bool,
) -> None:
"""Fire the async success log for ``result`` now, or park it on the logging object while
the proxy defers logging past its post-call guardrails.
Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses
bridge) each exit through here with the same logging object and their own shape of the same
response. The immediate path already logs one request once, since the first task marks
``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same
first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log
reads usage from, and a later wrapper never swaps in its client-shaped translation.
"""
def _enqueue_async_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
if not getattr(logging_obj, "_defer_async_logging", False):
_enqueue_async_logging()
return
if getattr(logging_obj, "_enqueue_deferred_logging", None) is not None:
return
logging_obj._enqueue_deferred_logging = _enqueue_async_logging
async def _client_async_logging_helper(
logging_obj: LiteLLMLoggingObject,
result,

View file

@ -4039,7 +4039,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp
cache_read_input_token_cost=6e-7,
cache_read_input_audio_token_cost=6e-7,
cache_creation_input_token_cost=7.5e-6,
cache_creation_input_token_cost_above_1hr=0.0,
cache_creation_input_token_cost_above_1hr=7.5e-6,
output_cost_per_reasoning_token=3e-5,
)
assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost)
@ -5334,3 +5334,56 @@ def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates(
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider)
assert prompt_cost == pytest.approx(expected_prompt_cost)
def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a_write_price():
"""Azure and OpenAI publish no cache-write price and bill cache writes as ordinary input.
A deployment priced with only input, output, and cache-read rates must bill the creation
tokens the provider reports at the input rate, never at 0. The numbers are a cold 7,336-token
prompt on a deployment that reports all but 3 of them as cache creation."""
model_info = {
"input_cost_per_token": 2e-7,
"output_cost_per_token": 1.25e-6,
"cache_read_input_token_cost": 2e-8,
}
usage = Usage(
prompt_tokens=7336,
completion_tokens=23,
total_tokens=7359,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_creation_tokens=7333),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="custom-priced-deployment", usage=usage, custom_llm_provider="azure", model_info=model_info
)
assert prompt_cost == pytest.approx(7336 * 2e-7)
assert completion_cost == pytest.approx(23 * 1.25e-6)
@pytest.mark.parametrize(
("cache_rates", "current_time", "expected_creation", "expected_creation_1h"),
(
pytest.param({}, None, 2e-7, 2e-7, id="no-write-price-uses-the-input-rate"),
pytest.param({"cache_creation_input_token_cost": 2.5e-7}, None, 2.5e-7, 2.5e-7, id="no-1h-price-uses-the-write-price"),
pytest.param({"cache_creation_input_token_cost": 0.0}, None, 0.0, 0.0, id="explicit-zero-stays-zero"),
pytest.param(
{"off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 1e-7}},
datetime(2026, 9, 14, 12, tzinfo=timezone.utc),
1e-7,
1e-7,
id="no-write-price-uses-the-off-peak-input-rate",
),
),
)
def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path(
cache_rates: dict, current_time: datetime | None, expected_creation: float, expected_creation_1h: float
) -> None:
model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates}
usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)
_, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time)
assert creation == pytest.approx(expected_creation)
assert creation_1h == pytest.approx(expected_creation_1h)

View file

@ -16,14 +16,21 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end.
import asyncio
import logging
from collections.abc import Callable
from datetime import datetime
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
import litellm
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.utils import StandardLoggingPayload
from litellm.utils import _dispatch_success_logging
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
@ -54,6 +61,25 @@ def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn):
mock_logging_obj.async_success_handler = async_success_fn
async def _wait_until(condition: Callable[[], bool]) -> None:
"""Give the logging worker a bounded window to run what the closure enqueued."""
for _ in range(200):
if condition():
return
await asyncio.sleep(0.01)
class _RecordingLogger(CustomLogger):
"""Keeps what the async success callback was handed, the way a spend logger sees it."""
def __init__(self) -> None:
super().__init__()
self.standard_logging_object: StandardLoggingPayload | None = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.standard_logging_object = kwargs["standard_logging_object"]
class PostCallGuardrail(CustomGuardrail):
"""A post-call guardrail."""
@ -259,6 +285,120 @@ async def test_deferred_flag_stores_and_executes_closure():
pass
@pytest.mark.asyncio
async def test_deferred_slot_keeps_the_innermost_wrapper_result():
"""Nested @client wrappers exit through _dispatch_success_logging with one shared logging
object. The deferred slot must keep the first stored result, the way the immediate path's
has_logged dedupe keeps the first fired task, so the spend log reads usage from the
innermost provider-shaped response and never from an outer wrapper's translation of it."""
logging_obj: Final = MagicMock()
logging_obj._defer_async_logging = True
logging_obj._enqueue_deferred_logging = None
logging_obj.async_success_handler = AsyncMock()
inner_result: Final = object()
outer_result: Final = object()
for result in (inner_result, outer_result):
_dispatch_success_logging(
logging_obj=logging_obj,
result=result,
start_time=datetime.now(),
end_time=datetime.now(),
is_completion_with_fallbacks=False,
is_litellm_internal_call=False,
)
logging_obj._enqueue_deferred_logging()
await _wait_until(lambda: logging_obj.async_success_handler.await_count > 0)
logging_obj.async_success_handler.assert_awaited_once()
assert logging_obj.async_success_handler.await_args.kwargs["result"] is inner_result
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_count == 2
@pytest.mark.asyncio
async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the_provider_usage(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
"""/v1/messages on an Azure gpt-5.4+ deployment with function tools runs three nested
wrappers: anthropic_messages, the chat adapter's acompletion, and the Responses bridge
acompletion hands the call to, which retags the call as ``responses``. With logging
deferred for a post-call guardrail the stored closure must carry the innermost provider
response: logging the Anthropic-shaped reply under Responses semantics books this
7,336-token prompt as 3 tokens, since Anthropic's input_tokens excludes the cache hit."""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
litellm.in_memory_llm_clients_cache.flush_cache()
respx_mock.post(url__regex=r"https://deferred-nested\.openai\.azure\.com/openai/.*responses.*").mock(
return_value=httpx.Response(
200,
json={
"id": "resp_deferred_nested",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.4-nano",
"output": [
{
"type": "message",
"id": "msg_deferred_nested",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "Hello!", "annotations": []}],
}
],
"usage": {
"input_tokens": 7336,
"input_tokens_details": {"cached_tokens": 7333},
"output_tokens": 23,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 7359,
},
},
)
)
recorder: Final = _RecordingLogger()
logging_obj: Final = Logging(
model="azure/gpt-5.4-nano",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="deferred-nested-anthropic-messages",
function_id="deferred-nested-anthropic-messages",
dynamic_async_success_callbacks=[recorder],
)
logging_obj._defer_async_logging = True
response: Final = await litellm.anthropic_messages(
model="azure/gpt-5.4-nano",
messages=[{"role": "user", "content": "hi"}],
max_tokens=16,
tools=[
{
"name": "lookup_volume",
"description": "Look up a storage volume by name",
"input_schema": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]},
}
],
api_key="sk-deferred-nested",
api_base="https://deferred-nested.openai.azure.com",
api_version="2025-04-01-preview",
litellm_logging_obj=logging_obj,
)
assert response["content"] == [{"type": "text", "text": "Hello!"}]
assert response["usage"]["input_tokens"] == 3
assert response["usage"]["cache_read_input_tokens"] == 7333
logging_obj._enqueue_deferred_logging()
await _wait_until(lambda: recorder.standard_logging_object is not None)
assert recorder.standard_logging_object is not None
assert recorder.standard_logging_object["prompt_tokens"] == 7336
assert recorder.standard_logging_object["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 7333
assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6)
# ---------------------------------------------------------------------------
# 3. Non-streaming regression: without flag, create_task fires normally
# ---------------------------------------------------------------------------