Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/response_polling/polling_handler.py: 14%
128 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Response Polling Handler for Background Responses with Cache
3"""
5import json
6from datetime import datetime, timezone
7from typing import Any, Final
9from litellm._logging import verbose_proxy_logger
10from litellm._uuid import uuid4
11from litellm.caching.redis_cache import RedisCache
12from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus
15class ResponsePollingHandler:
16 """Handles polling-based responses with Redis cache"""
18 CACHE_KEY_PREFIX = "litellm:polling:response:"
19 POLLING_ID_PREFIX = "litellm_poll_" # Clear prefix to identify polling IDs
21 def __init__(self, redis_cache: RedisCache | None = None, ttl: int = 3600):
22 self.redis_cache = redis_cache
23 self.ttl = ttl # Time-to-live for cache entries (default: 1 hour)
25 @classmethod
26 def generate_polling_id(cls) -> str:
27 """Generate a unique UUID for polling with clear prefix"""
28 return f"{cls.POLLING_ID_PREFIX}{uuid4()}"
30 @classmethod
31 def is_polling_id(cls, response_id: str) -> bool:
32 """Check if a response_id is a polling ID"""
33 return response_id.startswith(cls.POLLING_ID_PREFIX)
35 @classmethod
36 def get_cache_key(cls, polling_id: str) -> str:
37 """Get Redis cache key for a polling ID"""
38 return f"{cls.CACHE_KEY_PREFIX}{polling_id}"
40 async def create_initial_state(
41 self,
42 polling_id: str,
43 request_data: dict[str, Any],
44 ) -> ResponsesAPIResponse:
45 """
46 Create initial state in Redis for a polling request
48 Uses OpenAI ResponsesAPIResponse object:
49 https://platform.openai.com/docs/api-reference/responses/object
51 Args:
52 polling_id: Unique identifier for this polling request
53 request_data: Original request data
55 Returns:
56 ResponsesAPIResponse object following OpenAI spec
57 """
58 created_timestamp: Final = int(datetime.now(timezone.utc).timestamp())
60 # Create OpenAI-compliant response object
61 response: Final = ResponsesAPIResponse(
62 id=polling_id,
63 object="response",
64 status="queued", # OpenAI native status
65 created_at=created_timestamp,
66 output=[],
67 metadata=request_data.get("metadata", {}),
68 usage=None,
69 )
71 cache_key: Final = self.get_cache_key(polling_id)
73 if self.redis_cache:
74 # Store ResponsesAPIResponse directly in Redis
75 await self.redis_cache.async_set_cache(
76 key=cache_key,
77 value=response.model_dump_json(), # Pydantic v2 method
78 ttl=self.ttl,
79 )
80 verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl)
82 return response
84 async def update_state(
85 self,
86 polling_id: str,
87 status: ResponsesAPIStatus | None = None,
88 usage: dict | None = None,
89 error: dict | None = None,
90 incomplete_details: dict | None = None,
91 reasoning: dict | None = None,
92 tool_choice: object | None = None,
93 tools: list | None = None,
94 output: list | None = None,
95 # Additional ResponsesAPIResponse fields
96 model: str | None = None,
97 instructions: str | None = None,
98 temperature: float | None = None,
99 top_p: float | None = None,
100 max_output_tokens: int | None = None,
101 previous_response_id: str | None = None,
102 text: dict | None = None,
103 truncation: str | None = None,
104 parallel_tool_calls: bool | None = None,
105 user: str | None = None,
106 store: bool | None = None,
107 ) -> None:
108 """
109 Update the polling state in Redis
111 Uses OpenAI Response object format with native status types:
112 https://platform.openai.com/docs/api-reference/responses/object
114 Args:
115 polling_id: Unique identifier for this polling request
116 status: OpenAI ResponsesAPIStatus value
117 usage: Usage information
118 error: Error dict (automatically sets status to "failed")
119 incomplete_details: Details for incomplete responses
120 reasoning: Reasoning configuration from response.completed
121 tool_choice: Tool choice configuration from response.completed
122 tools: Tools list from response.completed
123 output: Full output list to replace current output
124 model: Model identifier
125 instructions: System instructions
126 temperature: Sampling temperature
127 top_p: Nucleus sampling parameter
128 max_output_tokens: Maximum output tokens
129 previous_response_id: ID of previous response in conversation
130 text: Text configuration
131 truncation: Truncation setting
132 parallel_tool_calls: Whether parallel tool calls are enabled
133 user: User identifier
134 store: Whether to store the response
135 """
136 if not self.redis_cache:
137 return
139 cache_key: Final = self.get_cache_key(polling_id)
141 # Get current state
142 cached_state: Final = await self.redis_cache.async_get_cache(cache_key)
143 if not cached_state:
144 verbose_proxy_logger.warning("No cached state found for polling_id: %s", polling_id)
145 return
147 # Parse existing ResponsesAPIResponse from cache
148 state: Final = json.loads(cached_state)
150 # Update status (using OpenAI native status values)
151 if status:
152 state["status"] = status
154 # Replace full output list if provided
155 if output is not None:
156 state["output"] = output
158 # Update usage
159 if usage:
160 state["usage"] = usage
162 # Handle error (sets status to OpenAI's "failed")
163 if error:
164 state["status"] = "failed"
165 state["error"] = error # Use OpenAI's 'error' field
167 # Handle incomplete details
168 if incomplete_details:
169 state["incomplete_details"] = incomplete_details
171 # Update reasoning, tool_choice, tools from response.completed
172 if reasoning is not None:
173 state["reasoning"] = reasoning
174 if tool_choice is not None:
175 state["tool_choice"] = tool_choice
176 if tools is not None:
177 state["tools"] = tools
179 # Update additional ResponsesAPIResponse fields
180 if model is not None:
181 state["model"] = model
182 if instructions is not None:
183 state["instructions"] = instructions
184 if temperature is not None:
185 state["temperature"] = temperature
186 if top_p is not None:
187 state["top_p"] = top_p
188 if max_output_tokens is not None:
189 state["max_output_tokens"] = max_output_tokens
190 if previous_response_id is not None:
191 state["previous_response_id"] = previous_response_id
192 if text is not None:
193 state["text"] = text
194 if truncation is not None:
195 state["truncation"] = truncation
196 if parallel_tool_calls is not None:
197 state["parallel_tool_calls"] = parallel_tool_calls
198 if user is not None:
199 state["user"] = user
200 if store is not None:
201 state["store"] = store
203 # Update cache with configured TTL
204 await self.redis_cache.async_set_cache(
205 key=cache_key,
206 value=json.dumps(state),
207 ttl=self.ttl,
208 )
210 output_count: Final = len(state.get("output", []))
211 verbose_proxy_logger.debug(
212 "Updated polling state for %s: status=%s, output_items=%s", polling_id, state["status"], output_count
213 )
215 async def get_state(self, polling_id: str) -> dict[str, Any] | None:
216 """Get current polling state from Redis"""
217 if not self.redis_cache:
218 return None
220 cache_key: Final = self.get_cache_key(polling_id)
221 cached_state: Final = await self.redis_cache.async_get_cache(cache_key)
223 if cached_state:
224 return json.loads(cached_state)
226 return None
228 async def cancel_polling(self, polling_id: str) -> bool:
229 """
230 Cancel a polling request
232 Following OpenAI Response object format for cancelled status
233 """
234 await self.update_state(
235 polling_id=polling_id,
236 status="cancelled",
237 )
238 return True
240 async def delete_polling(self, polling_id: str) -> bool:
241 """Delete a polling request from cache"""
242 if not self.redis_cache:
243 return False
245 cache_key: Final = self.get_cache_key(polling_id)
246 # Use RedisCache's async_delete_cache method which handles Redis/RedisCluster
247 await self.redis_cache.async_delete_cache(cache_key)
248 return True
251def should_use_polling_for_request(
252 background_mode: bool,
253 polling_via_cache_enabled, # Can be False, "all", or List[str]
254 redis_cache, # RedisCache or None
255 model: str,
256 llm_router, # Router instance or None
257 native_background_mode: list[str] | None = None, # List of models that should use native background mode
258) -> bool:
259 """
260 Determine if polling via cache should be used for a request.
262 Args:
263 background_mode: Whether background=true was set in the request
264 polling_via_cache_enabled: Config value - False, "all", or list of providers
265 redis_cache: Redis cache instance (required for polling)
266 model: Model name from the request (e.g., "gpt-5" or "openai/gpt-4o")
267 llm_router: LiteLLM router instance for looking up model deployments
268 native_background_mode: List of model names that should use native provider
269 background mode instead of polling via cache
271 Returns:
272 True if polling should be used, False otherwise
273 """
274 # All conditions must be met
275 if not (background_mode and polling_via_cache_enabled and redis_cache): 275 ↛ 279line 275 didn't jump to line 279 because the condition on line 275 was always true
276 return False
278 # Check if model is in native_background_mode list - these use native provider background mode
279 if native_background_mode and model in native_background_mode:
280 verbose_proxy_logger.debug("Model %s is in native_background_mode list, skipping polling via cache", model)
281 return False
283 # "all" enables polling for all providers
284 if polling_via_cache_enabled == "all":
285 return True
287 # Check if provider is in the enabled list
288 if isinstance(polling_via_cache_enabled, list):
289 # First, try to get provider from model string format "provider/model"
290 if "/" in model:
291 provider: Final = model.split("/")[0]
292 if provider in polling_via_cache_enabled:
293 return True
294 # Otherwise, check ALL deployments for this model_name in router
295 elif llm_router is not None:
296 try:
297 # Get all deployment indices for this model name
298 indices: Final = llm_router.model_name_to_deployment_indices.get(model, [])
299 for idx in indices:
300 deployment_dict = llm_router.model_list[idx]
301 litellm_params = deployment_dict.get("litellm_params", {})
303 # Check custom_llm_provider first
304 dep_provider = litellm_params.get("custom_llm_provider")
306 # Then try to extract from model (e.g., "openai/gpt-5")
307 if not dep_provider:
308 dep_model = litellm_params.get("model", "")
309 if "/" in dep_model:
310 dep_provider = dep_model.split("/")[0]
312 # If ANY deployment's provider matches, enable polling
313 if dep_provider and dep_provider in polling_via_cache_enabled:
314 verbose_proxy_logger.debug("Polling enabled for model=%s, provider=%s", model, dep_provider)
315 return True
316 except Exception as e:
317 verbose_proxy_logger.debug("Could not resolve provider for model %s: %s", model, e)
319 return False