diff --git a/litellm/constants.py b/litellm/constants.py index c33e5a53b76..aaaddd063e7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 283194bad7c..0c8c9a853ec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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.") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9978fa04f40..86d954c0913 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 935c78bc9a1..15267533957 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -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 [] diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 76a80ac2651..ca1827aa38e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 54d59e690f9..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -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() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 926e9e0af2a..4fc288a47cb 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -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(