""" Background Streaming Task for Polling Via Cache Feature Handles streaming responses from LLM providers and updates Redis cache with partial results for polling. Follows OpenAI Response Streaming format: https://platform.openai.com/docs/api-reference/responses-streaming """ import asyncio import json from collections.abc import Sequence from typing import TYPE_CHECKING, Final, TypedDict, cast from fastapi import Request, Response from fastapi.responses import StreamingResponse from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import ProxyLogging from litellm.router import Router class _StreamContentPart(TypedDict, total=False): text: ReadOnly[str] class _StreamOutputItem(TypedDict, total=False): id: ReadOnly[str] content: ReadOnly[Sequence[_StreamContentPart | None]] async def background_streaming_task( polling_id: str, data, polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, general_settings, llm_router: "Router | None", proxy_config: "ProxyConfig", proxy_logging_obj: "ProxyLogging", select_data_generator, user_model, user_temperature: float | None, user_request_timeout: float | None, user_max_tokens: int | None, user_api_base: str | None, version: str | None, ): """ Background task to stream response and update cache Follows OpenAI Response Streaming format: https://platform.openai.com/docs/api-reference/responses-streaming Processes streaming events and builds Response object: https://platform.openai.com/docs/api-reference/responses/object """ try: verbose_proxy_logger.info("Starting background streaming for %s", polling_id) # Update status to in_progress (OpenAI format) await polling_handler.update_state( polling_id=polling_id, status="in_progress", ) # Force streaming mode and remove background flag data["stream"] = True data.pop("background", None) # Create processor processor: Final = ProxyBaseLLMRequestProcessing(data=data) # Make streaming request. # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. response: Final[StreamingResponse] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, route_type="aresponses", proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, select_data_generator=select_data_generator, model=None, user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, user_max_tokens=user_max_tokens, user_api_base=user_api_base, version=version, skip_pre_call_logic=True, ) # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming output_items: Final[dict[str, _StreamOutputItem]] = {} # Track output items by ID # Track accumulated text deltas by (item_id, content_index) accumulated_text: Final[dict[tuple[str, int], str]] = {} # ResponsesAPIResponse fields to extract from response.completed usage_data = None reasoning_data = None tool_choice_data = None tools_data = None model_data = None instructions_data = None temperature_data = None top_p_data = None max_output_tokens_data = None previous_response_id_data = None text_data = None truncation_data = None parallel_tool_calls_data = None user_data = None store_data = None incomplete_details_data = None state_dirty = False # Track if state needs to be synced last_update_time = asyncio.get_event_loop().time() UPDATE_INTERVAL: Final = 0.150 # 150ms batching interval # Track the terminal event from the stream (may not be "completed") terminal_status: ResponsesAPIStatus | None = ( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None _event_to_status: Final = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", "response.cancelled": "cancelled", } async def flush_state_if_needed(force: bool = False) -> None: """Flush accumulated state to Redis if interval elapsed or forced""" nonlocal state_dirty, last_update_time current_time: Final = asyncio.get_event_loop().time() if state_dirty and (force or (current_time - last_update_time) >= UPDATE_INTERVAL): # Convert output_items dict to list for update output_list: Final = list(output_items.values()) await polling_handler.update_state( polling_id=polling_id, output=output_list, ) state_dirty = False last_update_time = current_time # Handle StreamingResponse if not hasattr(response, "body_iterator"): verbose_proxy_logger.warning( "background_streaming_task: response for %s has no body_iterator; this may indicate a misconfiguration or provider error", polling_id, ) if hasattr(response, "body_iterator"): async for chunk in response.body_iterator: # Parse chunk if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") if isinstance(chunk, str) and chunk.startswith("data: "): chunk_data = chunk[6:].strip() if chunk_data == "[DONE]": break try: event = json.loads(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec if event_type == "response.output_item.added": # New output item added item = event.get("item", {}) item_id = item.get("id") if item_id: output_items[item_id] = item state_dirty = True elif event_type == "response.content_part.added": # Content part added to an output item item_id = event.get("item_id") content_part = event.get("part", {}) if item_id and item_id in output_items: # Update the output item with new content current_item = output_items[item_id] appended_item: _StreamOutputItem = { **current_item, "content": (*current_item.get("content", ()), content_part), } output_items[item_id] = appended_item state_dirty = True elif event_type == "response.output_text.delta": # Text delta - accumulate text content # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta item_id = event.get("item_id") content_index: int = event.get("content_index", 0) delta = event.get("delta", "") if item_id and item_id in output_items: # Accumulate text delta key = (item_id, content_index) if key not in accumulated_text: accumulated_text[key] = "" accumulated_text[key] += delta # Update the content in output_items current_item = output_items[item_id] content_list: Sequence[_StreamContentPart | None] = current_item.get("content", ()) if content_index < len(content_list): # Update existing content part with accumulated text content_entry = content_list[content_index] if isinstance(content_entry, dict): delta_part: _StreamContentPart = { **content_entry, "text": accumulated_text[key], } delta_item: _StreamOutputItem = { **current_item, "content": tuple( delta_part if index == content_index else entry for index, entry in enumerate(content_list) ), } output_items[item_id] = delta_item state_dirty = True elif event_type == "response.content_part.done": # Content part completed item_id = event.get("item_id") content_part = event.get("part", {}) content_index = event.get("content_index", 0) if item_id and item_id in output_items: # Update with final content from event current_item = output_items[item_id] content_list = current_item.get("content", ()) if content_index < len(content_list): finalized_item: _StreamOutputItem = { **current_item, "content": tuple( content_part if index == content_index else entry for index, entry in enumerate(content_list) ), } output_items[item_id] = finalized_item state_dirty = True elif event_type == "response.output_item.done": # Output item completed - use final item data item = event.get("item", {}) item_id = item.get("id") if item_id: output_items[item_id] = item state_dirty = True elif event_type == "response.in_progress": # Response is now in progress # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress await polling_handler.update_state( polling_id=polling_id, status="in_progress", ) elif event_type in ( "response.completed", "response.failed", "response.incomplete", "response.cancelled", ): # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) terminal_status = cast( ResponsesAPIStatus, response_data.get( "status", _event_to_status.get(event_type, "completed"), ), ) # Extract error for failed and incomplete responses if event_type == "response.failed" or event_type == "response.incomplete": terminal_error = response_data.get("error") # Core response fields usage_data = response_data.get("usage") reasoning_data = response_data.get("reasoning") tool_choice_data = response_data.get("tool_choice") tools_data = response_data.get("tools") # Additional ResponsesAPIResponse fields model_data = response_data.get("model") instructions_data = response_data.get("instructions") temperature_data = response_data.get("temperature") top_p_data = response_data.get("top_p") max_output_tokens_data = response_data.get("max_output_tokens") previous_response_id_data = response_data.get("previous_response_id") text_data = response_data.get("text") truncation_data = response_data.get("truncation") parallel_tool_calls_data = response_data.get("parallel_tool_calls") user_data = response_data.get("user") store_data = response_data.get("store") incomplete_details_data = response_data.get("incomplete_details") # Also update output from final response if available if "output" in response_data: final_output = response_data.get("output", []) for item in final_output: item_id = item.get("id") if item_id: output_items[item_id] = item state_dirty = True # Flush state to Redis if interval elapsed await flush_state_if_needed() except json.JSONDecodeError as e: verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) # Final flush to ensure all accumulated state is saved await flush_state_if_needed(force=True) # Use the terminal status from the stream, default to "completed" final_status: Final = terminal_status or "completed" await polling_handler.update_state( polling_id=polling_id, status=final_status, usage=usage_data, error=terminal_error, reasoning=reasoning_data, tool_choice=tool_choice_data, tools=tools_data, model=model_data, instructions=instructions_data, temperature=temperature_data, top_p=top_p_data, max_output_tokens=max_output_tokens_data, previous_response_id=previous_response_id_data, text=text_data, truncation=truncation_data, parallel_tool_calls=parallel_tool_calls_data, user=user_data, store=store_data, incomplete_details=incomplete_details_data, ) verbose_proxy_logger.info( "Finished background streaming for %s, status=%s, error=%s, incomplete_details=%s, output_items=%s", polling_id, final_status, terminal_error, incomplete_details_data, len(output_items), ) except Exception as e: verbose_proxy_logger.error("Error in background streaming task for %s: %s", polling_id, e) import traceback verbose_proxy_logger.error(traceback.format_exc()) await polling_handler.update_state( polling_id=polling_id, status="failed", error={ "type": "internal_error", "message": str(e), "code": "background_streaming_error", }, )