mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 22:22:23 +00:00
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
This commit is contained in:
Xianzong Xie
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user