scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(responses): resolve previous_response_id for a just-written turn

Browse source 
The session lookup reads spend logs straight out of the database, so a
follow-up sent seconds after the turn it chains off found nothing while the
row was still queued in the worker that served it, and the conversation was
dropped without an error. Responses calls now ask the spend-log writer to
flush on its next pass instead of waiting out its poll interval, and the
lookup gives a just-finished turn a short second chance.

Replaying a session also accepted `input` only as a string or a single dict,
so the standard list shape dropped every user turn and left the model with
assistant messages alone.
This commit is contained in:
mateo-berri 2026-08-22 11:46:24 -07:00
parent 70e4273ba1
commit f89a3693ba
7 changed files with 192 additions and 184 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -1542,6 +1542,8 @@ SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BA
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000")))
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))

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

@ -65,6 +65,7 @@ from litellm.proxy.spend_tracking.savings import (
)
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.repositories.prisma_protocols import BatchTable
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -73,6 +74,9 @@ else:
ProxyLogging = Any
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
class _SpendBatch(Protocol):
litellm_usertable: BatchTable
litellm_verificationtoken: BatchTable
@ -820,9 +824,11 @@ class DBSpendUpdateWriter:
)
)
if prisma_client is not None and spend_logs_url is not None or prisma_client is not None:
from litellm.proxy.utils import enqueue_spend_logs
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 RESPONSES_SESSION_CALL_TYPES:
request_spend_log_flush()
else:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")

26
litellm/proxy/utils.py
View file

