scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

refactor(responses): mark a billed row completed without copying the response onto it

Browse source 
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
This commit is contained in:
ryan-crabbe-berri 2026-09-09 16:02:50 -07:00 • committed by jesus
parent 0235dbd7f2
commit ffa52f2f6a
2 changed files with 10 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -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(

54
tests/proxy_unit_tests/test_check_responses_cost.py
View file

@ -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