diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fc05796db5a..87bb73b3dd1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3876,6 +3876,7 @@ class SpendLogsMetadata(TypedDict): autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed litellm_gateway_injected_cache: ReadOnly[str | None] router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model + response_id: ReadOnly[str | None] # the provider's response id; request_id unless the row was re-keyed class SpendLogsPayload(TypedDict): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3ef21996b9c..40039f640fc 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -144,6 +144,7 @@ def _get_spend_logs_metadata( litellm_call_id: str | None = None, autorouter_savings: float | None = None, router_metadata: SpendLogsRouterMetadata | None = None, + response_id: str | None = None, ) -> SpendLogsMetadata: if metadata is None: return SpendLogsMetadata( @@ -184,6 +185,7 @@ def _get_spend_logs_metadata( litellm_gateway_injected_cache=None, litellm_call_id=litellm_call_id, router_metadata=router_metadata, + response_id=response_id, ) verbose_proxy_logger.debug( "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys())) @@ -191,8 +193,13 @@ def _get_spend_logs_metadata( # Filter the metadata dictionary to include only the specified keys clean_metadata: Final = SpendLogsMetadata( - **{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key != "router_metadata"}, + **{ + key: metadata.get(key) + for key in SpendLogsMetadata.__annotations__ + if key not in ("router_metadata", "response_id") + }, router_metadata=router_metadata, + response_id=response_id, ) _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") @@ -387,6 +394,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs usage = _combined_usage.model_dump() id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) + raw_response_id: Final = response_obj_dict.get("id") standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None)) end_user_id = get_end_user_id_for_cost_tracking(litellm_params) @@ -522,6 +530,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None ), litellm_call_id=litellm_call_id, + response_id=raw_response_id if isinstance(raw_response_id, str) else None, router_metadata=_get_router_metadata_for_spend_log( metadata=metadata, requested_model=_model_group, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c095586b6c9..83f403ab8c5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -26,6 +26,7 @@ from litellm.constants import ( DEFAULT_MODEL_CREATED_AT_TIME, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, MAX_TEAM_LIST_LIMIT, + NON_INFERENCE_CALL_TYPES, SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, @@ -142,7 +143,7 @@ from litellm.proxy.db.prisma_client import ( PrismaWrapper, parse_iam_endpoint_from_url, ) -from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient from litellm.proxy.db.spend_log_batching import ( spend_log_queue_within_budget, spend_log_row_bytes, @@ -7234,11 +7235,102 @@ async def _monitor_spend_logs_queue( MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 +_OBJECT_ID_KEYED_SPEND_LOG_CALL_TYPES: Final = NON_INFERENCE_CALL_TYPES | frozenset( + { + CallTypes.create_batch.value, + CallTypes.acreate_batch.value, + CallTypes.retrieve_batch.value, + CallTypes.aretrieve_batch.value, + CallTypes.create_file.value, + CallTypes.acreate_file.value, + } +) +"""Rows keyed on the id of the object the call addressed rather than on a fresh response id: +the zero-priced reads and management calls of a stored object, batch polls (every poll of one +batch collapses into its single cost row, the claim in ``_claim_batch_cost_spend_log``) and file +uploads. A second row for one of these collapses on purpose.""" + def _is_transient_spend_log_write_error(e: Exception) -> bool: return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) +def _spend_log_row_can_rekey_on_call_id(row: Mapping[str, object]) -> bool: + call_id: Final = row.get("litellm_call_id") + return isinstance(call_id, str) and call_id != "" and call_id != row.get("request_id") + + +def _spend_log_row_is_stored(row: Mapping[str, object], stored_identities: frozenset[tuple[object, object]]) -> bool: + """A stored row with this row's identity is this row or a replay of it, under either key.""" + call_id: Final = row.get("litellm_call_id") + return (row.get("request_id"), call_id) in stored_identities or (call_id, call_id) in stored_identities + + +async def _spend_logs_rekeyed_after_duplicate_skip( + repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]] +) -> tuple[Mapping[str, object], ...]: + """The rows ``create_many(skip_duplicates=True)`` skipped, re-keyed on their own ``litellm_call_id``. + + ``request_id`` is the provider's response id, so a provider that reuses one completion id + (a self-hosted OpenAI-compatible server, commonly) had every request after the first + served and charged but never logged (LIT-6666). The insert reports only how many rows + landed, so the rows are read back by identity from the writer (a lagging read replica + would report the rows this very insert landed as missing): a stored row carrying this + row's ``request_id`` AND ``litellm_call_id``, or already keyed on its call id, is this row + or a replay of it after a transport retry, and stays skipped. Rows keyed on an object id + collapse on purpose, and a row already keyed on its call id has nothing left to fall back + to; those are skipped and only logged. + """ + unverified: Final = tuple(row for row in rows if row.get("call_type") not in _OBJECT_ID_KEYED_SPEND_LOG_CALL_TYPES) + if not unverified: + return () + stored_identities: Final = await SpendLogsRepository(WriterPinnedClient(repo.prisma_client.db)).stored_identities( + str(key) for row in unverified for key in (row.get("request_id"), row.get("litellm_call_id")) if key + ) + dropped: Final = tuple(row for row in unverified if not _spend_log_row_is_stored(row, stored_identities)) + rekeyable: Final = tuple(row for row in dropped if _spend_log_row_can_rekey_on_call_id(row)) + for row in dropped: + if not _spend_log_row_can_rekey_on_call_id(row): + verbose_proxy_logger.error( + "Spend tracking - dropping spend log row whose request_id %s is already taken and whose " + "litellm_call_id %s offers no other key", + row.get("request_id"), + row.get("litellm_call_id"), + ) + if rekeyable: + verbose_proxy_logger.warning( + "Spend tracking - %d spend log row(s) shared a request_id with a stored row and were re-keyed on " + "their litellm_call_id: %s", + len(rekeyable), + sorted(frozenset(str(row.get("request_id")) for row in rekeyable)), + ) + return tuple( + {**row, "request_id": row["litellm_call_id"]} # mutable-ok: prisma create_many takes dict rows + for row in rekeyable + ) + + +async def _spend_logs_rekeyed_or_left_skipped( + repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]] +) -> tuple[Mapping[str, object], ...]: + """``_spend_logs_rekeyed_after_duplicate_skip``, leaving the rows skipped (the pre-existing + outcome) when the read-back fails on its data rather than on transport, so one bad read + cannot requeue a whole flush behind it forever.""" + try: + return await _spend_logs_rekeyed_after_duplicate_skip(repo, rows) + except Exception as e: + if _is_transient_spend_log_write_error(e): + raise + spend_log_error( + "Spend tracking - could not read back which of %d skipped spend log rows are stored; leaving them " + "skipped. error=%s", + len(rows), + str(e), + exc=e, + ) + return () + + async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], @@ -7267,8 +7359,7 @@ async def _create_spend_logs_with_poison_isolation( persist. Returns the budget left after this subtree. """ try: - await repo.table.create_many(data=rows, skip_duplicates=True) - return failure_budget + inserted: Final = await repo.table.create_many(data=rows, skip_duplicates=True) except Exception as e: if not PrismaDBExceptionHandler.is_prisma_data_error(e): raise @@ -7303,6 +7394,11 @@ async def _create_spend_logs_with_poison_isolation( ) return 0 return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining) + if isinstance(inserted, int) and inserted < len(rows): + rekeyed: Final = await _spend_logs_rekeyed_or_left_skipped(repo, rows) + if rekeyed: + return await _create_spend_logs_with_poison_isolation(repo, rekeyed, failure_budget) + return failure_budget def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging): diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 1ad7a735d96..f109619b4b1 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -7,6 +7,7 @@ These are thin wrappers for tables that do not (yet) need domain-specific query methods; richer repositories live in their own modules. """ +from collections.abc import Iterable from typing import TYPE_CHECKING, Any, Final, Generic from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync @@ -64,9 +65,19 @@ class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.Lite table_name = "litellm_organizationmembership" +_SPEND_LOG_IDENTITIES_SQL: Final = ( + 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1::text[])' +) + + class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]): table_name = "litellm_spendlogs" + async def stored_identities(self, request_ids: Iterable[str]) -> frozenset[tuple[object, object]]: + """``(request_id, litellm_call_id)`` of every stored row keyed on one of ``request_ids``.""" + rows: Final = await self.prisma_client.db.query_raw(_SPEND_LOG_IDENTITIES_SQL, sorted(frozenset(request_ids))) + return frozenset((row.get("request_id"), row.get("litellm_call_id")) for row in rows) + class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]): table_name = "litellm_budgetwindowspend" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 8283ee8395a..0a3aae959a0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -643,6 +643,7 @@ ignored_keys = [ "request_id", "litellm_call_id", "metadata.litellm_call_id", + "metadata.response_id", "session_id", "startTime", "endTime", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index a72b4e28143..a0b6d4f6bc6 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1356,6 +1356,43 @@ def test_get_logging_payload_populates_litellm_call_id_alongside_provider_reques assert payload["litellm_call_id"] == call_id +def test_get_logging_payload_keeps_the_provider_response_id_in_metadata(): + """LIT-6666: a row the flush re-keys on its call id (its provider id was already taken + by another request) still has to say which response it logged.""" + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {"user_api_key": "test-key", "response_id": "caller-supplied"}}, + }, + response_obj=litellm.ModelResponse( + id="chatcmpl-static-1", + choices=[], + usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert json.loads(payload["metadata"])["response_id"] == "chatcmpl-static-1" + + +def test_get_logging_payload_leaves_metadata_response_id_empty_without_a_response_id(): + """A failure row has no provider response; a caller cannot fill the slot through request metadata.""" + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {"user_api_key": "test-key", "response_id": "caller-supplied"}}, + }, + response_obj=None, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert json.loads(payload["metadata"])["response_id"] is None + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d671a4ffc1f..21d74e32a20 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -10,6 +10,9 @@ from __future__ import annotations import asyncio import json +import logging + +import httpx from collections.abc import Iterator from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -17,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm.proxy.utils as utils_mod +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.db.spend_log_batching import spend_log_row_bytes from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, enqueue_spend_logs @@ -475,7 +479,7 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage( def _deadlock_error() -> Exception: return _data_error( - 'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, ' + "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, " 'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })' ) @@ -508,9 +512,7 @@ async def test_update_spend_logs_retries_deadlock_and_keeps_every_row( logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], ) - attempts = tuple( - tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list - ) + attempts = tuple(tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list) assert attempts == (("a", "b"), ("a", "b"), ("a", "b")) assert mock_prisma_client.spend_log_transactions == [] @@ -883,3 +885,338 @@ def test_disable_spend_updates_error_when_general_settings_unavailable( monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False) with pytest.raises(ImportError): ProxyUpdateSpend.disable_spend_updates() + + +class _SpendLogsTable: + """``LiteLLM_SpendLogs`` as the flush sees it: ``create_many`` is ``INSERT ... ON CONFLICT + DO NOTHING`` on ``request_id`` and reports how many rows landed, and ``query_raw`` answers the + identity read-back with the stored ``(request_id, litellm_call_id)`` pairs.""" + + def __init__(self, seeded: List[Dict[str, Any]] | None = None) -> None: + self.rows: Dict[str, Dict[str, Any]] = {row["request_id"]: dict(row) for row in seeded or []} + self.inserts: List[List[str]] = [] + self.query_raw_calls = 0 + + async def create_many(self, *, data: Any, skip_duplicates: bool) -> int: + assert skip_duplicates is True + self.inserts.append([row["request_id"] for row in data]) + inserted = 0 + for row in data: + if row["request_id"] in self.rows: + continue + self.rows[row["request_id"]] = dict(row) + inserted += 1 + return inserted + + async def query_raw(self, sql: str, request_ids: Any) -> List[Dict[str, Any]]: + self.query_raw_calls += 1 + assert "WHERE request_id = ANY($1::text[])" in sql + return [ + {"request_id": rid, "litellm_call_id": row.get("litellm_call_id")} + for rid, row in self.rows.items() + if rid in request_ids + ] + + +def _wire_spend_logs_table(mock_prisma_client: Any, seeded: List[Dict[str, Any]] | None = None) -> _SpendLogsTable: + table = _SpendLogsTable(seeded) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=table.create_many) + mock_prisma_client.db.query_raw = AsyncMock(side_effect=table.query_raw) + return table + + +def _wire_routed_db(mock_prisma_client: Any, table: _SpendLogsTable) -> MagicMock: + """``prisma_client.db`` as a read-replica router: top-level reads go to the reader, which never + sees this flush's writes, and ``writer`` is the engine the rows landed on.""" + routed = MagicMock(spec=RoutingPrismaWrapper) + routed.writer_unavailable = False + routed.litellm_spendlogs = MagicMock() + routed.litellm_spendlogs.create_many = AsyncMock(side_effect=table.create_many) + routed.query_raw = AsyncMock(return_value=[]) + routed.writer = MagicMock() + routed.writer.query_raw = AsyncMock(side_effect=table.query_raw) + mock_prisma_client.db = routed + return routed + + +async def _flush(mock_prisma_client: Any, logs: List[Dict[str, Any]]) -> None: + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=logs, + ) + + +@pytest.mark.asyncio +async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_would_drop( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """Three requests answered with one provider id are three charged requests, so three rows + must land: the first keeps the provider id, the other two are re-keyed on their own call id. + Observed on a live proxy against a self-hosted server returning a fixed completion id: the + duplicate-tolerant flush kept one row while key and daily spend counted all three.""" + table = _wire_spend_logs_table(mock_prisma_client) + logs = [ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion") + for i in range(3) + ] + + await _flush(mock_prisma_client, logs) + + assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == { + "chatcmpl-static-1": "call-0", + "call-1": "call-1", + "call-2": "call-2", + } + assert table.inserts == [["chatcmpl-static-1"] * 3, ["call-1", "call-2"]] + + +@pytest.mark.asyncio +async def test_update_spend_logs_reads_stored_identities_back_from_the_writer( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """The read-back must see the rows this flush just landed, so it goes to the writer: with a + read replica configured, ``db.query_raw`` is routed to the reader, and a lagging reader would + report every just-inserted row as missing and re-insert it under its call id.""" + table = _wire_spend_logs_table(mock_prisma_client) + routed = _wire_routed_db(mock_prisma_client, table) + logs = [ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion") + for i in range(2) + ] + + await _flush(mock_prisma_client, logs) + + routed.writer.query_raw.assert_awaited_once() + routed.query_raw.assert_not_awaited() + assert sorted(table.rows) == ["call-1", "chatcmpl-static-1"] + + +@pytest.mark.asyncio +async def test_update_spend_logs_rekeys_only_the_colliding_rows_of_a_mixed_batch( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """Two providers reusing ids and one issuing unique ids share a flush: every row of the + unique provider keeps its key and only the later rows of each reused id are re-keyed.""" + table = _wire_spend_logs_table(mock_prisma_client) + logs = [ + make_spend_log_row(request_id="chatcmpl-unique-a", litellm_call_id="call-a", call_type="acompletion"), + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-b", call_type="acompletion"), + make_spend_log_row(request_id="chatcmpl-static-2", litellm_call_id="call-c", call_type="acompletion"), + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-d", call_type="acompletion"), + make_spend_log_row(request_id="chatcmpl-unique-e", litellm_call_id="call-e", call_type="acompletion"), + make_spend_log_row(request_id="chatcmpl-static-2", litellm_call_id="call-f", call_type="acompletion"), + ] + + await _flush(mock_prisma_client, logs) + + assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == { + "chatcmpl-unique-a": "call-a", + "chatcmpl-static-1": "call-b", + "chatcmpl-static-2": "call-c", + "call-d": "call-d", + "chatcmpl-unique-e": "call-e", + "call-f": "call-f", + } + + +@pytest.mark.asyncio +async def test_update_spend_logs_rekeys_a_row_whose_provider_id_an_earlier_flush_stored( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """The colliding row usually landed in an earlier flush, or from another worker.""" + table = _wire_spend_logs_table( + mock_prisma_client, + seeded=[ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-old", call_type="acompletion") + ], + ) + + await _flush( + mock_prisma_client, + [make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-new", call_type="acompletion")], + ) + + assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == { + "chatcmpl-static-1": "call-old", + "call-new": "call-new", + } + + +@pytest.mark.asyncio +async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A transport retry replays the whole batch; a row already stored under its own call id is + the same request, not a collision, and must not become a second row.""" + table = _wire_spend_logs_table( + mock_prisma_client, + seeded=[make_spend_log_row(request_id="chatcmpl-1", litellm_call_id="call-1", call_type="acompletion")], + ) + + await _flush( + mock_prisma_client, + [make_spend_log_row(request_id="chatcmpl-1", litellm_call_id="call-1", call_type="acompletion")], + ) + + assert list(table.rows) == ["chatcmpl-1"] + assert len(table.inserts) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", ["acreate_batch", "aretrieve_batch", "acreate_file", "avector_store_retrieve", "avector_store_delete"] +) +async def test_update_spend_logs_keeps_object_keyed_rows_collapsed_on_their_object_id( + mock_prisma_client: Any, make_spend_log_row: Any, call_type: str +) -> None: + """Every poll of one batch shares its cost row by design, and every read of a stored object + is keyed on that object's id, so a duplicate here is not a lost row and the identity read-back + is not even issued.""" + table = _wire_spend_logs_table(mock_prisma_client) + logs = [ + make_spend_log_row(request_id="batch_abc_batch_cost", litellm_call_id=f"call-{i}", call_type=call_type) + for i in range(2) + ] + + await _flush(mock_prisma_client, logs) + + assert list(table.rows) == ["batch_abc_batch_cost"] + assert table.query_raw_calls == 0 + assert len(table.inserts) == 1 + + +@pytest.mark.asyncio +async def test_update_spend_logs_drops_and_logs_a_row_already_keyed_on_its_call_id( + mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture +) -> None: + """A row whose key already is its call id (a client pinning x-litellm-call-id on a failure + row, say) has nothing to fall back to: it stays dropped, but loudly.""" + table = _wire_spend_logs_table( + mock_prisma_client, + seeded=[make_spend_log_row(request_id="call-pinned", litellm_call_id="call-other", call_type="acompletion")], + ) + + with caplog.at_level(logging.ERROR, logger=utils_mod.verbose_proxy_logger.name): + await _flush( + mock_prisma_client, + [make_spend_log_row(request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion")], + ) + + assert list(table.rows) == ["call-pinned"] + assert len(table.inserts) == 1 + dropped = [record.getMessage() for record in caplog.records if "dropping spend log row" in record.getMessage()] + assert dropped == [ + "Spend tracking - dropping spend log row whose request_id call-pinned is already taken and whose " + "litellm_call_id call-pinned offers no other key" + ] + + +@pytest.mark.asyncio +async def test_update_spend_logs_rekey_that_collides_again_stops( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A client pinning one x-litellm-call-id against a provider reusing one response id + collides on both keys. The stored row already keyed on that call id reads as this row's + replay, so the row stays skipped in the same pass, exactly as a pinned call id skips today.""" + table = _wire_spend_logs_table( + mock_prisma_client, + seeded=[ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-old", call_type="acompletion"), + make_spend_log_row(request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion"), + ], + ) + + await _flush( + mock_prisma_client, + [make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-pinned", call_type="acompletion")], + ) + + assert sorted(table.rows) == ["call-pinned", "chatcmpl-static-1"] + assert table.inserts == [["chatcmpl-static-1"]] + + +@pytest.mark.asyncio +async def test_update_spend_logs_treats_a_replay_of_a_rekeyed_row_as_stored( + mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture +) -> None: + """A transport retry replays a whole batch, including rows an earlier attempt already re-keyed. + Those are stored under their call id, so the read-back finds them there and nothing is re-keyed + or warned about a second time.""" + table = _wire_spend_logs_table(mock_prisma_client) + logs = [ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion") + for i in range(2) + ] + await _flush(mock_prisma_client, logs) + inserts_after_first_flush = len(table.inserts) + caplog.clear() + + with caplog.at_level(logging.WARNING, logger=utils_mod.verbose_proxy_logger.name): + await _flush(mock_prisma_client, logs) + + assert sorted(table.rows) == ["call-1", "chatcmpl-static-1"] + assert len(table.inserts) == inserts_after_first_flush + 1 + assert [record for record in caplog.records if "re-keyed" in record.getMessage()] == [] + + +@pytest.mark.asyncio +async def test_update_spend_logs_leaves_rows_skipped_when_the_read_back_fails_on_its_data( + mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture +) -> None: + """A read-back the database rejects (not a transport fault) must not requeue the whole flush + behind it: the skipped rows stay skipped, which is what happened before, and the failure is + logged.""" + table = _wire_spend_logs_table(mock_prisma_client) + mock_prisma_client.db.query_raw = AsyncMock(side_effect=_data_error("invalid input syntax for type text[]")) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + logs = [ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion") + for i in range(2) + ] + + with caplog.at_level(logging.ERROR, logger=utils_mod.verbose_proxy_logger.name): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=logs, + ) + + assert list(table.rows) == ["chatcmpl-static-1"] + assert len(table.inserts) == 1 + assert mock_prisma_client.spend_log_transactions == [] + assert any("could not read back" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_update_spend_logs_retries_the_flush_when_the_read_back_hits_a_transport_fault( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A transport fault during the read-back is the same outage as one during the insert, so + the flush retries (and finally requeues) instead of leaving the rows behind.""" + table = _wire_spend_logs_table(mock_prisma_client) + mock_prisma_client.db.query_raw = AsyncMock(side_effect=httpx.ReadError("network blip")) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + logs = [ + make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion") + for i in range(2) + ] + + with pytest.raises(httpx.ReadError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=logs, + ) + + assert mock_prisma_client.spend_log_transactions == logs