From b140dffd7e69f4648b67853b926286218504e732 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Tue, 15 Sep 2026 14:11:14 -0700 Subject: [PATCH] fix(spend_tracking): tell object reads apart by the request path instead of by call type A read, poll or download of a stored object answers with the id the request named in its path, so that id is the object's identity rather than one minted for this call. get_logging_payload now leaves metadata.response_id empty for those rows, and the flush re-keys only rows that carry a minted id or are already keyed on their own call id. The call-type list is gone, so a route this code has never heard of keeps collapsing on its object id. --- litellm/proxy/_types.py | 4 +- .../spend_tracking/spend_tracking_utils.py | 28 +++- litellm/proxy/utils.py | 88 +++---------- .../test_spend_tracking_utils.py | 36 +++++ .../test_proxy_update_spend.py | 124 +++++++++--------- 5 files changed, 147 insertions(+), 133 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 87bb73b3dd1..7badf041ef3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3876,7 +3876,9 @@ 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 + response_id: ReadOnly[ + str | None + ] # the id the provider minted for this response; empty for a row keyed on an object the request addressed class SpendLogsPayload(TypedDict): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index aa8d23a8d98..024715ac9f0 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -6,6 +6,7 @@ from datetime import datetime, timezone from datetime import datetime as dt from types import MappingProxyType from typing import Final, Literal, Protocol, cast, runtime_checkable +from urllib.parse import unquote, urlsplit from pydantic import BaseModel @@ -230,17 +231,33 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" +def _request_path_segments(litellm_params: Mapping[str, object]) -> frozenset[str]: + proxy_server_request: Final = litellm_params.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return frozenset() + request: Final = cast(Mapping[str, object], proxy_server_request) # cast-ok: built by add_litellm_data_to_request + url: Final = request.get("url") + if not isinstance(url, str): + return frozenset() + return frozenset(unquote(segment) for segment in urlsplit(url).path.split("/") if segment) + + def get_provider_response_id( - response_obj: Mapping[str, object], kwargs: Mapping[str, object], litellm_call_id: str | None + response_obj: Mapping[str, object], + kwargs: Mapping[str, object], + litellm_call_id: str | None, + litellm_params: Mapping[str, object], ) -> str | None: """The id the provider minted for this response: the response's own, else the one the standard - logging payload resolved. Never the proxy's call id, which is not a provider identity.""" + logging payload resolved. Never the proxy's call id, which is not a provider identity, and + never an id the request itself addressed in its path (a file, batch, response or vector store + read back by id), which identifies that object rather than this response.""" standard_logging_payload: Final = kwargs.get("standard_logging_object") candidate_ids: Final = ( response_obj.get("id"), standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None, ) - return next( + minted_id: Final = next( ( candidate for candidate in candidate_ids @@ -248,6 +265,9 @@ def get_provider_response_id( ), None, ) + if minted_id is None or minted_id in _request_path_segments(litellm_params): + return None + return minted_id def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None: @@ -418,7 +438,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) - provider_response_id: Final = get_provider_response_id(response_obj_dict, kwargs, litellm_call_id) + provider_response_id: Final = get_provider_response_id(response_obj_dict, kwargs, litellm_call_id, litellm_params) 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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7c376cabb03..1dc4471ea25 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7234,69 +7234,23 @@ async def _monitor_spend_logs_queue( MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 -_RESPONSE_ID_KEYED_SPEND_LOG_CALL_TYPES: Final = frozenset( - { - CallTypes.completion.value, - CallTypes.acompletion.value, - CallTypes.text_completion.value, - CallTypes.atext_completion.value, - CallTypes.embedding.value, - CallTypes.aembedding.value, - CallTypes.image_generation.value, - CallTypes.aimage_generation.value, - CallTypes.image_edit.value, - CallTypes.aimage_edit.value, - CallTypes.moderation.value, - CallTypes.amoderation.value, - CallTypes.transcription.value, - CallTypes.atranscription.value, - CallTypes.speech.value, - CallTypes.aspeech.value, - CallTypes.rerank.value, - CallTypes.arerank.value, - CallTypes.search.value, - CallTypes.asearch.value, - CallTypes.anthropic_messages.value, - CallTypes.aanthropic_messages.value, - CallTypes.responses.value, - CallTypes.aresponses.value, - CallTypes.generate_content.value, - CallTypes.agenerate_content.value, - CallTypes.generate_content_stream.value, - CallTypes.agenerate_content_stream.value, - CallTypes.ocr.value, - CallTypes.aocr.value, - CallTypes.vector_store_search.value, - CallTypes.avector_store_search.value, - CallTypes.create_interaction.value, - CallTypes.acreate_interaction.value, - CallTypes.create_video.value, - CallTypes.acreate_video.value, - CallTypes.video_remix.value, - CallTypes.avideo_remix.value, - CallTypes.video_edit.value, - CallTypes.avideo_edit.value, - CallTypes.video_extension.value, - CallTypes.avideo_extension.value, - CallTypes.create_container.value, - CallTypes.acreate_container.value, - CallTypes.run_code.value, - CallTypes.arun_code.value, - CallTypes.code_interpreter_tool.value, - CallTypes.acode_interpreter_tool.value, - CallTypes.call_mcp_tool.value, - CallTypes.send_message.value, - CallTypes.asend_message.value, - CallTypes.pass_through.value, - CallTypes.llm_passthrough_route.value, - CallTypes.allm_passthrough_route.value, - } -) -"""The inference and create calls, whose provider mints a response id per call: a stored row -with the same ``request_id`` is another request the provider gave the same id. Every other call -type (object reads, polls, cancels, lists and deletes, batch cost claims) is keyed on the id of the -object it addressed, and a second row for one of those collapses on purpose. Realtime and -Responses websocket sessions carry no provider id and are keyed on the call id already.""" + +def _spend_log_row_response_id(row: Mapping[str, object]) -> str | None: + """``metadata.response_id`` of a row: the id the provider minted for this response, absent + for rows keyed on an object the request addressed (a batch poll, a file or response read).""" + metadata: Final = row.get("metadata") + try: + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else metadata + except ValueError: + return None + response_id: Final = parsed.get("response_id") if isinstance(parsed, Mapping) else None + return response_id if isinstance(response_id, str) and response_id else None + + +def _spend_log_row_keyed_on_addressed_object(row: Mapping[str, object]) -> bool: + """A row keyed on the id of the object its request addressed rather than on a minted response id + or on its own call id; a second row for one of those collapses on purpose.""" + return _spend_log_row_response_id(row) is None and row.get("request_id") != row.get("litellm_call_id") def _is_transient_spend_log_write_error(e: Exception) -> bool: @@ -7325,11 +7279,11 @@ async def _spend_logs_rekeyed_after_duplicate_skip( 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 of a call type 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. + or a replay of it after a transport retry, and stays skipped. Rows keyed on the id of an + object the request addressed 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") in _RESPONSE_ID_KEYED_SPEND_LOG_CALL_TYPES) + unverified: Final = tuple(row for row in rows if not _spend_log_row_keyed_on_addressed_object(row)) if not unverified: return () stored_identities: Final = await SpendLogsRepository(WriterPinnedClient(repo.prisma_client.db)).stored_identities( 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 c765a961617..8947b1d1bab 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 @@ -1442,6 +1442,42 @@ def test_get_logging_payload_recognises_the_call_id_the_row_itself_resolves(): assert json.loads(payload["metadata"])["response_id"] is None +@pytest.mark.parametrize( + "url, call_type, expected", + [ + ("http://litellm/v1/files/file-abc", "afile_retrieve", None), + ("http://litellm/v1/files/file-abc/content", "afile_content", None), + ("http://litellm/openai/v1/batches/file-abc?limit=1", "aretrieve_batch", None), + ("http://litellm/v1/responses/file%2Dabc", "aget_responses", None), + ("http://litellm/v1/chat/completions", "acompletion", "file-abc"), + ("http://litellm/v1/responses?previous_response_id=file-abc", "aresponses", "file-abc"), + (None, "acompletion", "file-abc"), + ], +) +def test_get_logging_payload_leaves_metadata_response_id_empty_for_an_object_the_request_addressed( + url: str | None, call_type: str, expected: str | None +): + """A read, poll or download of a stored object answers with that object's id, which the request + named in its path: it is the object's identity, not an id minted for this call, and only a + minted id can mark a row as one the flush may re-key.""" + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "call_type": call_type, + "litellm_call_id": "call-1", + "litellm_params": { + "metadata": {"user_api_key": "test-key"}, + "proxy_server_request": {"url": url, "method": "GET"} if url is not None else None, + }, + }, + response_obj={"id": "file-abc"}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert json.loads(payload["metadata"])["response_id"] == expected + + @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 8cef2afe788..c66cfea9ede 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 @@ -948,6 +948,19 @@ def _wire_routed_db(mock_prisma_client: MagicMock, table: _SpendLogsTable) -> Ma return routed +def _inference_row( + make_spend_log_row: SpendLogRowFactory, *, request_id: str, litellm_call_id: str, response_id: str | None = None +) -> SpendLogRow: + """A row as ``get_logging_payload`` builds it for an inference call: ``metadata`` is a JSON + string carrying the id the provider minted, which is the row's key unless it was re-keyed.""" + return make_spend_log_row( + request_id=request_id, + litellm_call_id=litellm_call_id, + call_type="acompletion", + metadata=json.dumps({"response_id": response_id or request_id}), + ) + + async def _flush(mock_prisma_client: MagicMock, logs: list[SpendLogRow]) -> None: proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() @@ -970,7 +983,7 @@ async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_w 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") + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}") for i in range(3) ] @@ -985,27 +998,18 @@ async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_w @pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", - [ - "acompletion", - "atext_completion", - "aembedding", - "aresponses", - "aanthropic_messages", - "acreate_interaction", - "acreate_video", - "call_mcp_tool", - "allm_passthrough_route", - ], -) -async def test_update_spend_logs_rekeys_every_inference_call_type( - mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, call_type: str +@pytest.mark.parametrize("metadata", [json.dumps({"response_id": "static-1"}), {"response_id": "static-1"}]) +async def test_update_spend_logs_rekeys_a_row_carrying_the_id_the_provider_minted( + mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, metadata: str | dict[str, str] ) -> None: - """A self-hosted server reuses its id on every inference route it serves, not only chat.""" + """Whatever the route, a row whose key is the id the provider minted for that response is a + charged request of its own once another row holds the same id.""" table = _wire_spend_logs_table(mock_prisma_client) logs = [ - make_spend_log_row(request_id="static-1", litellm_call_id=f"call-{i}", call_type=call_type) for i in range(2) + make_spend_log_row( + request_id="static-1", litellm_call_id=f"call-{i}", call_type="aresponses", metadata=metadata + ) + for i in range(2) ] await _flush(mock_prisma_client, logs) @@ -1023,7 +1027,7 @@ async def test_update_spend_logs_reads_stored_identities_back_from_the_writer( 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") + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}") for i in range(2) ] @@ -1042,12 +1046,12 @@ async def test_update_spend_logs_rekeys_only_the_colliding_rows_of_a_mixed_batch 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"), + _inference_row(make_spend_log_row, request_id="chatcmpl-unique-a", litellm_call_id="call-a"), + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-b"), + _inference_row(make_spend_log_row, request_id="chatcmpl-static-2", litellm_call_id="call-c"), + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-d"), + _inference_row(make_spend_log_row, request_id="chatcmpl-unique-e", litellm_call_id="call-e"), + _inference_row(make_spend_log_row, request_id="chatcmpl-static-2", litellm_call_id="call-f"), ] await _flush(mock_prisma_client, logs) @@ -1069,14 +1073,12 @@ async def test_update_spend_logs_rekeys_a_row_whose_provider_id_an_earlier_flush """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") - ], + seeded=[_inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-old")], ) await _flush( mock_prisma_client, - [make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-new", call_type="acompletion")], + [_inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-new")], ) assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == { @@ -1093,12 +1095,12 @@ async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row( 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")], + seeded=[_inference_row(make_spend_log_row, request_id="chatcmpl-1", litellm_call_id="call-1")], ) await _flush( mock_prisma_client, - [make_spend_log_row(request_id="chatcmpl-1", litellm_call_id="call-1", call_type="acompletion")], + [_inference_row(make_spend_log_row, request_id="chatcmpl-1", litellm_call_id="call-1")], ) assert list(table.rows) == ["chatcmpl-1"] @@ -1107,40 +1109,34 @@ async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row( @pytest.mark.asyncio @pytest.mark.parametrize( - "call_type", + "call_type, request_id, metadata", [ - "acreate_batch", - "aretrieve_batch", - "acancel_batch", - "acreate_file", - "afile_retrieve", - "afile_content", - "afile_delete", - "aretrieve_fine_tuning_job", - "avideo_retrieve", - "aretrieve_container", - "avector_store_retrieve", - "avector_store_delete", - "aget_responses", - None, + ("aretrieve_batch", "batch_abc_batch_cost", json.dumps({"response_id": None})), + ("afile_retrieve", "file-abc", json.dumps({})), + ("avector_store_retrieve", "vs_abc", {}), + ("aget_responses", "resp_abc", "not json"), + (None, "obj_abc", None), ], ) async def test_update_spend_logs_keeps_object_keyed_rows_collapsed_on_their_object_id( - mock_prisma_client: MagicMock, make_spend_log_row: SpendLogRowFactory, call_type: str | None + mock_prisma_client: MagicMock, + make_spend_log_row: SpendLogRowFactory, + call_type: str | None, + request_id: str, + metadata: str | dict[str, str] | None, ) -> 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. Only the inference call types are re-keyed; anything else, a call type - this code has never heard of included, keeps collapsing.""" + is keyed on that object's id: such a row carries no minted response id, so a duplicate 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) + make_spend_log_row(request_id=request_id, litellm_call_id=f"call-{i}", call_type=call_type, metadata=metadata) for i in range(2) ] await _flush(mock_prisma_client, logs) - assert list(table.rows) == ["batch_abc_batch_cost"] + assert list(table.rows) == [request_id] assert table.query_raw_calls == 0 assert len(table.inserts) == 1 @@ -1153,13 +1149,17 @@ async def test_update_spend_logs_drops_and_logs_a_row_already_keyed_on_its_call_ 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")], + seeded=[_inference_row(make_spend_log_row, request_id="call-pinned", litellm_call_id="call-other")], ) 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")], + [ + make_spend_log_row( + request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion", metadata="{}" + ) + ], ) assert list(table.rows) == ["call-pinned"] @@ -1181,14 +1181,16 @@ async def test_update_spend_logs_rekey_that_collides_again_stops( 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"), + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-old"), + make_spend_log_row( + request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion", metadata="{}" + ), ], ) await _flush( mock_prisma_client, - [make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-pinned", call_type="acompletion")], + [_inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id="call-pinned")], ) assert sorted(table.rows) == ["call-pinned", "chatcmpl-static-1"] @@ -1204,7 +1206,7 @@ async def test_update_spend_logs_treats_a_replay_of_a_rekeyed_row_as_stored( 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") + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}") for i in range(2) ] await _flush(mock_prisma_client, logs) @@ -1231,7 +1233,7 @@ async def test_update_spend_logs_leaves_rows_skipped_when_the_read_back_fails_on 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") + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}") for i in range(2) ] @@ -1261,7 +1263,7 @@ async def test_update_spend_logs_retries_the_flush_when_the_read_back_hits_a_tra 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") + _inference_row(make_spend_log_row, request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}") for i in range(2) ]