diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4952dd28115..041d72a132d 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2,6 +2,7 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion API) """ +import os from collections.abc import Sequence from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast @@ -60,7 +61,26 @@ from litellm.types.utils import ( ) ########### Initialize Classes used for Responses API ########### -TOOL_CALLS_CACHE = InMemoryCache() +def _positive_int_env(name: str, default: int) -> int: + raw_value = os.getenv(name) + if raw_value is None: + return default + try: + parsed = int(raw_value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +TOOL_CALLS_CACHE = InMemoryCache( + max_size_in_memory=_positive_int_env( + "LITELLM_RESPONSES_TOOL_CALL_CACHE_MAX_ITEMS", 128 + ), + default_ttl=_positive_int_env("LITELLM_RESPONSES_TOOL_CALL_CACHE_TTL_SECONDS", 300), + max_size_per_item=_positive_int_env( + "LITELLM_RESPONSES_TOOL_CALL_CACHE_MAX_ITEM_KB", 256 + ), +) class ChatCompletionSession(TypedDict, total=False): diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index da8da1b486f..76e959c2b7e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import os import time import traceback from datetime import datetime @@ -30,6 +31,31 @@ from litellm.types.utils import CallTypes from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook +def _positive_int_env(name: str, default: int) -> int: + raw_value = os.getenv(name) + if raw_value is None: + return default + try: + parsed = int(raw_value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +_MANAGED_WS_HISTORY_MAX_RESPONSES = _positive_int_env( + "LITELLM_RESPONSES_WS_HISTORY_MAX_RESPONSES", 16 +) +_MANAGED_WS_HISTORY_MAX_MESSAGES = _positive_int_env( + "LITELLM_RESPONSES_WS_HISTORY_MAX_MESSAGES", 80 +) +_MANAGED_WS_HISTORY_MAX_BYTES = _positive_int_env( + "LITELLM_RESPONSES_WS_HISTORY_MAX_BYTES", 524_288 +) +_MANAGED_WS_HISTORY_TTL_SECONDS = _positive_int_env( + "LITELLM_RESPONSES_WS_HISTORY_TTL_SECONDS", 600 +) + + @lru_cache(maxsize=1) def _get_openai_response_types(): from litellm.types.llms import openai as openai_types @@ -1459,6 +1485,7 @@ class ManagedResponsesWebSocketHandler: # This avoids the async DB-write race condition where spend logs haven't # been committed yet when the next response.create arrives. self._session_history: Dict[str, List[Dict[str, Any]]] = {} + self._session_history_touched: Dict[str, float] = {} # ------------------------------------------------------------------ # Internal helpers @@ -1502,7 +1529,11 @@ class ManagedResponsesWebSocketHandler: previous_response_id ) raw_id = decoded.get("response_id", previous_response_id) - return list(self._session_history.get(raw_id, [])) + self._prune_session_history() + messages = self._session_history.get(raw_id, []) + if messages: + self._session_history_touched[raw_id] = time.time() + return list(messages) def _store_history(self, response_id: str, messages: List[Dict[str, Any]]) -> None: """ @@ -1511,7 +1542,62 @@ class ManagedResponsesWebSocketHandler: Replaces any prior value — callers are responsible for passing the full history (prior turns + current input + new output). """ - self._session_history[response_id] = messages + self._prune_session_history() + normalized_messages = self._trim_history_messages(messages) + if not normalized_messages: + return + self._session_history[response_id] = normalized_messages + self._session_history_touched[response_id] = time.time() + self._prune_session_history() + + def _delete_history(self, response_id: Optional[str]) -> None: + if not response_id: + return + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id) + raw_id = decoded.get("response_id", response_id) + self._session_history.pop(raw_id, None) + self._session_history_touched.pop(raw_id, None) + + @staticmethod + def _history_size_bytes(messages: List[Dict[str, Any]]) -> int: + try: + return len(json.dumps(messages, default=str, separators=(",", ":"))) + except Exception: + return sum(len(str(message)) for message in messages) + + def _trim_history_messages( + self, messages: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + normalized_messages = list(messages) + if len(normalized_messages) > _MANAGED_WS_HISTORY_MAX_MESSAGES: + normalized_messages = normalized_messages[-_MANAGED_WS_HISTORY_MAX_MESSAGES:] + + while ( + len(normalized_messages) > 1 + and self._history_size_bytes(normalized_messages) + > _MANAGED_WS_HISTORY_MAX_BYTES + ): + normalized_messages.pop(0) + return normalized_messages + + def _prune_session_history(self) -> None: + now = time.time() + expired_response_ids = [ + response_id + for response_id, touched_at in self._session_history_touched.items() + if now - touched_at > _MANAGED_WS_HISTORY_TTL_SECONDS + ] + for response_id in expired_response_ids: + self._delete_history(response_id) + + while len(self._session_history) > _MANAGED_WS_HISTORY_MAX_RESPONSES: + oldest_response_id = min( + self._session_history, + key=lambda response_id: self._session_history_touched.get( + response_id, 0.0 + ), + ) + self._delete_history(oldest_response_id) @staticmethod def _extract_response_id(completed_event: Dict[str, Any]) -> Optional[str]: @@ -1726,6 +1812,7 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, completed_event: Optional[Dict[str, Any]], + previous_response_id: Optional[str], prior_history: List[Dict[str, Any]], current_messages: List[Dict[str, Any]], ) -> None: @@ -1738,6 +1825,15 @@ class ManagedResponsesWebSocketHandler: output_msgs = self._extract_output_messages(completed_event) all_messages = prior_history + current_messages + output_msgs self._store_history(new_response_id, all_messages) + decoded_previous = ( + ResponsesAPIRequestUtils._decode_responses_api_response_id( + previous_response_id + ).get("response_id") + if previous_response_id + else None + ) + if decoded_previous and decoded_previous != new_response_id: + self._delete_history(decoded_previous) verbose_logger.debug( "ManagedResponsesWS: stored %d messages for response_id=%s", len(all_messages), @@ -1806,7 +1902,9 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(exc)) return - self._save_turn_history(completed_event, prior_history, current_messages) + self._save_turn_history( + completed_event, previous_response_id, prior_history, current_messages + ) # ------------------------------------------------------------------ # Main entry point @@ -1832,3 +1930,12 @@ class ManagedResponsesWebSocketHandler: except Exception as exc: verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc) await self._send_error(f"Internal server error: {exc}") + finally: + cleared = len(self._session_history) + self._session_history.clear() + self._session_history_touched.clear() + if cleared: + verbose_logger.debug( + "ManagedResponsesWS: cleared %d in-memory response histories", + cleared, + ) diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 1981651797d..fde5d5500dd 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -165,6 +165,90 @@ class TestManagedWebSocketHandlerIntegration: assert handler.timeout == 30.0 assert handler.custom_llm_provider == "test_provider" + def test_managed_handler_drops_superseded_turn_history(self): + """New response history supersedes the previous response key.""" + from unittest.mock import MagicMock + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + handler = ManagedResponsesWebSocketHandler( + websocket=MagicMock(), + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + ) + handler._store_history( + "first-response", + [{"type": "message", "role": "user", "content": "hello"}], + ) + + handler._save_turn_history( + { + "type": "response.completed", + "response": { + "id": "second-response", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "done"}], + } + ], + }, + }, + "first-response", + handler._get_history_messages("first-response"), + [{"type": "message", "role": "user", "content": "next"}], + ) + + assert "first-response" not in handler._session_history + assert "second-response" in handler._session_history + + @pytest.mark.asyncio + async def test_managed_handler_clears_history_on_disconnect(self): + """Per-connection response history is released when the WS run ends.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, + ) + + mock_websocket = MagicMock() + mock_websocket.receive_text = AsyncMock(side_effect=Exception("disconnect")) + handler = ManagedResponsesWebSocketHandler( + websocket=mock_websocket, + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + ) + handler._store_history( + "first-response", + [{"type": "message", "role": "user", "content": "hello"}], + ) + + await handler.run() + + assert handler._session_history == {} + assert handler._session_history_touched == {} + class TestChunkTransformation: """Test chunk serialization and transformation for WebSocket streaming"""