mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
0f4ac2143b
With polling via cache on, POST /v1/responses/{id}/cancel only wrote the
cancelled status into Redis. The background task kept streaming from the
provider, and its next state write merged over the cancel, so the response
came back as completed with full output and full billing
The background task now races the stream against a watcher that reads the
stored state every second and stops the stream when it reads cancelled, on
any worker or pod. Every state write goes through a Lua script that only
sets the key when it exists and is not cancelled, so a write already in
flight when the cancel lands can no longer overwrite it
355 lines
13 KiB
Python
355 lines
13 KiB
Python
"""
|
|
Response Polling Handler for Background Responses with Cache
|
|
"""
|
|
|
|
import json
|
|
from collections.abc import Mapping, Sequence
|
|
from datetime import datetime, timezone
|
|
from typing import Final
|
|
|
|
from pydantic import TypeAdapter
|
|
|
|
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.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])
|
|
_SCRIPT_REPLY: Final = TypeAdapter(int)
|
|
_WRITTEN: Final = 1
|
|
|
|
_WRITE_UNLESS_CANCELLED_SCRIPT: Final = """
|
|
local current = redis.call('GET', KEYS[1])
|
|
if current == false then
|
|
return 0
|
|
end
|
|
local decoded, stored = pcall(cjson.decode, current)
|
|
if decoded and type(stored) == 'string' then
|
|
decoded, stored = pcall(cjson.decode, stored)
|
|
end
|
|
if decoded and type(stored) == 'table' and stored['status'] == 'cancelled' then
|
|
return 0
|
|
end
|
|
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
|
|
return 1
|
|
"""
|
|
|
|
|
|
class ResponsePollingHandler:
|
|
"""Handles polling-based responses with Redis cache"""
|
|
|
|
CACHE_KEY_PREFIX = "litellm:polling:response:"
|
|
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)
|
|
|
|
@classmethod
|
|
def generate_polling_id(cls) -> str:
|
|
"""Generate a unique UUID for polling with clear prefix"""
|
|
return f"{cls.POLLING_ID_PREFIX}{uuid4()}"
|
|
|
|
@classmethod
|
|
def is_polling_id(cls, response_id: str) -> bool:
|
|
"""Check if a response_id is a polling ID"""
|
|
return response_id.startswith(cls.POLLING_ID_PREFIX)
|
|
|
|
@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}"
|
|
|
|
@with_service_target(_RESPONSE_POLLING_TARGET)
|
|
async def create_initial_state(
|
|
self,
|
|
polling_id: str,
|
|
request_data: Mapping[str, object],
|
|
) -> ResponsesAPIResponse:
|
|
"""
|
|
Create initial state in Redis for a polling request
|
|
|
|
Uses OpenAI ResponsesAPIResponse object:
|
|
https://platform.openai.com/docs/api-reference/responses/object
|
|
|
|
Args:
|
|
polling_id: Unique identifier for this polling request
|
|
request_data: Original request data
|
|
|
|
Returns:
|
|
ResponsesAPIResponse object following OpenAI spec
|
|
"""
|
|
created_timestamp: Final = int(datetime.now(timezone.utc).timestamp())
|
|
|
|
# Create OpenAI-compliant response object
|
|
response: Final = ResponsesAPIResponse.model_validate(
|
|
{
|
|
"id": polling_id,
|
|
"object": "response",
|
|
"status": "queued", # OpenAI native status
|
|
"created_at": created_timestamp,
|
|
"output": [],
|
|
"metadata": request_data.get("metadata", {}),
|
|
"usage": None,
|
|
}
|
|
)
|
|
|
|
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)
|
|
|
|
return response
|
|
|
|
@with_service_target(_RESPONSE_POLLING_TARGET)
|
|
async def update_state(
|
|
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,
|
|
tool_choice: object | None = None,
|
|
tools: Sequence[object] | None = None,
|
|
output: list | None = None,
|
|
# Additional ResponsesAPIResponse fields
|
|
model: str | None = None,
|
|
instructions: str | None = None,
|
|
temperature: float | None = None,
|
|
top_p: float | None = None,
|
|
max_output_tokens: int | None = None,
|
|
previous_response_id: str | None = None,
|
|
text: dict | None = None,
|
|
truncation: str | None = None,
|
|
parallel_tool_calls: bool | None = None,
|
|
user: str | None = None,
|
|
store: bool | None = None,
|
|
) -> None:
|
|
"""
|
|
Update the polling state in Redis
|
|
|
|
Uses OpenAI Response object format with native status types:
|
|
https://platform.openai.com/docs/api-reference/responses/object
|
|
|
|
Args:
|
|
polling_id: Unique identifier for this polling request
|
|
status: OpenAI ResponsesAPIStatus value
|
|
usage: Usage information
|
|
error: Error dict (automatically sets status to "failed")
|
|
incomplete_details: Details for incomplete responses
|
|
reasoning: Reasoning configuration from response.completed
|
|
tool_choice: Tool choice configuration from response.completed
|
|
tools: Tools list from response.completed
|
|
output: Full output list to replace current output
|
|
model: Model identifier
|
|
instructions: System instructions
|
|
temperature: Sampling temperature
|
|
top_p: Nucleus sampling parameter
|
|
max_output_tokens: Maximum output tokens
|
|
previous_response_id: ID of previous response in conversation
|
|
text: Text configuration
|
|
truncation: Truncation setting
|
|
parallel_tool_calls: Whether parallel tool calls are enabled
|
|
user: User identifier
|
|
store: Whether to store the response
|
|
"""
|
|
if not self.redis_cache:
|
|
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:
|
|
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))
|
|
|
|
# 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
|
|
|
|
write_unless_cancelled: Final = self.redis_cache.async_register_script(_WRITE_UNLESS_CANCELLED_SCRIPT)
|
|
try:
|
|
written: Final = _SCRIPT_REPLY.validate_python(
|
|
await write_unless_cancelled(keys=(cache_key,), args=(json.dumps(json.dumps(state)), int(self.ttl)))
|
|
)
|
|
except Exception as write_error: # noqa: BLE001 # a Redis fault must not end the background stream
|
|
verbose_proxy_logger.error("Failed to write the polling state for %s: %s", polling_id, write_error)
|
|
return
|
|
if written != _WRITTEN:
|
|
verbose_proxy_logger.debug(
|
|
"Skipped the polling state write for %s: the response was cancelled or has expired", polling_id
|
|
)
|
|
return
|
|
|
|
output_count: Final = len(_JSON_LIST.validate_python(state.get("output", [])))
|
|
verbose_proxy_logger.debug(
|
|
"Updated polling state for %s: status=%s, output_items=%s", polling_id, state["status"], output_count
|
|
)
|
|
|
|
@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:
|
|
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
|
|
|
|
async def cancel_polling(self, polling_id: str) -> bool:
|
|
"""
|
|
Cancel a polling request
|
|
|
|
Following OpenAI Response object format for cancelled status
|
|
"""
|
|
await self.update_state(
|
|
polling_id=polling_id,
|
|
status="cancelled",
|
|
)
|
|
return True
|
|
|
|
@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:
|
|
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
|
|
|
|
|
|
def should_use_polling_for_request(
|
|
background_mode: bool,
|
|
polling_via_cache_enabled, # Can be False, "all", or List[str]
|
|
redis_cache, # RedisCache or None
|
|
model: str,
|
|
llm_router, # Router instance or None
|
|
native_background_mode: Sequence[str] | None = None, # List of models that should use native background mode
|
|
) -> bool:
|
|
"""
|
|
Determine if polling via cache should be used for a request.
|
|
|
|
Args:
|
|
background_mode: Whether background=true was set in the request
|
|
polling_via_cache_enabled: Config value - False, "all", or list of providers
|
|
redis_cache: Redis cache instance (required for polling)
|
|
model: Model name from the request (e.g., "gpt-5" or "openai/gpt-4o")
|
|
llm_router: LiteLLM router instance for looking up model deployments
|
|
native_background_mode: List of model names that should use native provider
|
|
background mode instead of polling via cache
|
|
|
|
Returns:
|
|
True if polling should be used, False otherwise
|
|
"""
|
|
# All conditions must be met
|
|
if not (background_mode and polling_via_cache_enabled and redis_cache):
|
|
return False
|
|
|
|
# Check if model is in native_background_mode list - these use native provider background mode
|
|
if native_background_mode and model in native_background_mode:
|
|
verbose_proxy_logger.debug("Model %s is in native_background_mode list, skipping polling via cache", model)
|
|
return False
|
|
|
|
# "all" enables polling for all providers
|
|
if polling_via_cache_enabled == "all":
|
|
return True
|
|
|
|
# Check if provider is in the enabled list
|
|
if isinstance(polling_via_cache_enabled, list):
|
|
# First, try to get provider from model string format "provider/model"
|
|
if "/" in model:
|
|
provider: Final = model.split("/")[0]
|
|
if provider in polling_via_cache_enabled:
|
|
return True
|
|
# Otherwise, check ALL deployments for this model_name in router
|
|
elif llm_router is not None:
|
|
try:
|
|
for deployment_dict in llm_router.deployments_named(model) or ():
|
|
litellm_params = deployment_dict.get("litellm_params", {})
|
|
|
|
# Check custom_llm_provider first
|
|
dep_provider = litellm_params.get("custom_llm_provider")
|
|
|
|
# Then try to extract from model (e.g., "openai/gpt-5")
|
|
if not dep_provider:
|
|
dep_model = litellm_params.get("model", "")
|
|
if "/" in dep_model:
|
|
dep_provider = dep_model.split("/")[0]
|
|
|
|
# If ANY deployment's provider matches, enable polling
|
|
if dep_provider and dep_provider in polling_via_cache_enabled:
|
|
verbose_proxy_logger.debug("Polling enabled for model=%s, provider=%s", model, dep_provider)
|
|
return True
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug("Could not resolve provider for model %s: %s", model, e)
|
|
|
|
return False
|