scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(batches): claim the batch cost spend row in the database before charging

Browse source 
The cost callback used to look for an existing `<batch id>_batch_cost` row before charging a
completed batch, which left a window where concurrent retrieves on any instance all charged the
key, and it would honor a row any request had written under that id. The spend update writer now
inserts the batch cost row itself with `create_many(skip_duplicates=True)` and only the retrieve
whose insert lands charges the key, team, and user. An existing row only takes the charge when it
is a successful `aretrieve_batch` row, so a client-chosen `x-litellm-call-id` on another endpoint
cannot suppress billing. Batch cost rows no longer get their own immediate flush path

`batch_cost_is_final` now treats the proxy's normalized `complete` status like `completed`, which
the enterprise batch cost poller relies on when it decides whether a completed batch is safe to
retire. Tests build that status with `model_copy` since the OpenAI `Batch` model rejects it

The `test-quality-ok` markers sit on the `patch(` lines the gate keys on, and the logging tests no
longer wrap the priced retrieve in `contextlib.suppress`
This commit is contained in:
mateo-berri 2026-09-05 18:58:11 -07:00
parent 4c00a6e189
commit b067e836f8
9 changed files with 249 additions and 192 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/batches/batch_utils.py
View file

@ -25,7 +25,8 @@ class BatchCostUsageResult:
failed_requests: int
_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"})
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"})
def batch_cost_is_final(batch: Batch) -> bool:
@ -39,7 +40,7 @@ def batch_cost_is_final(batch: Batch) -> bool:
"""
if batch.status not in _TERMINAL_BATCH_STATUSES:
return False
if batch.status != "completed" or batch.output_file_id is not None:
if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None:
return True
request_counts: Final = batch.request_counts
return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0

73
litellm/proxy/db/db_spend_update_writer.py
View file

@ -82,7 +82,10 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value})
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
class _SpendBatch(Protocol):
@ -216,7 +219,12 @@ class DBSpendUpdateWriter:
start_time: datetime | None,
end_time: datetime | None,
response_cost: float | None,
) -> None:
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
False only for a batch retrieve whose cost row another retrieve already wrote,
so the caller leaves the key, team, and user counters alone (LIT-7048).
"""
from litellm.proxy.proxy_server import (
disable_spend_logs,
litellm_proxy_budget_name,
@ -233,7 +241,7 @@ class DBSpendUpdateWriter:
team_id,
)
if ProxyUpdateSpend.disable_spend_updates() is True:
return
return True
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
@ -264,10 +272,8 @@ class DBSpendUpdateWriter:
payload["team_id"] = team_id
if disable_spend_logs is False:
await self._insert_spend_log_to_db(
payload=payload,
prisma_client=prisma_client,
)
if not await self._record_spend_log(payload=payload, prisma_client=prisma_client):
return False
await self._enqueue_tool_usage_transaction(
payload=payload,
completion_response=completion_response,
@ -307,6 +313,7 @@ class DBSpendUpdateWriter:
)
verbose_proxy_logger.debug("Runs spend update on all tables")
return True
except Exception:
spend_log_error(
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
@ -319,7 +326,55 @@ class DBSpendUpdateWriter:
org_id,
end_user_id,
)
return
return True
async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool:
if prisma_client is None or not _is_batch_cost_row(payload):
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
return True
return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client)
async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool:
"""Write the batch's cost row now, or learn that another retrieve already did.
Every retrieve of one batch shares this row, so the insert that lands first owns
the charge and every later one finds the row and charges nothing (LIT-7048). Only
a row a successful retrieve wrote counts: a failed retrieve, or any request whose
client picked the batch id as its call id, cannot take the charge away.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
request_id: Final = payload["request_id"]
spend_logs: Final = SpendLogsRepository(prisma_client).table
try:
claimed: Final = await spend_logs.create_many(
data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list
skip_duplicates=True,
)
if claimed == 1:
return True
existing: Final = await spend_logs.find_unique(
where={"request_id": request_id} # mutable-ok: prisma where clause
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
verbose_proxy_logger.warning(
"Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e
)
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
return True
if (
existing is not None
and existing.call_type == CallTypes.aretrieve_batch.value
and existing.status == "success"
):
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
return False
verbose_proxy_logger.warning(
"Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own",
request_id,
getattr(existing, "call_type", None),
)
return True
async def _enqueue_tool_usage_transaction(
self,
@ -940,7 +995,7 @@ class DBSpendUpdateWriter:
from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush
await enqueue_spend_logs(prisma_client, (payload,))
if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES:
if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES:
request_spend_log_flush()
else:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")

68
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -34,7 +34,6 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import (
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_sanitize_error_information_for_spend_logs,
get_request_model_access_groups,
get_spend_logs_id,
)
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
@ -46,9 +45,7 @@ from litellm.types.utils import (
from litellm.utils import get_end_user_id_for_cost_tracking
if TYPE_CHECKING:
from prisma.types import LiteLLM_SpendLogsWhereUniqueInput
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.proxy.utils import ProxyLogging
_UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
{
@ -229,7 +226,6 @@ class _ProxyDBLogger(CustomLogger):
):
from litellm.proxy.proxy_server import (
increment_spend_counters,
prisma_client,
proxy_logging_obj,
update_cache,
)
@ -257,15 +253,15 @@ class _ProxyDBLogger(CustomLogger):
if (
isinstance(completion_response, LiteLLMBatch)
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
and not batch_cost_is_final(completion_response)
):
batch_spend_log_id: Final = get_spend_logs_id(
CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs
verbose_proxy_logger.debug(
"Cost tracking deferred for batch %s still in status %s",
completion_response.id,
completion_response.status,
)
if not await _batch_cost_is_trackable_now(
batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client
):
await _release_budget_reservation(budget_reservation=budget_reservation)
return
await _release_budget_reservation(budget_reservation=budget_reservation)
return
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
@ -307,7 +303,7 @@ class _ProxyDBLogger(CustomLogger):
call_type=call_type,
):
## UPDATE DATABASE
await _update_database_and_spend_counters(
charged: Final = await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key=user_api_key,
@ -324,6 +320,8 @@ class _ProxyDBLogger(CustomLogger):
request_tags=tags,
model_access_groups=model_access_groups,
)
if not charged:
return
# update cache (fire-and-forget for backward compat:
# cached object fields, soft budget alerts, etc.)
@ -509,42 +507,6 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None:
bucket[key] = value
async def _batch_cost_is_trackable_now(
batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None"
) -> bool:
"""A batch is billed exactly once, from the first retrieve that sees it final.
Every retrieve of one batch shares a single spend row (its id plus the batch cost
suffix), so a poll that lands before the output exists would write that row at $0
and pin it there, and every retrieve after the first would add the cost to the
key, team, and user counters again.
"""
if not batch_cost_is_final(batch):
verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status)
return False
if prisma_client is None or spend_log_id is None:
return True
if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id):
return True
verbose_proxy_logger.debug(
"Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id
)
return False
async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool:
from litellm.proxy.utils import spend_log_is_queued
if await spend_log_is_queued(prisma_client, request_id):
return True
spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id}
try:
return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row
verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e)
return False
def _is_unbilled_interaction_response(completion_response: object) -> bool:
from litellm.interactions.background_cost_polling import missing_usage_is_expected
from litellm.types.interactions import InteractionsAPIResponse
@ -636,9 +598,9 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
) -> None:
) -> bool:
try:
await proxy_logging_obj.db_spend_update_writer.update_database(
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,
@ -663,6 +625,9 @@ async def _update_database_and_spend_counters(
"Failed to invalidate budget reservation counters after release failed"
)
raise
if not charged:
await _release_budget_reservation(budget_reservation=budget_reservation)
return False
try:
await increment_spend_counters(
@ -688,6 +653,7 @@ async def _update_database_and_spend_counters(
finally:
budget_reservation["finalized"] = True
raise
return True
async def _release_budget_reservation(budget_reservation: dict | None) -> None:

10
litellm/proxy/utils.py
View file

@ -6251,9 +6251,7 @@ def request_spend_log_flush() -> None:
The Responses API hands the client an id it can chain from straight away, and that
lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval.
A batch's cost row is what every other worker checks before charging the same batch
again, so it cannot wait either. Repeated requests coalesce into the monitor's next
pass, so the batching holds.
Repeated requests coalesce into the monitor's next pass, so the batching holds.
"""
PrismaClient.spend_log_flush_requested.set()
@ -6268,12 +6266,6 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool:
return True
async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool:
"""Whether a spend log with ``request_id`` is still waiting for the next flush."""
async with prisma_client._spend_log_transactions_lock:
return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions)
async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]:
"""Take up to ``limit`` of the oldest queued spend logs off the queue.

14
tests/test_litellm/batches/test_batch_utils.py
View file

@ -1735,10 +1735,10 @@ def _retrieved_batch(
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status=status,
status="validating",
output_file_id=output_file_id,
request_counts=counts,
)
).model_copy(update={"status": status})
class TestBatchCostIsFinal:
@ -1750,8 +1750,9 @@ class TestBatchCostIsFinal:
def test_in_flight_batch_is_not_final(self, status):
assert bu.batch_cost_is_final(_retrieved_batch(status)) is False
def test_completed_with_output_is_final(self):
assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True
@pytest.mark.parametrize("status", ["completed", "complete"])
def test_completed_with_output_is_final(self, status):
assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True
def test_completed_without_output_and_unknown_counts_is_not_final(self):
assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False
@ -1764,9 +1765,10 @@ class TestBatchCostIsFinal:
counts = BatchRequestCounts(total=2, completed=2, failed=0)
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
def test_completed_without_output_and_every_line_failed_is_final(self):
@pytest.mark.parametrize("status", ["completed", "complete"])
def test_completed_without_output_and_every_line_failed_is_final(self, status):
counts = BatchRequestCounts(total=2, completed=0, failed=2)
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True
assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True
@pytest.mark.parametrize("status", ["failed", "expired", "cancelled"])
def test_other_terminal_statuses_are_final(self, status):

12
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -665,14 +665,14 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status=status,
status="validating",
output_file_id=output_file_id,
)
).model_copy(update={"status": status})
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status", "output_file_id"),
[("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)],
[("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)],
)
async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None:
from litellm.litellm_core_utils import litellm_logging as logging_module
@ -681,8 +681,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
batch = self._batch(status, output_file_id)
with contextlib.suppress(Exception):
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
handle_completed_batch.assert_not_awaited()
assert "response_cost" not in batch._hidden_params
@ -705,8 +704,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
batch = self._batch("completed", "file-out")
with contextlib.suppress(Exception):
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
handle_completed_batch.assert_awaited_once()
assert batch._hidden_params["response_cost"] == 8e-06

128
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -7,6 +7,7 @@ import re
from collections.abc import Callable
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
@ -2934,16 +2935,14 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type, expects_flush",
[("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)],
[("aresponses", True), ("responses", True), ("acompletion", False)],
)
async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back(
call_type: str, expects_flush: bool
):
"""
A `previous_response_id` chained straight off the previous turn reads the DB, so a
Responses row cannot sit in this worker's queue until the monitor's next poll. A
batch's cost row is what another worker checks before charging the same batch again
(LIT-7048), so it cannot wait either.
Responses row cannot sit in this worker's queue until the monitor's next poll.
"""
from litellm.proxy.utils import PrismaClient
@ -2961,6 +2960,127 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker
PrismaClient.spend_log_flush_requested.clear()
def _batch_cost_payload() -> dict:
return {
**_minimal_spend_payload(),
"request_id": "batch_abc_batch_cost",
"call_type": "aretrieve_batch",
"status": "success",
}
def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock:
prisma = _tool_usage_prisma()
prisma.jsonify_object = lambda data: dict(data)
prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted)
prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing)
return prisma
async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool:
with (
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
"litellm.proxy.proxy_server.disable_spend_logs", False
),
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
),
patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
return_value=payload,
),
):
charged = await db_writer.update_database(
token="test-token",
user_id="test-user",
end_user_id=None,
team_id=None,
org_id=None,
kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"},
completion_response=None,
start_time=datetime.now(timezone.utc),
end_time=datetime.now(timezone.utc),
response_cost=0.25,
)
await asyncio.sleep(0)
return charged
@pytest.mark.asyncio
@pytest.mark.parametrize(
("inserted", "existing", "charged"),
[
(1, None, True),
(0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False),
(0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True),
(0, SimpleNamespace(call_type="aembedding", status="success"), True),
(0, None, True),
],
ids=[
"first_retrieve_owns_the_row",
"another_retrieve_already_charged",
"failed_retrieve_holds_the_row",
"client_chosen_call_id_holds_the_row",
"row_gone_between_insert_and_lookup",
],
)
async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row(
inserted: int, existing: object, charged: bool
):
"""
Every retrieve of one batch shares one spend row, so the insert that lands first is
the charge and every later retrieve must leave the counters alone (LIT-7048). A row
written by anything but a successful retrieve, say a request whose client picked the
batch id as its call id, must not be able to take the charge away.
"""
db_writer = DBSpendUpdateWriter()
db_writer._batch_database_updates = AsyncMock()
prisma = _spend_logs_prisma(inserted, existing)
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged
claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs
assert claimed_rows["skip_duplicates"] is True
assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)]
assert prisma.spend_log_transactions == []
assert db_writer._batch_database_updates.await_count == (1 if charged else 0)
@pytest.mark.asyncio
async def test_update_database_queues_a_batch_cost_row_it_could_not_claim():
"""An unreachable DB must not drop the batch's only spend row, nor its charge."""
db_writer = DBSpendUpdateWriter()
db_writer._batch_database_updates = AsyncMock()
prisma = _spend_logs_prisma(0, None)
prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable"))
assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True
assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"]
assert db_writer._batch_database_updates.await_count == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
"payload",
[{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}],
ids=["not_a_batch_retrieve", "failed_batch_retrieve"],
)
async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict):
db_writer = DBSpendUpdateWriter()
db_writer._batch_database_updates = AsyncMock()
prisma = _spend_logs_prisma(1, None)
assert await _update_database_with(db_writer, prisma, payload) is True
prisma.db.litellm_spendlogs.create_many.assert_not_called()
assert prisma.spend_log_transactions == [payload]
assert db_writer._batch_database_updates.await_count == 1
@pytest.mark.asyncio
@pytest.mark.parametrize(
"injected_deployment, attributed",

116
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -1,3 +1,4 @@
import asyncio
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
@ -6,7 +7,6 @@ import pytest
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.proxy_track_cost_callback import (
_batch_cost_is_trackable_now,
_get_budget_reservation_from_metadata,
_ProxyDBLogger,
_should_track_cost_callback,
@ -783,110 +783,46 @@ def _retrieved_batch(status: str, output_file_id: str | None):
)
def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock:
import asyncio
prisma_client = MagicMock()
prisma_client._spend_log_transactions_lock = asyncio.Lock()
prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids]
prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row)
return prisma_client
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"),
("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"),
[
("in_progress", None, "batch_abc_batch_cost", None, False),
("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False),
("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False),
("completed", "file-out", "batch_abc_batch_cost", None, True),
("completed", "file-out", None, _prisma_client_with((), None), True),
("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False),
(
"completed",
"file-out",
"batch_abc_batch_cost",
_prisma_client_with((), {"request_id": "batch_abc_batch_cost"}),
False,
),
("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True),
("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True),
("aretrieve_batch", "in_progress", None, True, False, False),
("aretrieve_batch", "completed", None, True, False, False),
("aretrieve_batch", "completed", "file-out", False, True, False),
("aretrieve_batch", "completed", "file-out", True, True, True),
("aretrieve_batch", "failed", None, True, True, True),
("acreate_batch", "validating", None, True, True, True),
],
ids=[
"in_progress_without_db",
"in_progress_never_consults_db",
"completed_without_output_yet",
"final_without_db",
"final_without_spend_log_id",
"final_row_queued_for_flush",
"final_row_already_stored",
"final_first_sighting",
"failed_first_sighting",
"retrieve_before_final",
"retrieve_completed_without_output_yet",
"retrieve_after_another_retrieve_charged",
"retrieve_first_final",
"retrieve_failed_batch",
"create_before_final",
],
)
async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable):
"""
A batch is billed from the first retrieve that sees it final and never again:
a poll before that wrote the shared spend row at $0 and pinned it there, and
every completed retrieve after the first charged the key again (LIT-7048).
"""
assert (
await _batch_cost_is_trackable_now(
batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client
)
is trackable
)
@pytest.mark.asyncio
async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails():
"""An unreadable spend log table must not drop the batch's only spend row."""
prisma_client = _prisma_client_with((), None)
prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable"))
assert (
await _batch_cost_is_trackable_now(
batch=_retrieved_batch("completed", "file-out"),
spend_log_id="batch_abc_batch_cost",
prisma_client=prisma_client,
)
is True
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call_type", "status", "output_file_id", "stored_row", "charged"),
[
("aretrieve_batch", "in_progress", None, None, False),
("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False),
("aretrieve_batch", "completed", "file-out", None, True),
("acreate_batch", "validating", None, None, True),
],
ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"],
)
async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate
call_type, status, output_file_id, stored_row, charged
async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate
call_type, status, output_file_id, row_claimed, spend_written, charged
):
"""
Only retrieves are gated, since creating a batch is its own billable request.
A retrieve that writes nothing hands its budget reservation back instead.
A poll before the batch is final used to pin its shared spend row at $0, and every
completed retrieve after the first charged the key again (LIT-7048). Only retrieves
are gated, since creating a batch is its own billable request, and a retrieve that
charges nothing hands its budget reservation back instead.
"""
logger = _ProxyDBLogger()
budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []}
kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation)
with (
patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row)
),
patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
),
) as mock_increment_spend_counters,
patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock
),
) as mock_update_cache,
patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam
"litellm.proxy.proxy_server.proxy_logging_obj"
) as mock_proxy_logging,
@ -895,7 +831,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( #
) as mock_release_budget_reservation,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed)
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
await logger._PROXY_track_cost_callback(
@ -904,13 +840,15 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( #
start_time=datetime.now(),
end_time=datetime.now(),
)
await asyncio.sleep(0)
mock_proxy_logging.failed_tracking_alert.assert_not_called()
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0)
assert mock_increment_spend_counters.await_count == (1 if charged else 0)
assert mock_update_cache.await_count == (1 if charged else 0)
if charged:
mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once()
mock_release_budget_reservation.assert_not_awaited()
else:
mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called()
mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation)

15
tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
View file

@ -6,7 +6,6 @@ Symbols pinned here:
- ``update_spend_logs_job``
- ``_monitor_spend_logs_queue``
- ``_raise_failed_update_spend_exception``
- ``spend_log_is_queued``
"""
from __future__ import annotations
@ -23,7 +22,6 @@ from litellm.proxy.utils import (
_monitor_spend_logs_queue,
_raise_failed_update_spend_exception,
drain_spend_logs_queue,
spend_log_is_queued,
update_daily_tag_spend,
update_spend,
update_spend_logs_job,
@ -631,16 +629,3 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None:
with pytest.raises(ValueError, match="specific"):
asyncio.run(_runner())
@pytest.mark.asyncio
async def test_spend_log_is_queued_matches_only_rows_awaiting_flush(
mock_prisma_client: Any, make_spend_log_row: Any
) -> None:
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")]
assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True
assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False
mock_prisma_client.spend_log_transactions = []
assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False