diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index fe224a52cfa..9bf4874ac80 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -2,9 +2,8 @@ Response Polling Handler for Background Responses with Cache """ -import json -from collections.abc import Mapping -from datetime import datetime, timezone +from collections.abc import Mapping, Sequence +from datetime import datetime, timedelta, timezone from typing import Final from pydantic import TypeAdapter @@ -13,22 +12,26 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid4 from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores import response_polling +from litellm.redis.stores.response_polling import ResponsePollingStore, TransportResponsePollingStore +from litellm.redis.transport.contract import RedisFailure from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus _RESPONSE_POLLING_TARGET: Final = "response_polling" -_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_LIST: Final = TypeAdapter(list[object]) class ResponsePollingHandler: """Handles polling-based responses with Redis cache""" - CACHE_KEY_PREFIX = "litellm:polling:response:" + CACHE_KEY_PREFIX = response_polling.KEY_PREFIX POLLING_ID_PREFIX = "litellm_poll_" # Clear prefix to identify polling IDs def __init__(self, redis_cache: RedisCache | None = None, ttl: int = 3600): - self.redis_cache = redis_cache - self.ttl = ttl # Time-to-live for cache entries (default: 1 hour) + self.ttl = ttl + self._store: Final[ResponsePollingStore | None] = ( + TransportResponsePollingStore(redis_cache.transport()) if redis_cache is not None else None + ) @classmethod def generate_polling_id(cls) -> str: @@ -43,7 +46,10 @@ class ResponsePollingHandler: @classmethod def get_cache_key(cls, polling_id: str) -> str: """Get Redis cache key for a polling ID""" - return f"{cls.CACHE_KEY_PREFIX}{polling_id}" + return response_polling.state_key(polling_id) + + def _ttl(self) -> timedelta: + return timedelta(seconds=self.ttl) @with_service_target(_RESPONSE_POLLING_TARGET) async def create_initial_state( @@ -79,16 +85,10 @@ class ResponsePollingHandler: } ) - cache_key: Final = self.get_cache_key(polling_id) - - if self.redis_cache: - # Store ResponsesAPIResponse directly in Redis - await self.redis_cache.async_set_cache( - key=cache_key, - value=response.model_dump_json(), # Pydantic v2 method - ttl=self.ttl, - ) - verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl) + if self._store is not None: + written: Final = await self._store.put(polling_id, response.model_dump(mode="json"), ttl=self._ttl()) + if written is None: + verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl) return response @@ -97,13 +97,13 @@ class ResponsePollingHandler: self, polling_id: str, status: ResponsesAPIStatus | None = None, - usage: dict | None = None, - error: dict | None = None, - incomplete_details: dict | None = None, - reasoning: dict | None = None, + usage: Mapping[str, object] | None = None, + error: Mapping[str, object] | None = None, + incomplete_details: Mapping[str, object] | None = None, + reasoning: Mapping[str, object] | None = None, tool_choice: object | None = None, - tools: list | None = None, - output: list | None = None, + tools: Sequence[object] | None = None, + output: Sequence[object] | None = None, # Additional ResponsesAPIResponse fields model: str | None = None, instructions: str | None = None, @@ -111,7 +111,7 @@ class ResponsePollingHandler: top_p: float | None = None, max_output_tokens: int | None = None, previous_response_id: str | None = None, - text: dict | None = None, + text: Mapping[str, object] | None = None, truncation: str | None = None, parallel_tool_calls: bool | None = None, user: str | None = None, @@ -145,79 +145,45 @@ class ResponsePollingHandler: user: User identifier store: Whether to store the response """ - if not self.redis_cache: + if self._store is None: return - cache_key: Final = self.get_cache_key(polling_id) - - # Get current state - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) - if not cached_state: + cached_state: Final = await self._store.get(polling_id) + if isinstance(cached_state, RedisFailure): + return + if cached_state is None: verbose_proxy_logger.warning("No cached state found for polling_id: %s", polling_id) return - # Parse existing ResponsesAPIResponse from cache - state: Final = _JSON_OBJECT.validate_python(json.loads(cached_state)) + updates: Final[Mapping[str, object]] = { + "status": "failed" if error else status or None, + "output": output, + "usage": usage or None, + "error": error or None, + "incomplete_details": incomplete_details or None, + "reasoning": reasoning, + "tool_choice": tool_choice, + "tools": tools, + "model": model, + "instructions": instructions, + "temperature": temperature, + "top_p": top_p, + "max_output_tokens": max_output_tokens, + "previous_response_id": previous_response_id, + "text": text, + "truncation": truncation, + "parallel_tool_calls": parallel_tool_calls, + "user": user, + "store": store, + } + state: Final[Mapping[str, object]] = { + **cached_state, + **{key: value for key, value in updates.items() if value is not None}, + } - # Update status (using OpenAI native status values) - if status: - state["status"] = status - - # Replace full output list if provided - if output is not None: - state["output"] = output - - # Update usage - if usage: - state["usage"] = usage - - # Handle error (sets status to OpenAI's "failed") - if error: - state["status"] = "failed" - state["error"] = error # Use OpenAI's 'error' field - - # Handle incomplete details - if incomplete_details: - state["incomplete_details"] = incomplete_details - - # Update reasoning, tool_choice, tools from response.completed - if reasoning is not None: - state["reasoning"] = reasoning - if tool_choice is not None: - state["tool_choice"] = tool_choice - if tools is not None: - state["tools"] = tools - - # Update additional ResponsesAPIResponse fields - if model is not None: - state["model"] = model - if instructions is not None: - state["instructions"] = instructions - if temperature is not None: - state["temperature"] = temperature - if top_p is not None: - state["top_p"] = top_p - if max_output_tokens is not None: - state["max_output_tokens"] = max_output_tokens - if previous_response_id is not None: - state["previous_response_id"] = previous_response_id - if text is not None: - state["text"] = text - if truncation is not None: - state["truncation"] = truncation - if parallel_tool_calls is not None: - state["parallel_tool_calls"] = parallel_tool_calls - if user is not None: - state["user"] = user - if store is not None: - state["store"] = store - - # Update cache with configured TTL - await self.redis_cache.async_set_cache( - key=cache_key, - value=json.dumps(state), - ttl=self.ttl, - ) + written: Final = await self._store.put(polling_id, state, ttl=self._ttl()) + if written is not None: + return output_count: Final = len(_JSON_LIST.validate_python(state.get("output", []))) verbose_proxy_logger.debug( @@ -227,16 +193,11 @@ class ResponsePollingHandler: @with_service_target(_RESPONSE_POLLING_TARGET) async def get_state(self, polling_id: str) -> Mapping[str, object] | None: """Get current polling state from Redis""" - if not self.redis_cache: + if self._store is None: return None - cache_key: Final = self.get_cache_key(polling_id) - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) - - if cached_state: - return _JSON_OBJECT.validate_python(json.loads(cached_state)) - - return None + cached_state: Final = await self._store.get(polling_id) + return None if isinstance(cached_state, RedisFailure) else cached_state async def cancel_polling(self, polling_id: str) -> bool: """ @@ -253,13 +214,10 @@ class ResponsePollingHandler: @with_service_target(_RESPONSE_POLLING_TARGET) async def delete_polling(self, polling_id: str) -> bool: """Delete a polling request from cache""" - if not self.redis_cache: + if self._store is None: return False - cache_key: Final = self.get_cache_key(polling_id) - # Use RedisCache's async_delete_cache method which handles Redis/RedisCluster - await self.redis_cache.async_delete_cache(cache_key) - return True + return await self._store.delete(polling_id) is None def should_use_polling_for_request( diff --git a/litellm/redis/stores/__init__.py b/litellm/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/stores/response_polling.py b/litellm/redis/stores/response_polling.py new file mode 100644 index 00000000000..df1f6d78736 --- /dev/null +++ b/litellm/redis/stores/response_polling.py @@ -0,0 +1,48 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +KEY_PREFIX: Final = "litellm:polling:response:" + +_DOCUMENT: Final = TypeAdapter(str) +_STATE: Final = TypeAdapter(dict[str, JsonValue]) + + +def state_key(polling_id: str) -> str: + return f"{KEY_PREFIX}{polling_id}" + + +class ResponsePollingStore(Protocol): + async def put(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ... + + async def get(self, polling_id: str) -> Mapping[str, JsonValue] | None | RedisFailure: ... + + async def delete(self, polling_id: str) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportResponsePollingStore: + transport: RedisTransport + + async def put(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: + written: Final = await self.transport.set(state_key(polling_id), json.dumps(json.dumps(state)), ttl=ttl) + return written if isinstance(written, RedisFailure) else None + + async def get(self, polling_id: str) -> Mapping[str, JsonValue] | None | RedisFailure: + raw: Final = await self.transport.get(state_key(polling_id)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + return _STATE.validate_json(_DOCUMENT.validate_json(raw)) + except ValidationError as e: + return RedisFailure(kind="decode", message=f"polling state is not a JSON object: {e}") + + async def delete(self, polling_id: str) -> None | RedisFailure: + deleted: Final = await self.transport.delete(state_key(polling_id)) + return deleted if isinstance(deleted, RedisFailure) else None diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 087b273548a..d7b9fe44885 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -20,11 +20,24 @@ from datetime import datetime, timezone from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch +import fakeredis import pytest from fastapi import Request from pydantic import ValidationError +from litellm.caching.redis_cache import RedisCache from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler +from tests._support.fake_redis import fake_redis_cache + + +async def _seed(cache: RedisCache, state: dict[str, object]) -> None: + await cache.async_set_cache( + key=ResponsePollingHandler.get_cache_key(str(state["id"])), value=json.dumps(state), ttl=3600 + ) + + +async def _legacy_read(cache: RedisCache, polling_id: str) -> dict[str, object]: + return json.loads(await cache.async_get_cache(ResponsePollingHandler.get_cache_key(polling_id))) from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -58,15 +71,9 @@ class TestResponsePollingHandler: # OpenAI format assert ResponsePollingHandler.is_polling_id("resp_abc123") is False # Anthropic format - assert ( - ResponsePollingHandler.is_polling_id("msg_01XFDUDYJgAACzvnptvVoYEL") - is False - ) + assert ResponsePollingHandler.is_polling_id("msg_01XFDUDYJgAACzvnptvVoYEL") is False # Generic UUID - assert ( - ResponsePollingHandler.is_polling_id("550e8400-e29b-41d4-a716-446655440000") - is False - ) + assert ResponsePollingHandler.is_polling_id("550e8400-e29b-41d4-a716-446655440000") is False def test_get_cache_key_format(self): """Test that cache keys have the correct format""" @@ -80,8 +87,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_create_initial_state_returns_queued_status(self): """Test that create_initial_state returns response with queued status""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer()), ttl=3600) polling_id = "litellm_poll_test123" request_data = { @@ -105,8 +111,9 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_create_initial_state_stores_in_redis(self): """Test that create_initial_state stores state in Redis with correct TTL""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=7200) + server = fakeredis.FakeServer() + cache = fake_redis_cache(server, namespace="ns") + handler = ResponsePollingHandler(redis_cache=cache, ttl=7200) polling_id = "litellm_poll_test123" request_data = {"model": "gpt-4o", "input": "Hello"} @@ -116,26 +123,18 @@ class TestResponsePollingHandler: request_data=request_data, ) - # Verify Redis was called with correct parameters - mock_redis.async_set_cache.assert_called_once() - call_args = mock_redis.async_set_cache.call_args + raw = fakeredis.FakeAsyncRedis(server=server) + assert await raw.keys("*") == [b"ns:litellm:polling:response:litellm_poll_test123"] + assert 7_190_000 < await raw.pttl("ns:litellm:polling:response:litellm_poll_test123") <= 7_200_000 - assert ( - call_args.kwargs["key"] == "litellm:polling:response:litellm_poll_test123" - ) - assert call_args.kwargs["ttl"] == 7200 - - # Verify the stored value is valid JSON - stored_value = call_args.kwargs["value"] - parsed = json.loads(stored_value) + parsed = await _legacy_read(cache, polling_id) assert parsed["id"] == polling_id assert parsed["status"] == "queued" @pytest.mark.asyncio async def test_create_initial_state_sets_created_at_timestamp(self): """Test that create_initial_state sets a valid created_at timestamp""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) before_time = int(datetime.now(timezone.utc).timestamp()) @@ -153,46 +152,45 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_changes_status_to_in_progress(self): """Test that update_state can change status to in_progress""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "queued", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=cache, ttl=3600) await handler.update_state( polling_id="litellm_poll_test", status="in_progress", ) - # Verify the update was saved - mock_redis.async_set_cache.assert_called_once() - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "in_progress" @pytest.mark.asyncio async def test_update_state_replaces_full_output_list(self): """Test that update_state replaces the full output list""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [{"id": "old_item", "type": "message"}], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=cache, ttl=3600) new_output = [ { @@ -212,8 +210,7 @@ class TestResponsePollingHandler: output=new_output, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert len(stored["output"]) == 2 assert stored["output"][0]["id"] == "item_1" @@ -222,18 +219,19 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_usage(self): """Test that update_state correctly stores usage data""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) usage_data = {"input_tokens": 10, "output_tokens": 50, "total_tokens": 60} @@ -243,8 +241,7 @@ class TestResponsePollingHandler: usage=usage_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "completed" assert stored["usage"] == usage_data @@ -252,24 +249,23 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_reasoning_tools_tool_choice(self): """Test that update_state stores reasoning, tools, and tool_choice from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) reasoning_data = {"effort": "medium", "summary": "Step by step analysis"} tool_choice_data = {"type": "function", "function": {"name": "get_weather"}} - tools_data = [ - {"type": "function", "function": {"name": "get_weather", "parameters": {}}} - ] + tools_data = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] await handler.update_state( polling_id="litellm_poll_test", @@ -279,8 +275,7 @@ class TestResponsePollingHandler: tools=tools_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["reasoning"] == reasoning_data assert stored["tool_choice"] == tool_choice_data @@ -289,18 +284,19 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_all_responses_api_fields(self): """Test that update_state stores all ResponsesAPIResponse fields from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) # All ResponsesAPIResponse fields that can be updated await handler.update_state( @@ -324,8 +320,7 @@ class TestResponsePollingHandler: incomplete_details={"reason": "max_output_tokens"}, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") # Verify all fields are stored correctly assert stored["status"] == "completed" @@ -353,8 +348,9 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_preserves_existing_fields(self): """Test that update_state preserves fields not being updated""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", @@ -363,10 +359,10 @@ class TestResponsePollingHandler: "created_at": 1234567890, "model": "gpt-4o", "temperature": 0.5, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) # Only update status await handler.update_state( @@ -374,8 +370,7 @@ class TestResponsePollingHandler: status="completed", ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") # Verify existing fields are preserved assert stored["status"] == "completed" @@ -386,18 +381,19 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_error_sets_failed_status(self): """Test that providing an error automatically sets status to failed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) error_data = { "type": "internal_error", @@ -407,30 +403,62 @@ class TestResponsePollingHandler: await handler.update_state( polling_id="litellm_poll_test", + status="completed", error=error_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "failed" assert stored["error"] == error_data + @pytest.mark.asyncio + async def test_update_state_preserves_false_empty_and_zero_values(self): + cache: Final = fake_redis_cache(fakeredis.FakeServer()) + await _seed(cache, {"id": "litellm_poll_test", "status": "in_progress", "usage": {"total_tokens": 3}}) + handler: Final = ResponsePollingHandler(redis_cache=cache) + + await handler.update_state( + polling_id="litellm_poll_test", + usage={}, + output=[], + tools=[], + reasoning={}, + instructions="", + temperature=0, + parallel_tool_calls=False, + store=False, + ) + + assert await _legacy_read(cache, "litellm_poll_test") == { + "id": "litellm_poll_test", + "status": "in_progress", + "usage": {"total_tokens": 3}, + "output": [], + "tools": [], + "reasoning": {}, + "instructions": "", + "temperature": 0, + "parallel_tool_calls": False, + "store": False, + } + @pytest.mark.asyncio async def test_update_state_with_incomplete_details(self): """Test that update_state stores incomplete_details""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) incomplete_details = {"reason": "max_output_tokens"} @@ -440,8 +468,7 @@ class TestResponsePollingHandler: incomplete_details=incomplete_details, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "incomplete" assert stored["incomplete_details"] == incomplete_details @@ -460,27 +487,23 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_handles_missing_cached_state(self): """Test that update_state handles case when cached state doesn't exist""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = None # Cache miss + server = fakeredis.FakeServer() + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(server)) - handler = ResponsePollingHandler(redis_cache=mock_redis) - - # Should not raise an exception await handler.update_state( polling_id="litellm_poll_test", status="in_progress", ) - # Should not try to set cache if nothing was found - mock_redis.async_set_cache.assert_not_called() + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [] # ==================== Get State Tests ==================== @pytest.mark.asyncio async def test_get_state_returns_cached_state(self): """Test that get_state returns the cached state""" - mock_redis = AsyncMock() - cached_state = { + cache = fake_redis_cache(fakeredis.FakeServer()) + cached_state: dict[str, object] = { "id": "litellm_poll_test", "object": "response", "status": "in_progress", @@ -488,9 +511,9 @@ class TestResponsePollingHandler: "created_at": 1234567890, "usage": {"input_tokens": 10, "output_tokens": 20}, } - mock_redis.async_get_cache.return_value = json.dumps(cached_state) + await _seed(cache, cached_state) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.get_state("litellm_poll_test") @@ -499,10 +522,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_get_state_returns_none_for_missing_state(self): """Test that get_state returns None when state doesn't exist""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = None - - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) result = await handler.get_state("litellm_poll_nonexistent") @@ -522,25 +542,25 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_cancel_polling_updates_status_to_cancelled(self): """Test that cancel_polling sets status to cancelled""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.cancel_polling("litellm_poll_test") assert result is True - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "cancelled" # ==================== Delete Polling Tests ==================== @@ -548,27 +568,38 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_delete_polling_removes_from_cache(self): """Test that delete_polling removes the entry from Redis""" - mock_redis = AsyncMock() - mock_async_client = AsyncMock() - mock_redis.redis_async_client = True # hasattr check - # init_async_client is a sync method that returns an async client - mock_redis.init_async_client = Mock(return_value=mock_async_client) - - # Mock async_delete_cache to actually call init_async_client and delete - async def mock_async_delete_cache(key): - client = mock_redis.init_async_client() - await client.delete(key) - - mock_redis.async_delete_cache = mock_async_delete_cache - - handler = ResponsePollingHandler(redis_cache=mock_redis) + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed(cache, {"id": "litellm_poll_test", "status": "queued"}) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.delete_polling("litellm_poll_test") assert result is True - mock_async_client.delete.assert_called_once_with( - "litellm:polling:response:litellm_poll_test" + assert await handler.get_state("litellm_poll_test") is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored", [b"not json", b'{"id": "single-encoded"}', b'"[1, 2]"']) + async def test_a_state_that_does_not_decode_reads_as_missing(self, stored: bytes): + server = fakeredis.FakeServer() + await fakeredis.FakeAsyncRedis(server=server).set( + ResponsePollingHandler.get_cache_key("litellm_poll_x"), stored ) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(server)) + + assert await handler.get_state("litellm_poll_x") is None + + @pytest.mark.asyncio + async def test_an_unreachable_redis_reads_as_missing_and_fails_the_delete(self): + server = fakeredis.FakeServer() + cache = fake_redis_cache(server) + await _seed(cache, {"id": "litellm_poll_test", "status": "queued"}) + handler = ResponsePollingHandler(redis_cache=cache) + server.connected = False + + assert await handler.get_state("litellm_poll_test") is None + assert await handler.delete_polling("litellm_poll_test") is False + queued = await handler.create_initial_state(polling_id="litellm_poll_new", request_data={}) + assert queued.status == "queued" @pytest.mark.asyncio async def test_delete_polling_returns_false_without_redis(self): @@ -596,26 +627,30 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_uses_configured_ttl(self): """Test that update_state uses the configured TTL""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + server = fakeredis.FakeServer() + cache = fake_redis_cache(server) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "queued", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=1800) + handler = ResponsePollingHandler(redis_cache=cache, ttl=1800) await handler.update_state( polling_id="litellm_poll_test", status="in_progress", ) - call_args = mock_redis.async_set_cache.call_args - assert call_args.kwargs["ttl"] == 1800 + ttl_ms = await fakeredis.FakeAsyncRedis(server=server).pttl( + ResponsePollingHandler.get_cache_key("litellm_poll_test") + ) + assert 1_790_000 < ttl_ms <= 1_800_000 class TestStreamingEventProcessing: @@ -689,6 +724,7 @@ class TestBackgroundStreamingModule: def test_background_streaming_task_is_async(self): """Test that background_streaming_task is an async function""" import inspect + from litellm.proxy.response_polling.background_streaming import ( background_streaming_task, ) @@ -886,9 +922,7 @@ class TestProviderResolutionForPolling: model_name_to_deployment_indices = { "gpt-5": [0], } - model_list = [ - {"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}} - ] + model_list = [{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}}] model = "unknown-model" # Not in router polling_via_cache_enabled = ["openai"] @@ -1282,9 +1316,7 @@ class TestStreamingEventParsing: content_list = output_items[item_id]["content"] if content_index < len(content_list): if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[ - key - ] + content_list[content_index]["text"] = accumulated_text[key] assert accumulated_text[("item_123", 0)] == "Hello World!" assert output_items["item_123"]["content"][0]["text"] == "Hello World!" @@ -1346,9 +1378,7 @@ class TestStreamingEventParsing: def test_parse_sse_format(self): """Test parsing Server-Sent Events format""" - raw_chunk = ( - b'data: {"type": "response.output_item.added", "item": {"id": "123"}}' - ) + raw_chunk = b'data: {"type": "response.output_item.added", "item": {"id": "123"}}' # Decode bytes to string if isinstance(raw_chunk, bytes): @@ -1476,9 +1506,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) # Find the final update_state call (last one) @@ -1515,9 +1543,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1558,17 +1584,13 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] assert final_call.kwargs["status"] == "incomplete" assert final_call.kwargs["error"] == error_payload - assert final_call.kwargs["incomplete_details"] == { - "reason": "max_output_tokens" - } + assert final_call.kwargs["incomplete_details"] == {"reason": "max_output_tokens"} assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 4096} @pytest.mark.asyncio @@ -1597,9 +1619,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1632,9 +1652,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1699,9 +1717,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1725,9 +1741,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1763,9 +1777,7 @@ class TestBackgroundStreamingTerminalEvents: yield f"data: {json.dumps(event)}\n\n" async def upstream_call_behind_disconnect_guard(**kwargs): - return await create_response( - slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] - ) + return await create_response(slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"]) handler = AsyncMock(spec=ResponsePollingHandler) kwargs = _make_background_streaming_kwargs("poll_7", handler) @@ -1838,8 +1850,7 @@ class TestEdgeCases: @pytest.mark.asyncio async def test_create_initial_state_with_empty_metadata(self): """Test create_initial_state handles missing metadata gracefully""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) response = await handler.create_initial_state( polling_id="litellm_poll_test", @@ -1851,26 +1862,26 @@ class TestEdgeCases: @pytest.mark.asyncio async def test_update_state_with_none_output_clears_output(self): """Test that output=[] explicitly sets empty output""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [{"id": "item_1"}], # Has existing output "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) await handler.update_state( polling_id="litellm_poll_test", output=[], # Explicitly set empty ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["output"] == []