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 c616d9e8723..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,6 +6006,26 @@ async def enqueue_spend_logs( ) +def request_spend_log_flush() -> None: + """Wake the queue monitor now rather than leaving the rows for its next poll. + + 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. + """ + 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]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. @@ -6450,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 a53d7c68b0a..1566bb1bdd7 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,3 +1,4 @@ +import asyncio import json from typing import TYPE_CHECKING, Any, Final, cast @@ -104,10 +105,13 @@ class ResponsesSessionHandler: if proxy_server_request_dict: _response_input_param: Final = proxy_server_request_dict.get("input", None) _messages = proxy_server_request_dict.get("messages", None) - if isinstance(_response_input_param, str): + if isinstance(_response_input_param, (str, list)): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast(ResponseInputParam, _response_input_param) + response_input_param = cast( + ResponseInputParam, + [_response_input_param], # mutable-ok: a lone input item still has to arrive as a list + ) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -256,13 +260,22 @@ 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. Deployments that write no spend + logs at all have nothing to wait for, so they keep the single original query. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.constants import ( + RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS, + RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL, + ) + from litellm.proxy.proxy_server import disable_spend_logs, prisma_client verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) decoded_response_id: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) - previous_response_id = decoded_response_id.get("response_id", previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) if prisma_client is None: return [] @@ -278,12 +291,17 @@ class ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - spend_logs: Final = await prisma_client.db.query_raw(query, previous_response_id) + max_attempts: Final = 1 if disable_spend_logs else RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS + for attempt in range(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", - previous_response_id, - json.dumps(spend_logs, indent=4, default=str), - ) - - return spend_logs + verbose_proxy_logger.debug("Found no spend logs for previous response id %s", response_id) + return [] # mutable-ok: an empty result the caller only reads diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 67be8dedd55..92bbca9ee5b 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -102,6 +102,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.final_text: str = "" self._cached_item_id: str | None = None self._cached_response_id: str | None = None + self._buffered_chunk: ModelResponseStream | None = None + self._upstream_exhausted: bool = False + self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} @@ -346,6 +349,59 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) self._pending_tool_events.append(item_done_event) + def _adopt_response_id_from_chunk(self, chunk: ModelResponseStream) -> None: + if self._cached_response_id is not None: + return + chunk_id: Final = getattr(chunk, "id", None) + if chunk_id and isinstance(chunk_id, str): + self._cached_response_id = chunk_id + + async def _aprime_response_id(self) -> None: + """ + Pull the first upstream chunk before `response.created` is emitted so every event + carries the chat completion id that spend tracking stores as `request_id`. + """ + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = await self.litellm_custom_stream_wrapper.__anext__() + except StopAsyncIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _prime_response_id(self) -> None: + if self._response_id_primed: + return + self._response_id_primed = True + while True: + try: + chunk = self.litellm_custom_stream_wrapper.__next__() + except StopIteration: + self._upstream_exhausted = True + return + if chunk is not None: + self._buffered_chunk = chunk + self._adopt_response_id_from_chunk(chunk) + return + + def _take_buffered_chunk(self) -> ModelResponseStream | None: + buffered: Final = self._buffered_chunk + self._buffered_chunk = None + return buffered + + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: + return ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response, + custom_llm_provider=self.custom_llm_provider, + litellm_metadata=self.litellm_metadata, + ) + def _default_response_created_event_data(self) -> dict: # Use cached response ID if available, otherwise generate a new one if self._cached_response_id is None: @@ -404,7 +460,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, - response=ResponsesAPIResponse(**response_created_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -415,7 +471,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 event: Final = ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, - response=ResponsesAPIResponse(**response_in_progress_event_data), + response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)), ) event.__dict__["sequence_number"] = self._sequence_number return event @@ -827,6 +883,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self.finished is True: raise StopAsyncIteration + await self._aprime_response_id() result = self.return_default_initial_events() if result: return result @@ -838,7 +895,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._pending_tool_events.pop(0) try: - chunk = await self.litellm_custom_stream_wrapper.__anext__() + chunk = self._take_buffered_chunk() + if chunk is None: + if self._upstream_exhausted: + raise StopAsyncIteration + chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) self._ensure_output_item_for_chunk(chunk) @@ -929,6 +990,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): while True: if self.finished is True: raise StopIteration + self._prime_response_id() result = self.return_default_initial_events() if result: return result @@ -939,7 +1001,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._pending_tool_events: return self._pending_tool_events.pop(0) try: - chunk = self.litellm_custom_stream_wrapper.__next__() + buffered_chunk = self._take_buffered_chunk() + if buffered_chunk is not None: + chunk = buffered_chunk + elif self._upstream_exhausted: + raise StopIteration + else: + chunk = self.litellm_custom_stream_wrapper.__next__() self._ensure_output_item_for_chunk(chunk) # Accumulate provider_specific_fields from chunk and delta for src in ( @@ -1117,11 +1185,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): responses_api_response.output = list(self._output_with_streamed_item_ids(responses_api_response)) # Encode the response ID to match non-streaming behavior - encoded_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=responses_api_response, - custom_llm_provider=self.custom_llm_provider, - litellm_metadata=self.litellm_metadata, - ) + encoded_response: Final = self._with_encoded_response_id(responses_api_response) return ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, 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 19f240fa3d4..df477f6d01e 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 @@ -10,6 +10,7 @@ from litellm.responses.litellm_completion_transformation import session_handler from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) +from litellm.responses.utils import ResponsesAPIRequestUtils @pytest.mark.asyncio @@ -430,3 +431,210 @@ async def test_get_chat_completion_message_history_empty_response_dict(): # Verify the session was still created correctly assert result["litellm_session_id"] == "test-session" + + +def _chat_completion_response(request_id: str, content: str) -> dict: + return { + "id": request_id, + "object": "chat.completion", + "created": 1748575031, + "model": "claude-haiku-4-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + } + + +class _FakePrismaDB: + def __init__(self, results): + self._results = list(results) + self.calls = [] + + async def query_raw(self, query, *args): + self.calls.append(args) + if not self._results: + return [] + return list(self._results.pop(0)) + + +class _FakePrismaClient: + 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 +async def test_message_history_reconstructs_list_shaped_input(): + """ + The Responses API sends `input` as a list of items, which is what lands in the stored + proxy_server_request. The user turns have to survive session reconstruction. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + mock_spend_logs = [ + _spend_log( + request_id, + "a96757c4-c6dc-4c76-b37e-e7dfa526b701", + "Remember this: my favorite color is chartreuse.", + "OK", + ) + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs: + mock_get_spend_logs.return_value = mock_spend_logs + + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + assert result["litellm_session_id"] == "a96757c4-c6dc-4c76-b37e-e7dfa526b701" + + +@pytest.mark.asyncio +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 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" + 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( + request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "Remember this: my favorite color is chartreuse."), + ("assistant", "OK"), + ] + 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_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" + fake_prisma_client = _FakePrismaClient( + 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): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + second_request_id + ) + + messages = result["messages"] + assert [(message.get("role"), message.get("content")) for message in messages] == [ + ("user", "My favorite color is chartreuse."), + ("assistant", "Got it."), + ("user", "And my favorite city is Lisbon."), + ("assistant", "Noted."), + ] + assert result["litellm_session_id"] == session_id + + +@pytest.mark.asyncio +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): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" + ) + + assert spend_logs == [] + assert len(fake_prisma_client.db.calls) == litellm.constants.RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS + + +@pytest.mark.asyncio +async def test_message_history_looks_up_the_decoded_chat_completion_id(): + """ + A `previous_response_id` handed back by the proxy is base64 encoded; spend logs store + the bare chat completion id, so that is what the lookup has to query on. + """ + request_id = "chatcmpl-935b8dad-fdc2-466e-a8ca-e26e5a8a21bb" + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="anthropic", + model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", + response_id=request_id, + ) + 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( + encoded_response_id + ) + + assert fake_prisma_client.db.calls == [(request_id,)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( + instant_session_lookup_retries: None, +): + """ + A deployment that writes no spend logs has nothing to wait for, so the miss path keeps + the single query it always had. + """ + fake_prisma_client = _FakePrismaClient(results=[]) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( + "litellm.proxy.proxy_server.disable_spend_logs", True + ): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + "chatcmpl-does-not-exist" + ) + + assert spend_logs == [] + assert fake_prisma_client.db.calls == [("chatcmpl-does-not-exist",)] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py similarity index 76% rename from tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py rename to tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index fa6f42609ca..823f656ddc5 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1,18 +1,23 @@ """ -Tests for streaming tool-calls in Responses API transformation. +Tests for the Responses API streaming bridge in +litellm/responses/litellm_completion_transformation/streaming_iterator.py. Ensures that when the underlying chat-completions stream includes tool_calls deltas, LiteLLM emits Responses API streaming events (output_item.added + function_call_arguments.*). Also ensures that tool calls that only appear in the final built response still get emitted -before response.completed. +before response.completed, and that every event of a bridged stream carries the response id +spend tracking stores, so a follow-up previous_response_id still finds the conversation. """ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock + +import pytest from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import ( Delta, @@ -21,6 +26,68 @@ from litellm.types.utils import ( StreamingChoices, ) +CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" +RESPONSE_ID_EVENT_TYPES = frozenset( + {"response.created", "response.in_progress", "response.completed"} +) + + +def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-4-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +class _FakeStreamWrapper: + def __init__(self, chunks): + self._chunks = list(chunks) + self.logging_obj = MagicMock() + + def __iter__(self): + return self + + def __next__(self): + if not self._chunks: + raise StopIteration + return self._chunks.pop(0) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: + return LiteLLMCompletionStreamingIterator( + model="claude-haiku-4-5", + litellm_custom_stream_wrapper=_FakeStreamWrapper(chunks), + request_input="What is the weather in San Francisco?", + responses_api_request={}, + custom_llm_provider="anthropic", + litellm_metadata={}, + ) + + +def _response_ids(events) -> list[str]: + return [ + event.response.id + for event in events + if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES + ] + def test_tool_call_delta_is_emitted_as_responses_events(): iterator = LiteLLMCompletionStreamingIterator( @@ -397,3 +464,62 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): assert arguments_by_call_id["call_b"] == '{"b":' assert arguments_by_call_id["call_a"] != '{"a":1}' assert arguments_by_call_id["call_b"] != '{"b":1}' + + +@pytest.mark.asyncio +async def test_streaming_events_share_the_chat_completion_response_id(): + """ + Every event of a bridged stream has to carry the same id, and that id has to decode + to the chat completion id spend tracking stores as `request_id`. Otherwise a + follow-up `previous_response_id` matches no session and the conversation is dropped. + """ + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0]) + assert decoded["response_id"] == CHAT_COMPLETION_ID + assert decoded["custom_llm_provider"] == "anthropic" + + +def test_sync_streaming_events_share_the_chat_completion_response_id(): + iterator = _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = list(iterator) + + response_ids = _response_ids(events) + assert len(response_ids) == 3 + assert len(set(response_ids)) == 1 + assert ( + ResponsesAPIRequestUtils._decode_responses_api_response_id(response_ids[0])["response_id"] + == CHAT_COMPLETION_ID + ) + + +@pytest.mark.asyncio +async def test_streaming_emits_every_chunk_after_priming_the_response_id(): + iterator = _build_iterator( + [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] + ) + + events = [event async for event in iterator] + + deltas = "".join( + event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert deltas == "Hello!" + + +@pytest.mark.asyncio +async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): + iterator = _build_iterator([]) + + events = [event async for event in iterator] + + response_ids = _response_ids(events) + assert response_ids + assert len(set(response_ids)) == 1 + assert response_ids[0].startswith("resp_")