From 1c3c12bb1be52f2333ed00e7ea8a328076dad7f6 Mon Sep 17 00:00:00 2001 From: Xianzong Xie Date: Wed, 3 Dec 2025 22:50:26 -0800 Subject: [PATCH] refactor: move background_streaming_task to separate module - Create new background_streaming.py in response_polling/ - Update endpoints.py to import from new location - Update __init__.py to export background_streaming_task - Add tests for module imports and structure Committed-By-Agent: cursor --- .../proxy/response_api_endpoints/endpoints.py | 251 +---------------- litellm/proxy/response_polling/__init__.py | 9 +- .../response_polling/background_streaming.py | 263 ++++++++++++++++++ .../test_response_polling_handler.py | 32 +++ 4 files changed, 307 insertions(+), 248 deletions(-) create mode 100644 litellm/proxy/response_polling/background_streaming.py diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c19c6555d2..d435f0a34c 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,5 +1,4 @@ import asyncio -import json from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -11,250 +10,6 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin router = APIRouter() -async def _background_streaming_task( # noqa: PLR0915 - polling_id: str, - data: dict, - polling_handler, - 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 - 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, - ) - - # Process streaming response following OpenAI events format - # https://platform.openai.com/docs/api-reference/responses-streaming - output_items = {} # Track output items by ID - accumulated_text = {} # Track accumulated text deltas by (item_id, content_index) - usage_data = None - reasoning_data = None - tool_choice_data = None - tools_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 - - 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 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 == "response.completed": - # Response completed - includes usage, reasoning, tools, tool_choice - # https://platform.openai.com/docs/api-reference/responses-streaming/response-completed - response_data = event.get("response", {}) - 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") - - # 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) - - # Mark as completed with all response data - await polling_handler.update_state( - polling_id=polling_id, - status="completed", - usage=usage_data, - reasoning=reasoning_data, - tool_choice=tool_choice_data, - tools=tools_data, - ) - - verbose_proxy_logger.info( - f"Completed background streaming for {polling_id}, 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" - }, - ) - - @router.post( "/v1/responses", dependencies=[Depends(user_api_key_auth)], @@ -346,6 +101,9 @@ async def responses_api( from litellm.proxy.response_polling.polling_handler import ( ResponsePollingHandler, ) + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) verbose_proxy_logger.info( f"Starting background response with polling for model={data.get('model')}" @@ -367,9 +125,8 @@ async def responses_api( ) # Start background task to stream and update cache - import asyncio asyncio.create_task( - _background_streaming_task( + background_streaming_task( polling_id=polling_id, data=data.copy(), polling_handler=polling_handler, diff --git a/litellm/proxy/response_polling/__init__.py b/litellm/proxy/response_polling/__init__.py index 5d8f053536..b014286b9e 100644 --- a/litellm/proxy/response_polling/__init__.py +++ b/litellm/proxy/response_polling/__init__.py @@ -1,5 +1,12 @@ """ Response Polling Module for Background Responses with Cache """ +from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, +) +from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler - +__all__ = [ + "ResponsePollingHandler", + "background_streaming_task", +] diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py new file mode 100644 index 0000000000..a0ce4d8221 --- /dev/null +++ b/litellm/proxy/response_polling/background_streaming.py @@ -0,0 +1,263 @@ +""" +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 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 + + +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 + 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, + ) + + # Process streaming response following OpenAI events format + # https://platform.openai.com/docs/api-reference/responses-streaming + output_items = {} # Track output items by ID + accumulated_text = {} # Track accumulated text deltas by (item_id, content_index) + usage_data = None + reasoning_data = None + tool_choice_data = None + tools_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 + + 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 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 == "response.completed": + # Response completed - includes usage, reasoning, tools, tool_choice + # https://platform.openai.com/docs/api-reference/responses-streaming/response-completed + response_data = event.get("response", {}) + 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") + + # 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) + + # Mark as completed with all response data + await polling_handler.update_state( + polling_id=polling_id, + status="completed", + usage=usage_data, + reasoning=reasoning_data, + tool_choice=tool_choice_data, + tools=tools_data, + ) + + verbose_proxy_logger.info( + f"Completed background streaming for {polling_id}, 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" + }, + ) + diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 352fe3e424..81231c61df 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -528,3 +528,35 @@ class TestStreamingEventProcessing: assert UPDATE_INTERVAL == 0.150 assert UPDATE_INTERVAL * 1000 == 150 # 150 milliseconds + +class TestBackgroundStreamingModule: + """Test cases for background_streaming module imports and structure""" + + def test_background_streaming_task_can_be_imported(self): + """Test that background_streaming_task can be imported from the module""" + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + assert background_streaming_task is not None + assert callable(background_streaming_task) + + def test_module_exports_from_init(self): + """Test that the module exports are available from __init__""" + from litellm.proxy.response_polling import ( + ResponsePollingHandler, + background_streaming_task, + ) + + assert ResponsePollingHandler is not None + assert background_streaming_task is not None + + def test_background_streaming_task_is_async(self): + """Test that background_streaming_task is an async function""" + import asyncio + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + assert asyncio.iscoroutinefunction(background_streaming_task) +