mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #37956 from BerriAI/litellm_fix_26167_bridged_session_lookup
fix(responses): keep the conversation when chaining previous_response_id on the bridge
This commit is contained in:
Mateo Wang
GitHub
commit
aae36f4bd4
9 changed files with 549 additions and 27 deletions
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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",)]
|
||||
|
|
|
|||
|
|
@ -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_")
|
||||
Loading…
Add table
Reference in a new issue