@ -3341,6 +3341,7 @@ class _StaleReadEngine:
class PrismaClient:
spend_log_transactions: list = []
_spend_log_transactions_lock = asyncio.Lock()
spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event()
spend_log_queue_bytes: ClassVar[int] = 0
spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None
tool_usage_transactions: list["ToolUsageTransaction"] = []
@ -6005,14 +6006,24 @@ async def enqueue_spend_logs(
)
async def peek_spend_logs(prisma_client: PrismaClient) -> tuple[SpendLogsPayload, ...]:
"""Snapshot the spend logs still waiting for the next flush, leaving the queue intact.
def request_spend_log_flush() -> None:
"""Wake the queue monitor now rather than leaving the rows for its next poll.
Reads that need a just-finished request use this, since the batch writer only
reaches the DB every ``PROXY_BATCH_WRITE_AT`` seconds.
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.
Repeated requests coalesce into the monitor's next pass, so the batching holds.
"""
async with prisma_client._spend_log_transactions_lock:
return tuple(prisma_client.spend_log_transactions)
PrismaClient.spend_log_flush_requested.set()
async def _wait_for_spend_log_flush_request(interval: float) -> bool:
"""Wait out ``interval``, returning early and True when a flush was requested."""
try:
await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval)
except asyncio.TimeoutError:
return False
PrismaClient.spend_log_flush_requested.clear()
return True
async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]:
@ -6460,7 +6471,8 @@ async def _monitor_spend_logs_queue(
# Exponential backoff when no logs to process
current_interval = min(current_interval * backoff_multiplier, max_backoff)
await asyncio.sleep(current_interval)
if await _wait_for_spend_log_flush_request(current_interval):
current_interval = base_interval
except Exception as e:
spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e)
# Continue monitoring even if there's an error, with exponential backoff

86
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -1,5 +1,5 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
import litellm
@ -132,8 +132,8 @@ class ResponsesSessionHandler:
############################################################
# Add Output messages for this Spend Log
############################################################
_response_output: Final = ResponsesSessionHandler._get_response_dict_from_spend_log(spend_log)
if _response_output:
_response_output: Final = spend_log.get("response", "{}")
if isinstance(_response_output, dict) and _response_output and _response_output != {}:
# transform `ChatCompletion Response` to `ResponsesAPIResponse`
model_response: Final = ModelResponse(**_response_output)
for choice in model_response.choices:
@ -141,23 +141,6 @@ class ResponsesSessionHandler:
chat_completion_message_history.append(getattr(choice, "message"))
return chat_completion_message_history
@staticmethod
def _get_response_dict_from_spend_log(spend_log: SpendLogsPayload) -> Mapping[str, Any] | None:
"""
Spend logs read from the DB hold `response` as a dict, ones still queued in memory
hold it as a JSON string.
"""
_response_output: Final = spend_log.get("response")
if isinstance(_response_output, dict):
return _response_output or None
if isinstance(_response_output, str):
try:
parsed: Final = json.loads(_response_output)
except json.JSONDecodeError:
return None
return parsed if isinstance(parsed, dict) and parsed else None
return None
@staticmethod
async def get_proxy_server_request_from_spend_log(
spend_log: SpendLogsPayload,
@ -272,9 +255,16 @@ class ResponsesSessionHandler:
SQL query
SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id
A just-finished turn gets a short second chance: the worker that served it may
still be writing its spend log when the follow-up arrives, and an empty result
drops the whole conversation instead of erroring.
"""
from litellm.constants import (
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS,
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL,
)
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.utils import peek_spend_logs
verbose_proxy_logger.debug("decoding response id=%s", previous_response_id)
@ -295,46 +285,16 @@ class ResponsesSessionHandler:
ORDER BY "endTime" ASC;
"""
written_spend_logs: Final = await prisma_client.db.query_raw(query, response_id)
queued_spend_logs: Final = await peek_spend_logs(prisma_client)
spend_logs: Final = list(
ResponsesSessionHandler._merge_queued_spend_logs(
response_id=response_id,
written_spend_logs=written_spend_logs,
queued_spend_logs=queued_spend_logs,
)
)
for attempt in range(RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS):
if attempt:
await asyncio.sleep(RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL)
if spend_logs := await prisma_client.db.query_raw(query, response_id):
verbose_proxy_logger.debug(
"Found the following spend logs for previous response id %s: %s",
response_id,
json.dumps(spend_logs, indent=4, default=str),
)
return spend_logs
verbose_proxy_logger.debug(
"Found the following spend logs for previous response id %s: %s",
response_id,
json.dumps(spend_logs, indent=4, default=str),
)
return spend_logs
@staticmethod
def _merge_queued_spend_logs(
response_id: str,
written_spend_logs: Sequence[SpendLogsPayload],
queued_spend_logs: Sequence[SpendLogsPayload],
) -> tuple[SpendLogsPayload, ...]:
"""
Append the session's spend logs that the batch writer has not flushed to the DB yet.
Without this a follow-up sent inside the ``PROXY_BATCH_WRITE_AT`` window sees an
empty session and silently drops the conversation. The queue is FIFO, so anything
still on it is newer than every row already written.
"""
session_ids: Final = frozenset(
session_id
for spend_log in (*written_spend_logs, *queued_spend_logs)
if spend_log.get("request_id") == response_id and (session_id := spend_log.get("session_id"))
) | frozenset(session_id for spend_log in written_spend_logs if (session_id := spend_log.get("session_id")))
written_request_ids: Final = frozenset(spend_log.get("request_id") for spend_log in written_spend_logs)
unflushed: Final = tuple(
spend_log
for spend_log in queued_spend_logs
if spend_log.get("session_id") in session_ids and spend_log.get("request_id") not in written_request_ids
)
return (*written_spend_logs, *unflushed)
verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id)
return []

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

@ -2712,3 +2712,31 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
assert mock_prisma_client.db.tx.call_count == 2
proxy_logging.failure_handler.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type, expects_flush",
[("aresponses", True), ("responses", True), ("acompletion", False)],
)
async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(
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.
"""
from litellm.proxy.utils import PrismaClient
db_writer = DBSpendUpdateWriter()
prisma = _tool_usage_prisma()
PrismaClient.spend_log_flush_requested.clear()
await db_writer._insert_spend_log_to_db(
payload={"request_id": "req-1", "call_type": call_type},
prisma_client=prisma,
)
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush
PrismaClient.spend_log_flush_requested.clear()

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

