81474c17fe
* Handle response.failed, response.incomplete, and response.cancelled terminal events in background streaming
Previously the background streaming task only handled response.completed and
hardcoded the final status to "completed". This missed three other terminal
event types from the OpenAI streaming spec, causing failed/incomplete/cancelled
responses to be incorrectly marked as completed.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
* Remove unused terminal_response_data variable
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
* Address code review: derive fallback status from event type, rewrite tests as integration tests
1. Replace hardcoded "completed" fallback in response_data.get("status")
with _event_to_status lookup so that response.incomplete and
response.cancelled events get the correct fallback if the response
body ever omits the status field.
2. Replace duplicated-logic unit tests with integration tests that
exercise background_streaming_task directly using mocked streaming
responses and assert on the final update_state call arguments.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
* Remove dead mock_processor and unused mock_response parameter from test helper
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
* Remove FastAPI and UserAPIKeyAuth imports from test file
These types were only used as Mock(spec=...) arguments. Drop the spec
constraints and remove the top-level imports to avoid pulling FastAPI
into test files outside litellm/proxy/.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
* Log warning when streaming response has no body_iterator
If base_process_llm_request returns a non-streaming response (no
body_iterator), log a warning since this likely indicates a
misconfiguration or provider error rather than a successful completion.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Committed-By-Agent: claude
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
356 lines
15 KiB
Python
356 lines
15 KiB
Python
"""
|
|
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
|
|
|
|
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: 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 = 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 = response_data.get(
|
|
"status",
|
|
_event_to_status.get(event_type, "completed"),
|
|
)
|
|
|
|
# Extract error for failed responses
|
|
if event_type == "response.failed":
|
|
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}, 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",
|
|
},
|
|
)
|