""" 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 typing import Any, Optional, cast from fastapi import Request, Response 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 async def background_streaming_task( # noqa: PLR0915 polling_id: str, data: dict, polling_handler: ResponsePollingHandler, request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, general_settings: dict, llm_router, proxy_config, proxy_logging_obj, select_data_generator, user_model, user_temperature, user_request_timeout, user_max_tokens, user_api_base, version, ): """ 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(f"Starting background streaming for {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 = 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 = 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: dict[str, dict[str, Any]] = {} # Track output items by ID accumulated_text = ( {} ) # Track accumulated text deltas by (item_id, content_index) # 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 = 0.150 # 150ms batching interval # Track the terminal event from the stream (may not be "completed") terminal_status: Optional[ ResponsesAPIStatus ] = None # Will be set by response.completed/failed/incomplete/cancelled terminal_error = None _event_to_status = { "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 = 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 = 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( f"background_streaming_task: response for {polling_id} has no " "body_iterator; this may indicate a misconfiguration or provider error" ) 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 if "content" not in output_items[item_id]: output_items[item_id]["content"] = [] output_items[item_id]["content"].append(content_part) 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 = 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 if "content" in output_items[item_id]: content_list = output_items[item_id]["content"] if content_index < len(content_list): # Update existing content part with accumulated text if isinstance( content_list[content_index], dict ): content_list[content_index][ "text" ] = accumulated_text[key] 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 if "content" in output_items[item_id]: content_list = output_items[item_id]["content"] if content_index < len(content_list): content_list[content_index] = content_part 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( f"Failed to parse streaming chunk: {e}" ) pass # 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 = 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( f"Finished background streaming for {polling_id}, status={final_status}, error={terminal_error}, incomplete_details={incomplete_details_data}, output_items={len(output_items)}" ) except Exception as e: verbose_proxy_logger.error( f"Error in background streaming task for {polling_id}: {str(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", }, )