From ffa52f2f6aa9e5b7c5caa69effe191be5b8b3e29 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 16:02:50 -0700 Subject: [PATCH] refactor(responses): mark a billed row completed without copying the response onto it The previous commit stored the finished ResponsesAPIResponse in file_object. That duplicates content the provider still serves from its own copy, and the usage and spend it was meant to preserve already land in LiteLLM_SpendLogs on every billed call regardless of store_prompts_in_spend_logs, which gates only the messages and response body columns. The poller now writes status alone, as it did before. The write stays per job rather than one bulk update so a single failure cannot strand the rest of the cycle. Claude-Session: https://claude.ai/code/session_01Hn5E8Jz1LjGLFyiYxBRcBW --- .../common_utils/check_responses_cost.py | 24 ++++----- .../test_check_responses_cost.py | 54 ------------------- 2 files changed, 10 insertions(+), 68 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 62e742be680..c756832a7b7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -178,16 +178,12 @@ class CheckResponsesCost: f"so its cost will not be retried: {db_err}" ) - async def _persist_terminal_response( - self, job: "LiteLLM_ManagedObjectTable", response: ResponsesAPIResponse - ) -> None: - """Store the finished response on its managed row and retire the row from polling. + async def _mark_job_completed(self, job: "LiteLLM_ManagedObjectTable") -> None: + """Retire a billed row from polling, per job so one failure can't strand the rest. - The row is the only copy of a background generation's usage that outlives the poll, so - ``GET /v1/responses/{id}`` can serve a terminal job from here instead of re-reading it - from the provider. Every provider re-read replays the same usage and hands back a - freshly encoded id, which is what made this route bill per read and defeated id-based - dedup in the first place. + Only ``status`` is written. The generation's usage and spend already land in + ``LiteLLM_SpendLogs`` unconditionally, so copying the response body onto this row would + duplicate content the provider still serves, on a table nothing ever deletes from. ``status`` stays the literal "completed" for every terminal provider status, matching what this poller has always written, so stale-row expiry keeps skipping these rows. @@ -195,11 +191,11 @@ class CheckResponsesCost: try: await self.prisma_client.db.litellm_managedobjecttable.update_many( where={"id": job.id}, - data={"status": "completed", "file_object": response.model_dump_json()}, + data={"status": "completed"}, ) except Exception as db_err: verbose_proxy_logger.error( - f"CheckResponsesCost: failed to persist terminal response for job {job.id}: {db_err}" + f"CheckResponsesCost: failed to mark job {job.id} completed: {db_err}" ) async def check_responses_cost(self): @@ -293,10 +289,10 @@ class CheckResponsesCost: verbose_proxy_logger.info( f"Response {unified_object_id} has terminal status {response.status}, marking as complete" ) - completed_jobs.append((job, response)) + completed_jobs.append(job) - for job, response in completed_jobs: - await self._persist_terminal_response(job, response) + for job in completed_jobs: + await self._mark_job_completed(job) if len(completed_jobs) > 0: verbose_proxy_logger.info( diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py index 2d832975a11..bbf2b88e5b7 100644 --- a/tests/proxy_unit_tests/test_check_responses_cost.py +++ b/tests/proxy_unit_tests/test_check_responses_cost.py @@ -3,7 +3,6 @@ Unit tests for CheckResponsesCost class """ import asyncio -import json from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -31,10 +30,6 @@ def _completed_job_ids(mock_prisma_client): return [call.kwargs["where"]["id"] for call in _completion_calls(mock_prisma_client)] -def _persisted_response(completion_call): - return json.loads(completion_call.kwargs["data"]["file_object"]) - - def _claim_calls(mock_prisma_client): return _update_many_calls_writing( mock_prisma_client, lambda data: data == {"batch_processed": True} @@ -1119,55 +1114,6 @@ class TestCheckResponsesCost: mock_aget.assert_awaited_once() assert _completed_job_ids(mock_prisma_client) == ["job-old-schema"] - @pytest.mark.asyncio - async def test_terminal_response_replaces_the_queued_file_object_on_the_row( - self, check_responses_cost_instance, mock_prisma_client - ): - """The row has to become the durable copy of the finished generation. Leaving the queued - placeholder there forces the retrieve endpoint back to the provider, and every re-read - replays the same usage under a fresh id, which is what billed the job again per read.""" - mock_job = MagicMock() - mock_job.unified_object_id = "resp_test_persisted" - mock_job.created_by = "test-user" - mock_job.id = "job-persisted" - mock_job.file_object = { - "model": "gpt-5", - "id": "resp_test_persisted", - "status": "queued", - "usage": None, - } - - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) - - mock_response = ResponsesAPIResponse( - id="resp_finished_upstream", - object="response", - status="completed", - created_at=int(datetime.now().timestamp()), - output=[], - usage=ResponseAPIUsage( - input_tokens=100, output_tokens=50, total_tokens=150 - ), - ) - - with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: - mock_aget.return_value = mock_response - await check_responses_cost_instance.check_responses_cost() - - completion_calls = _completion_calls(mock_prisma_client) - assert len(completion_calls) == 1 - assert completion_calls[0].kwargs["where"] == {"id": "job-persisted"} - - persisted = _persisted_response(completion_calls[0]) - assert persisted["id"] == "resp_finished_upstream" - assert persisted["status"] == "completed" - assert persisted["usage"]["total_tokens"] == 150 - @pytest.mark.asyncio async def test_a_failed_persist_does_not_abort_the_rest_of_the_poll_cycle( self, check_responses_cost_instance, mock_prisma_client