@ -11,7 +11,8 @@ Symbols pinned here:
from __future__ import annotations
import asyncio
from typing import Any, Dict, List
from contextlib import suppress
from typing import Any, Dict, Final, List
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -526,6 +527,53 @@ async def test_monitor_spend_logs_queue_swallows_errors_and_backs_off(
assert sleep_count["n"] == 3
@pytest.mark.asyncio
async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested(
mock_prisma_client: Any,
make_spend_log_row: Any,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A requested flush wakes the monitor mid-poll, so a Responses row reaches the DB
before the client can chain a `previous_response_id` off it.
"""
import litellm.constants as constants_mod
import litellm.proxy.utils as utils_mod
from litellm.proxy.utils import PrismaClient, request_spend_log_flush
monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False)
PrismaClient.spend_log_flush_requested.clear()
mock_prisma_client.spend_log_transactions = []
mock_prisma_client.tool_usage_transactions = []
flushed: Final = asyncio.Event()
async def _fake_job(*args: Any, **kwargs: Any) -> None:
flushed.set()
monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job)
monitor: Final = asyncio.create_task(
_monitor_spend_logs_queue(
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=MagicMock(),
)
)
try:
await asyncio.sleep(0.05)
assert not flushed.is_set()
mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1"))
request_spend_log_flush()
await asyncio.wait_for(flushed.wait(), timeout=5.0)
finally:
monitor.cancel()
with suppress(asyncio.CancelledError):
await monitor
PrismaClient.spend_log_flush_requested.clear()
def test_raise_failed_update_spend_exception_emits_failure_handler() -> None:
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()

176
tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py
View file

@ -1,4 +1,3 @@
import asyncio
import json
from unittest.mock import AsyncMock, patch
@ -451,20 +450,38 @@ def _chat_completion_response(request_id: str, content: str) -> dict:
class _FakePrismaDB:
def __init__(self, rows):
self._rows = rows
def __init__(self, results):
self._results = list(results)
self.calls = []
async def query_raw(self, query, *args):
self.calls.append(args)
return list(self._rows)
if not self._results:
return []
return list(self._results.pop(0))
class _FakePrismaClient:
def __init__(self, written_rows, queued_rows):
self.db = _FakePrismaDB(written_rows)
self.spend_log_transactions = list(queued_rows)
self._spend_log_transactions_lock = asyncio.Lock()
def __init__(self, results):
self.db = _FakePrismaDB(results)
def _spend_log(request_id: str, session_id: str, prompt: str, answer: str) -> dict:
return {
"request_id": request_id,
"call_type": "aresponses",
"session_id": session_id,
"proxy_server_request": {
"input": [{"role": "user", "content": prompt}],
"model": "claude-bridge",
},
"response": _chat_completion_response(request_id, answer),
}
@pytest.fixture
def instant_session_lookup_retries(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm.constants, "RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.0)
@pytest.mark.asyncio
@ -475,21 +492,12 @@ async def test_message_history_reconstructs_list_shaped_input():
"""
request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb"
mock_spend_logs = [
{
"request_id": request_id,
"call_type": "aresponses",
"session_id": "a96757c4-c6dc-4c76-b37e-e7dfa526b701",
"proxy_server_request": {
"input": [
{
"role": "user",
"content": "Remember this: my favorite color is chartreuse.",
}
],
"model": "claude-bridge",
},
"response": _chat_completion_response(request_id, "OK"),
}
_spend_log(
request_id,
"a96757c4-c6dc-4c76-b37e-e7dfa526b701",
"Remember this: my favorite color is chartreuse.",
"OK",
)
]
with patch.object(
@ -512,31 +520,22 @@ async def test_message_history_reconstructs_list_shaped_input():
@pytest.mark.asyncio
async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_writer():
async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flushed_yet(
instant_session_lookup_retries: None,
):
"""
A follow-up sent right after the previous turn arrives before the batch writer has
flushed that turn's spend log, so the row is only in memory. The history has to
include it anyway.
A follow-up sent right after the previous turn can beat that turn's spend log to the
DB. The lookup has to try again instead of handing back an empty conversation.
"""
request_id = "chatcmpl-6c1f5f6c-6a2b-4c62-8d1f-0d9d4ce0a1b2"
queued_spend_log = {
"request_id": request_id,
"call_type": "aresponses",
"session_id": "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3",
"proxy_server_request": json.dumps(
{
"input": [
{
"role": "user",
"content": "Remember this: my favorite color is chartreuse.",
}
],
"model": "claude-bridge",
}
),
"response": json.dumps(_chat_completion_response(request_id, "OK")),
}
fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[queued_spend_log])
session_id = "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3"
spend_log = _spend_log(
request_id,
session_id,
"Remember this: my favorite color is chartreuse.",
"OK",
)
fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]])
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
@ -548,44 +547,22 @@ async def test_message_history_includes_spend_logs_still_waiting_on_the_batch_wr
("user", "Remember this: my favorite color is chartreuse."),
("assistant", "OK"),
]
assert result["litellm_session_id"] == "b7d0a5b0-6d20-4a68-9d24-6ba0f6d1f1a3"
assert fake_prisma_client.spend_log_transactions == [queued_spend_log]
assert result["litellm_session_id"] == session_id
assert fake_prisma_client.db.calls == [(request_id,), (request_id,)]
@pytest.mark.asyncio
async def test_message_history_merges_written_and_queued_turns_in_order():
"""
Turn 1 already flushed to the DB, turn 2 still queued: the follow-up sees the whole
conversation, in order, with no row counted twice.
"""
async def test_message_history_reconstructs_every_turn_of_the_session_in_order():
session_id = "5c5f9a3e-1c86-4c0e-9d7c-0a54b8a0f2f1"
first_request_id = "chatcmpl-1111"
second_request_id = "chatcmpl-2222"
written_spend_log = {
"request_id": first_request_id,
"call_type": "aresponses",
"session_id": session_id,
"proxy_server_request": {
"input": [{"role": "user", "content": "My favorite color is chartreuse."}],
"model": "claude-bridge",
},
"response": _chat_completion_response(first_request_id, "Got it."),
}
queued_spend_log = {
"request_id": second_request_id,
"call_type": "aresponses",
"session_id": session_id,
"proxy_server_request": json.dumps(
{
"input": [{"role": "user", "content": "And my favorite city is Lisbon."}],
"model": "claude-bridge",
}
),
"response": json.dumps(_chat_completion_response(second_request_id, "Noted.")),
}
fake_prisma_client = _FakePrismaClient(
written_rows=[written_spend_log],
queued_rows=[written_spend_log, queued_spend_log],
results=[
[
_spend_log(first_request_id, session_id, "My favorite color is chartreuse.", "Got it."),
_spend_log(second_request_id, session_id, "And my favorite city is Lisbon.", "Noted."),
]
]
)
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
@ -604,45 +581,18 @@ async def test_message_history_merges_written_and_queued_turns_in_order():
@pytest.mark.asyncio
async def test_message_history_ignores_queued_spend_logs_from_other_sessions():
request_id = "chatcmpl-3333"
written_spend_log = {
"request_id": request_id,
"call_type": "aresponses",
"session_id": "session-a",
"proxy_server_request": {
"input": [{"role": "user", "content": "Hello from session a."}],
"model": "claude-bridge",
},
"response": _chat_completion_response(request_id, "Hi."),
}
other_session_spend_log = {
"request_id": "chatcmpl-4444",
"call_type": "aresponses",
"session_id": "session-b",
"proxy_server_request": json.dumps(
{
"input": [{"role": "user", "content": "Hello from session b."}],
"model": "claude-bridge",
}
),
"response": json.dumps(_chat_completion_response("chatcmpl-4444", "Hi there.")),
}
fake_prisma_client = _FakePrismaClient(
written_rows=[written_spend_log],
queued_rows=[other_session_spend_log],
)
async def test_session_lookup_stops_retrying_once_the_budget_is_spent(
instant_session_lookup_retries: None,
):
fake_prisma_client = _FakePrismaClient(results=[])
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
request_id
spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(
"chatcmpl-does-not-exist"
)
messages = result["messages"]
assert [(message.get("role"), message.get("content")) for message in messages] == [
("user", "Hello from session a."),
("assistant", "Hi."),
]
assert spend_logs == []
assert len(fake_prisma_client.db.calls) == litellm.constants.RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS
@pytest.mark.asyncio
@ -657,7 +607,9 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id():
model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282",
response_id=request_id,
)
fake_prisma_client = _FakePrismaClient(written_rows=[], queued_rows=[])
fake_prisma_client = _FakePrismaClient(
results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]]
)
with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(