scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90
litellm/litellm/responses/streaming_iterator.py
yucheng 8b86362703 fix(caching): replay cache hits for converted streams as streams 
A deployment hook (Headroom, code interpreter, web search) can downgrade
kwargs["stream"] to False while the caller still expects to iterate the
result. The cache handler keyed stream replay and callback deferral off
the raw flag, so a cache hit returned a plain object to a caller that
iterates, and the Responses iterator never persisted the converted
stream in the first place. Key both off the conversion marker as well

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-15 07:51:19 +00:00

2639 lines
111 KiB
Python
Raw Blame History

from __future__ import annotations
import asyncio
import json
import time
import traceback
import uuid
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable
import httpx
from openai._streaming import SSEDecoder
from pydantic import BaseModel, ValidationError
from typing_extensions import TypeIs
import litellm
from litellm.constants import (
EMPTY_MAPPING,
LITELLM_MAX_STREAMING_DURATION_SECONDS,
STREAM_SSE_DONE_STRING,
)
from litellm.exceptions import MidStreamFallbackError, RateLimitError
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
update_response_metadata,
)
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import (
PART_UNION_TYPES,
ResponseAPIUsage,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.responses.streaming_websocket import (
PresidioGuardrailCallback,
ResponsesBackendWebSocket,
ResponsesClientWebSocket,
)
from litellm.types.router import LiteLLM_Params
class ProjectQuotaCallback(Protocol):
async def enforce_project_io_token_quota_for_frame(
self,
user_api_key_dict: UserAPIKeyAuth | None,
requested_model: str | None,
estimated_input_tokens: int,
estimated_output_tokens: int,
) -> None: ...
@lru_cache(maxsize=1)
def _get_openai_response_types():
from litellm.types.llms import openai as openai_types
return openai_types
def _is_json_object(value: object) -> TypeIs[dict[str, object]]: # guard-ok: trivial isinstance; JSON keys are str
return isinstance(value, dict)
def _is_json_array(value: object) -> TypeIs[list[object]]: # guard-ok: trivial isinstance narrowing
return isinstance(value, list)
def _optional_str(value: object) -> str | None:
"""Keep a JSON payload entry only when it is a string, since the wire format is caller-controlled."""
return value if isinstance(value, str) else None
def _json_array_or_empty(value: object) -> Sequence[object]:
"""Narrow a JSON payload entry that the caller iterates, tolerating a missing or malformed value."""
return value if _is_json_array(value) else ()
def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verifies every value is str
return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
class _MutableJsonObject(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@overload
def get(self, key: str, default: object, /) -> object: ...
def __getitem__(self, key: str, /) -> object: ...
def __setitem__(self, key: str, value: object, /) -> None: ...
def __contains__(self, key: object, /) -> bool: ...
def items(self) -> Iterable[tuple[str, object]]: ...
class _GetsLitellmParams(Protocol):
def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ...
class _UnmasksPiiText(Protocol):
def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
class _ShouldStoreResultInCache(Protocol):
def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ...
class _PostStreamingDeploymentHook(Protocol):
def __call__(
self,
*,
request_data: Mapping[str, object],
response_chunk: ResponsesAPIStreamingResponse,
call_type: CallTypes | None,
) -> Awaitable[ResponsesAPIStreamingResponse | None]: ...
@runtime_checkable
class _HasPostStreamingDeploymentHook(Protocol):
async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook
def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams:
return fn
_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache"
_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text"
def _load_json_object(payload: str | bytes) -> dict[str, object]:
"""Parse a JSON payload that the caller consumes as an object."""
return json.loads(payload)
def _load_json_value(payload: str | bytes) -> object:
"""Parse a JSON payload whose top-level shape the caller narrows itself."""
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
return model_id if isinstance(model_id, str) else None
def _log_background_task_failure(task: asyncio.Task[object], *, task_name: str) -> None:
if task.cancelled():
return
exception: Final = task.exception()
if exception is not None:
verbose_logger.error("%s failed: %s", task_name, exception)
_ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(
{ # mutable-ok: immediately frozen by MappingProxyType
"server_error": 500,
"rate_limit_exceeded": 429,
"insufficient_quota": 429,
"vector_store_timeout": 504,
"invalid_prompt": 400,
"invalid_image": 400,
"invalid_image_format": 400,
"invalid_base64_image": 400,
"invalid_image_url": 400,
"image_too_large": 400,
"image_too_small": 400,
"image_parse_error": 400,
"image_content_policy_violation": 400,
"invalid_image_mode": 400,
"image_file_too_large": 400,
"unsupported_image_media_type": 400,
"empty_image_file": 400,
"failed_to_download_image": 400,
"image_file_not_found": 400,
"invalid_request_error": 400,
"context_length_exceeded": 400,
"content_policy_violation": 400,
"model_not_found": 400,
}
)
def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]:
if _is_json_object(error_obj):
raw_message = error_obj.get("message")
raw_type = error_obj.get("type")
raw_code = error_obj.get("code")
elif error_obj is not None:
raw_message = getattr(error_obj, "message", None)
raw_type = getattr(error_obj, "type", None)
raw_code = getattr(error_obj, "code", None)
else:
raw_message = None
raw_type = None
raw_code = None
message: Final = str(raw_message) if raw_message is not None else "Response API in-stream error"
error_type: Final = raw_type if isinstance(raw_type, str) else None
code: Final = raw_code if isinstance(raw_code, str) else None
return message, error_type, code
def _status_code_for_error_fields(error_type: str | None, error_code: str | None) -> int:
fields: Final = tuple(field for field in (error_code, error_type) if field is not None)
if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields):
return 429
return next(
(_ERROR_CODE_HTTP_STATUS[field] for field in fields if field in _ERROR_CODE_HTTP_STATUS),
500,
)
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
return True
status_code: Final = getattr(mapped_exception, "status_code", None)
return not isinstance(status_code, int) or status_code >= 500 or status_code == 429
class BaseResponsesAPIStreamingIterator:
"""
Base class for streaming iterators that process responses from the Responses API.
This class contains shared logic for both synchronous and asynchronous iterators.
"""
def __init__(
self,
response: httpx.Response,
model: str,
responses_api_provider_config: BaseResponsesAPIConfig | None,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
self.response = response
self.model = model
self.logging_obj = logging_obj
self.finished = False
self.responses_api_provider_config = responses_api_provider_config
self.completed_response: ResponsesAPIStreamingResponse | None = None
self.start_time = getattr(logging_obj, "start_time", datetime.now())
self._failure_handled = False # Track if failure handler has been called
self._yielded_first_chunk = False
self._generated_content = ""
self._completed_response_cached = False
self._completed_response_logged = False
self._completed_response_cache_hit: bool | None = None
self._persist_completed_response_before_logging = True
self._stream_created_time: float = time.time()
# track request context for hooks
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
self.request_data: dict[str, object] = request_data or {}
self.call_type: str | None = call_type
# set hidden params for response headers (e.g., x-litellm-model-id)
# This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py
_api_base: Final = get_api_base(
model=model or "",
optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}),
)
self._hidden_params: dict[str, object] = {
"model_id": _model_id_from_metadata(litellm_metadata),
"api_base": _api_base,
"custom_llm_provider": custom_llm_provider,
}
self._hidden_params["additional_headers"] = process_response_headers(
self.response.headers or {}
) # GUARANTEE OPENAI HEADERS IN RESPONSE
self._raw_response_headers: Mapping[str, str] = MappingProxyType(
dict(self.response.headers or {}) # mutable-ok: immediately frozen by MappingProxyType
)
def _check_max_streaming_duration(self) -> None:
"""Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS."""
if LITELLM_MAX_STREAMING_DURATION_SECONDS is None:
return
elapsed: Final = time.time() - self._stream_created_time
if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS:
raise litellm.Timeout(
message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)",
model=self.model or "",
llm_provider=self.custom_llm_provider or "",
)
def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None:
"""Process a single chunk of data from the stream"""
if not chunk:
return None
# NOTE: ``SSEDecoder`` already strips the SSE ``data:`` field prefix, so
# the value passed in here is the raw field content. Do not re-run
# ``_strip_sse_data_from_chunk`` on it — doing so would incorrectly mangle
# payloads whose actual JSON value happens to start with ``data:``.
# Handle "[DONE]" marker
if chunk == STREAM_SSE_DONE_STRING:
self.finished = True
return None
if self.logging_obj.completion_start_time is None:
self.logging_obj._update_completion_start_time(completion_start_time=datetime.now())
try:
# Parse the JSON chunk
parsed_chunk: Final = _load_json_value(chunk)
# Format as ResponsesAPIStreamingResponse
if _is_json_object(parsed_chunk):
if self.responses_api_provider_config is None:
raise ValueError("responses_api_provider_config is required to process live streaming chunks")
openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response(
model=self.model,
parsed_chunk=parsed_chunk,
logging_obj=self.logging_obj,
)
# Only when the SSE JSON carries a response body (delta events do not).
# Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a
# truthy child Mock for any attribute, which breaks tests and is wrong on stream.
if "response" in parsed_chunk:
response_object: Final[ResponsesAPIResponse | None] = getattr(
openai_responses_api_chunk, "response", None
)
if response_object is not None:
response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
responses_api_response=response_object,
litellm_metadata=self.litellm_metadata,
custom_llm_provider=self.custom_llm_provider,
)
setattr(openai_responses_api_chunk, "response", response)
# Encode container_id on streaming events so proxy/UI follow-ups route correctly
_event_type: Final = getattr(openai_responses_api_chunk, "type", None)
if _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA:
_delta: Final = getattr(openai_responses_api_chunk, "delta", None)
if isinstance(_delta, str):
self._generated_content += _delta
_stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata)
if _event_type in (
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
):
_item: Final[object] = getattr(openai_responses_api_chunk, "item", None)
if _item is not None:
ResponsesAPIRequestUtils._encode_container_id_on_output_item(
item=_item,
custom_llm_provider=self.custom_llm_provider,
model_id=_stream_model_id,
)
elif _event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED:
_annotation: Final[object] = getattr(openai_responses_api_chunk, "annotation", None)
if _annotation is not None:
ResponsesAPIRequestUtils._encode_container_id_on_output_item(
item=_annotation,
custom_llm_provider=self.custom_llm_provider,
model_id=_stream_model_id,
)
elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE:
_part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr(
openai_responses_api_chunk, "part", None
)
if _part is not None:
if isinstance(_part, dict):
ResponsesAPIRequestUtils._encode_container_ids_in_annotations(
_part.get("annotations"),
self.custom_llm_provider,
_stream_model_id,
)
else:
ResponsesAPIRequestUtils._encode_container_ids_in_annotations(
getattr(_part, "annotations", None),
self.custom_llm_provider,
_stream_model_id,
)
# Wrap encrypted_content in streaming events (output_item.added, output_item.done)
if self.litellm_metadata and self.litellm_metadata.get("encrypted_content_affinity_enabled"):
openai_types = _get_openai_response_types()
event_type: Final = getattr(openai_responses_api_chunk, "type", None)
if event_type in (
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
):
item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None)
if item:
encrypted_content: Final = getattr(item, "encrypted_content", None)
if encrypted_content and isinstance(encrypted_content, str):
model_id: Final = _model_id_from_metadata(self.litellm_metadata)
if model_id:
wrapped_content: Final = (
ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id(
encrypted_content, model_id
)
)
setattr(item, "encrypted_content", wrapped_content)
# Store the completed response (also for incomplete/failed so logging still fires)
_chunk_type: Final = getattr(openai_responses_api_chunk, "type", None)
openai_types = _get_openai_response_types()
if openai_responses_api_chunk and _chunk_type in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
_stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
else:
self._handle_logging_completed_response()
return openai_responses_api_chunk
return None
except json.JSONDecodeError:
# If we can't parse the chunk, continue
return None
except Exception as e:
# Trigger failure hooks before re-raising
# This ensures failures are logged even when _process_chunk is called directly
self._handle_failure(e)
raise
def _log_completed_response(self, *, is_async: bool) -> None:
if self._completed_response_logged:
return
self._completed_response_logged = True
if self._persist_completed_response_before_logging:
self._persist_completed_response_to_cache(is_async=is_async)
logging_response: Final[object] = _logging_copy(self.completed_response)
self._restore_provider_response_headers(logging_response)
end_time: Final = datetime.now()
if is_async:
logging_coroutine: Final = self.logging_obj.dispatch_success_handlers(
logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
prefer_async_handlers=True,
)
deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None
if deferred_dispatch_armed:
# End-of-stream guardrail scans write guardrail_information after
# the terminal event; dispatching now would snapshot metadata early.
self.logging_obj._deferred_stream_complete_args = (logging_coroutine,)
else:
asyncio.create_task(logging_coroutine)
else:
run_async_function(
async_function=self.logging_obj.async_success_handler,
result=logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
)
executor.submit(
self.logging_obj.success_handler,
result=logging_response,
cache_hit=self._completed_response_cache_hit,
start_time=self.start_time,
end_time=end_time,
)
self._run_post_success_hooks(end_time=end_time)
def _restore_provider_response_headers(self, logging_response: object) -> None:
"""Re-apply the provider's response headers to the copy handed to logging callbacks.
``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the
``_hidden_params`` the provider transform set on the nested response are lost. Returns early
when the event was not a pydantic model and logging got the original, so logging-only state
never lands on the object the caller is iterating.
"""
if logging_response is self.completed_response:
return
target: Final[object] = getattr(logging_response, "response", None)
if not isinstance(target, ResponsesAPIResponse):
return
existing: Final[Mapping[str, object]] = target._hidden_params
source_hidden: Final[object] = getattr(
getattr(self.completed_response, "response", None), "_hidden_params", None
)
source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING
processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers")
raw: Final[object] = source.get("headers") or self._raw_response_headers
headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING
raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING
# rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy
# splats into the client's HTTP headers, and copying non-header keys would carry response_cost
target._hidden_params = { # mutable-ok: the cost calculator writes optional_params into _hidden_params
"additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it
"headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it
**existing,
}
def _handle_logging_completed_response(self):
"""Base implementation - should be overridden by subclasses"""
def _handle_logging_failed_response(self):
"""
Handle logging for RESPONSE_FAILED events by routing to failure handlers.
Unlike _handle_logging_completed_response (which calls success handlers),
this constructs an exception from the response error and routes to
async_failure_handler / failure_handler so logging integrations correctly
record the call as failed.
"""
response_obj: Final[ResponsesAPIResponse | None] = (
getattr(self.completed_response, "response", None) if self.completed_response else None
)
error_info: Final = getattr(response_obj, "error", None) if response_obj else None
self._record_failed_response_usage(response_obj)
self._handle_failure(self._map_error_event_exception(error_info))
def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
if response_obj is None or self.logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
try:
self.logging_obj.model_call_details["combined_usage_object"] = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage_obj)
)
except (TypeError, ValueError) as usage_error:
verbose_logger.debug(
"could not record usage for failed responses stream: %s",
usage_error,
)
return
self.logging_obj.model_call_details["response_cost"] = (
self.logging_obj._response_cost_calculator(result=response_obj) or 0.0
)
def _map_error_event_exception(self, error_obj: object) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=self.model or "",
custom_llm_provider=self.custom_llm_provider or "",
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
def _maybe_raise_for_error_event(self, result: object) -> None:
chunk_type: Final = getattr(result, "type", None)
if chunk_type not in ("error", "response.failed"):
return
error_obj: Final[object] = (
getattr(getattr(result, "response", None), "error", None)
if chunk_type == "response.failed"
else getattr(result, "error", None)
)
mapped_exception: Final = self._map_error_event_exception(error_obj)
if not _mid_stream_fallback_eligible(mapped_exception):
raise mapped_exception
raise MidStreamFallbackError(
message=str(mapped_exception),
model=self.model or "",
llm_provider=self.custom_llm_provider or "",
original_exception=mapped_exception,
generated_content=self._generated_content,
is_pre_first_chunk=not self._yielded_first_chunk,
)
def _get_completed_response_object(self) -> ResponsesAPIResponse | None:
openai_types: Final = _get_openai_response_types()
completed_response: Final = self.completed_response
if isinstance(completed_response, openai_types.ResponsesAPIResponse):
return completed_response
response_obj: Final = getattr(completed_response, "response", None)
if isinstance(response_obj, openai_types.ResponsesAPIResponse):
return response_obj
return None
def _persist_completed_response_to_cache(self, *, is_async: bool) -> None:
if self._completed_response_cached:
return
completed_response: Final = self.completed_response
openai_types: Final = _get_openai_response_types()
if getattr(completed_response, "type", None) != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
return
response_obj: Final = self._get_completed_response_object()
if response_obj is None:
return
caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
request_kwargs = getattr(caching_handler, "request_kwargs", None)
if not _is_json_object(request_kwargs):
return
if request_kwargs.get("stream") is not True and not converted_stream_requested(request_kwargs):
return
request_kwargs = request_kwargs.copy()
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
request_cache_key: Final = request_kwargs.pop("cache_key", None)
if preset_cache_key is None:
preset_cache_key = request_cache_key
if request_kwargs.get("metadata") is None:
request_kwargs.pop("metadata", None)
request_kwargs.pop("custom_llm_provider", None)
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr(
caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR
)
if not should_store_result_in_cache(
original_function=getattr(caching_handler, "original_function", None),
kwargs=request_kwargs,
):
return
cache: Final = litellm.cache
if cache is None:
return
cached_response: Final = response_obj.model_dump_json()
if is_async:
from litellm.caching.caching_handler import create_cache_write_task
cache_write_task: Final = create_cache_write_task(
lambda: cache.async_add_cache(
cached_response,
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
**request_kwargs,
)
)
cache_write_task.add_done_callback(
lambda task: _log_background_task_failure(
task,
task_name="Responses stream cache write",
)
)
else:
cache.add_cache(
cached_response,
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
**request_kwargs,
)
self._completed_response_cached = True
async def _call_post_streaming_deployment_hook(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Allow callbacks to modify streaming chunks before returning (parity with chat).
"""
try:
# Align with chat pipeline: use logging_obj model_call_details + call_type
typed_call_type: CallTypes | None = None
if self.call_type is not None:
try:
typed_call_type = CallTypes(self.call_type)
except ValueError:
typed_call_type = None
if typed_call_type is None:
try:
typed_call_type = CallTypes(getattr(self.logging_obj, "call_type", None))
except Exception:
typed_call_type = None
request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {})
callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or []
hooks_ran = False
for callback in callbacks:
if isinstance(callback, _HasPostStreamingDeploymentHook):
hooks_ran = True
post_streaming_hook: _PostStreamingDeploymentHook = (
callback.async_post_call_streaming_deployment_hook
)
result = await post_streaming_hook(
request_data=request_data,
response_chunk=chunk,
call_type=typed_call_type,
)
if result is not None:
chunk = result
if hooks_ran:
setattr(chunk, "_post_streaming_hooks_ran", True)
return chunk
except Exception:
return chunk
async def call_post_streaming_hooks_for_testing(
self, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Helper to invoke streaming deployment hooks explicitly (used in tests).
"""
return await self._call_post_streaming_deployment_hook(chunk)
def _run_post_success_hooks(self, end_time: datetime):
"""
Run post-call deployment hooks and update metadata similar to chat pipeline.
"""
if self.completed_response is None:
return
request_payload: Final[dict[str, object]] = {}
if isinstance(self.request_data, dict):
request_payload.update(self.request_data)
try:
if hasattr(self.logging_obj, "model_call_details"):
request_payload.update(self.logging_obj.model_call_details)
except Exception:
pass
if "litellm_params" not in request_payload:
try:
request_payload["litellm_params"] = getattr(self.logging_obj, "model_call_details", {}).get(
"litellm_params", {}
)
except Exception:
request_payload["litellm_params"] = {}
try:
update_response_metadata(
result=self.completed_response,
logging_obj=self.logging_obj,
model=self.model,
kwargs=request_payload,
start_time=self.start_time,
end_time=end_time,
# the provider call was timed to first byte, so the whole stream minus it is not overhead
include_overhead=False,
)
except Exception:
# Non-blocking
pass
try:
typed_call_type: CallTypes | None = None
if self.call_type is not None:
try:
typed_call_type = CallTypes(self.call_type)
except ValueError:
typed_call_type = None
except Exception:
typed_call_type = None
if typed_call_type is None:
try:
typed_call_type = CallTypes.responses
except Exception:
typed_call_type = None
try:
# Call synchronously; async hook will be executed via asyncio.run in a new loop
run_async_function(
async_function=async_post_call_success_deployment_hook,
request_data=request_payload,
response=self.completed_response,
call_type=typed_call_type,
)
except Exception:
pass
def _handle_failure(self, exception: Exception):
"""
Trigger failure handlers before bubbling the exception.
Only calls handlers once even if called multiple times.
"""
# Prevent double-calling failure handlers
if self._failure_handled:
return
self._failure_handled = True
traceback_exception: Final = traceback.format_exc()
try:
run_async_function(
async_function=self.logging_obj.async_failure_handler,
exception=exception,
traceback_exception=traceback_exception,
start_time=self.start_time,
end_time=datetime.now(),
)
except Exception:
pass
try:
executor.submit(
self.logging_obj.failure_handler,
exception,
traceback_exception,
self.start_time,
datetime.now(),
)
except Exception:
pass
async def call_post_streaming_hooks_for_testing(
iterator: object, chunk: ResponsesAPIStreamingResponse
) -> ResponsesAPIStreamingResponse:
"""
Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped.
"""
hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = (
getattr(iterator, "_call_post_streaming_deployment_hook", None)
)
if hook_fn is None:
return chunk
return await hook_fn(chunk)
class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
"""
Async iterator for processing streaming responses from the Responses API.
"""
def __init__(
self,
response: httpx.Response,
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
super().__init__(
response,
model,
responses_api_provider_config,
logging_obj,
litellm_metadata,
custom_llm_provider,
request_data,
call_type,
)
self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes())
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
try:
self._check_max_streaming_duration()
while True:
# Get the next chunk from the stream
try:
sse = await self.stream_iterator.__anext__()
except StopAsyncIteration:
self.finished = True
raise StopAsyncIteration
self._check_max_streaming_duration()
result = self._process_chunk(sse.data)
if self.finished:
raise StopAsyncIteration
elif result is not None:
self._maybe_raise_for_error_event(result)
# Await hook directly instead of run_async_function
# (which spawns a thread + event loop per call)
result = await self._call_post_streaming_deployment_hook(
chunk=result,
)
self._yielded_first_chunk = True
return result
# If result is None, continue the loop to get the next chunk
except StopAsyncIteration:
# Normal end of stream - don't log as failure
raise
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
self.finished = True
if self.completed_response is None:
self._handle_failure(e)
raise
raise StopAsyncIteration from e
except httpx.HTTPError as e:
# Handle HTTP errors
self.finished = True
self._handle_failure(e)
raise e
except Exception as e:
self.finished = True
self._handle_failure(e)
raise e
def _handle_logging_completed_response(self):
"""Handle logging for completed responses in async context"""
self._log_completed_response(is_async=True)
class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
"""
Synchronous iterator for processing streaming responses from the Responses API.
"""
def __init__(
self,
response: httpx.Response,
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
super().__init__(
response,
model,
responses_api_provider_config,
logging_obj,
litellm_metadata,
custom_llm_provider,
request_data,
call_type,
)
self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes())
def __iter__(self):
return self
def __next__(self):
try:
self._check_max_streaming_duration()
while True:
# Get the next chunk from the stream
try:
sse = next(self.stream_iterator)
except StopIteration:
self.finished = True
raise StopIteration
self._check_max_streaming_duration()
result = self._process_chunk(sse.data)
if self.finished:
raise StopIteration
elif result is not None:
self._maybe_raise_for_error_event(result)
# Sync path: use run_async_function for the hook
result = run_async_function(
async_function=self._call_post_streaming_deployment_hook,
chunk=result,
)
self._yielded_first_chunk = True
return result
# If result is None, continue the loop to get the next chunk
except StopIteration:
# Normal end of stream - don't log as failure
raise
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
self.finished = True
if self.completed_response is None:
self._handle_failure(e)
raise
raise StopIteration from e
except httpx.HTTPError as e:
# Handle HTTP errors
self.finished = True
self._handle_failure(e)
raise e
except Exception as e:
self.finished = True
self._handle_failure(e)
raise e
def _handle_logging_completed_response(self):
"""Handle logging for completed responses in sync context"""
self._log_completed_response(is_async=False)
class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
"""
Mock iterator—fake a stream by slicing the full response text into
5 char deltas, then emit a completed event.
Models like o1-pro don't support streaming, so we fake it.
"""
CHUNK_SIZE = 5
def __init__(
self,
response: httpx.Response,
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
logging_obj: LiteLLMLoggingObj,
litellm_metadata: dict[str, object] | None = None,
custom_llm_provider: str | None = None,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
transformed: Final = responses_api_provider_config.transform_response_api_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)
super().__init__(
response=httpx.Response(200),
model=model,
responses_api_provider_config=None,
logging_obj=logging_obj,
litellm_metadata=litellm_metadata,
custom_llm_provider=custom_llm_provider,
request_data=request_data,
call_type=call_type,
)
self._set_events_from_response(transformed=transformed, logging_obj=logging_obj)
def _set_events_from_response(
self,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events: Sequence[ResponsesAPIStreamingResponse] = build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=self.CHUNK_SIZE,
)
self._idx = 0
self.completed_response = self._events[-1]
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt: Final = self._events[self._idx]
self._idx += 1
openai_types: Final = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
def __iter__(self):
return self
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt: Final = self._events[self._idx]
self._idx += 1
openai_types: Final = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt
class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __init__(
self,
response: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict[str, object] | None = None,
call_type: str | None = None,
):
BaseResponsesAPIStreamingIterator.__init__(
self,
response=httpx.Response(200),
model=getattr(response, "model", ""),
responses_api_provider_config=None,
logging_obj=logging_obj,
litellm_metadata=None,
custom_llm_provider="cached_response",
request_data=request_data,
call_type=call_type,
)
self._completed_response_cache_hit = True
self._persist_completed_response_before_logging = False
self._events: list[ResponsesAPIStreamingResponse] = []
self._idx = 0
self._set_events_from_response(transformed=response, logging_obj=logging_obj)
def _set_events_from_response(
self,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE,
)
self._idx = 0
self.completed_response = self._events[-1]
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt: Final = self._events[self._idx]
self._idx += 1
openai_types: Final = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
def __iter__(self):
return self
def __next__(self) -> ResponsesAPIStreamingResponse:
if self._idx >= len(self._events):
raise StopIteration
evt: Final = self._events[self._idx]
self._idx += 1
openai_types: Final = _get_openai_response_types()
if getattr(evt, "type", None) == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt
@runtime_checkable
class _HasModelDump(Protocol):
def model_dump(self, *, exclude_none: bool = ...) -> dict[str, object]: ...
@runtime_checkable
class _HasModelDumpJson(Protocol):
def model_dump_json(self, *, exclude_none: bool = ...) -> str: ...
def _dump_response_object(obj: object) -> Mapping[str, object]:
if isinstance(obj, _HasModelDump):
return obj.model_dump()
if _is_json_object(obj):
return obj
return {}
def _build_response_status_event(
event_type: Literal[
"response.created",
"response.in_progress",
],
transformed: ResponsesAPIResponse,
) -> ResponsesAPIStreamingResponse:
openai_types: Final = _get_openai_response_types()
in_progress_response: Final = transformed.model_copy(
deep=True,
update={"status": "in_progress", "output": []},
)
if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
return openai_types.ResponseCreatedEvent(type=event_type, response=in_progress_response)
return openai_types.ResponseInProgressEvent(type=event_type, response=in_progress_response)
def _build_content_part_done_event(
*,
item_id: str,
output_index: int,
content_index: int,
part_payload: Mapping[str, object],
) -> ResponsesAPIStreamingResponse | None:
openai_types: Final = _get_openai_response_types()
part_type: Final = part_payload.get("type")
part: PART_UNION_TYPES
if part_type == "output_text":
raw_annotations: Final[object] = part_payload.get("annotations", []) or []
part = openai_types.ContentPartDonePartOutputText.model_validate(
{
"type": "output_text",
"text": str(part_payload.get("text") or ""),
"annotations": raw_annotations,
"logprobs": part_payload.get("logprobs"),
}
)
elif part_type == "refusal":
part = openai_types.ContentPartDonePartRefusal(
type="refusal",
refusal=str(part_payload.get("refusal") or ""),
)
elif part_type == "reasoning_text":
part = openai_types.ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=str(part_payload.get("reasoning") or ""),
)
else:
return None
return openai_types.ContentPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part=part,
)
def _add_text_like_part_events(
*,
events: list[ResponsesAPIStreamingResponse],
item_id: str,
output_index: int,
content_index: int,
part_payload: Mapping[str, object],
chunk_size: int,
) -> None:
openai_types: Final = _get_openai_response_types()
part_type: Final = part_payload.get("type")
if part_type == "output_text":
text: Final = str(part_payload.get("text") or "")
for i in range(0, len(text), chunk_size):
events.append(
openai_types.OutputTextDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=output_index,
content_index=content_index,
delta=text[i : i + chunk_size],
)
)
raw_annotation_items: Final = part_payload.get("annotations")
annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else []
for annotation_index, annotation in enumerate(annotation_items):
events.append(
openai_types.OutputTextAnnotationAddedEvent.model_validate(
{
"type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
"item_id": item_id,
"output_index": output_index,
"content_index": content_index,
"annotation_index": annotation_index,
"annotation": annotation,
}
)
)
events.append(
openai_types.OutputTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
text=text,
)
)
elif part_type == "refusal":
refusal: Final = str(part_payload.get("refusal") or "")
for i in range(0, len(refusal), chunk_size):
events.append(
openai_types.RefusalDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA,
item_id=item_id,
output_index=output_index,
content_index=content_index,
delta=refusal[i : i + chunk_size],
)
)
events.append(
openai_types.RefusalDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
refusal=refusal,
)
)
def _logging_copy(event: object) -> object:
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow
copies of the event and its nested response still keep the caller's ``usage`` attribute separate."""
if not isinstance(event, BaseModel):
return event
try:
return type(event).model_validate(event.model_dump())
except Exception:
return _detached_shallow_copy(event)
def _detached_shallow_copy(event: BaseModel) -> BaseModel:
nested: Final[object] = getattr(event, "response", None)
if isinstance(nested, BaseModel):
return event.model_copy(update={"response": nested.model_copy()})
return event.model_copy()
def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
if isinstance(usage, ResponseAPIUsage):
return usage
if not isinstance(usage, dict):
return None
try:
return ResponseAPIUsage.model_validate(usage)
except ValidationError:
return None
def _stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:
if response_obj is None or logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives
if isinstance(getattr(usage_obj, "cost", None), (int, float)):
return
try:
cost: Final[float | None] = logging_obj._response_cost_calculator(result=response_obj)
except Exception:
return
if isinstance(cost, (int, float)) and cost > 0:
setattr(usage_obj, "cost", cost)
def build_synthetic_response_events(
*,
transformed: ResponsesAPIResponse,
logging_obj: LiteLLMLoggingObj | None,
chunk_size: int,
) -> list[ResponsesAPIStreamingResponse]:
openai_types: Final = _get_openai_response_types()
_stamp_responses_usage_cost(transformed, logging_obj)
events: Final[list[ResponsesAPIStreamingResponse]] = [
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed),
]
sequence_number = 0
output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or []
for output_index, output_item in enumerate(output_items):
output_item_payload = _dump_response_object(output_item)
item_id = str(output_item_payload.get("id") or transformed.id)
item_type = output_item_payload.get("type")
events.append(
openai_types.OutputItemAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload),
)
)
if item_type == "message":
content_parts: Sequence[object] = _json_array_or_empty(output_item_payload.get("content"))
for content_index, part in enumerate(content_parts):
part_payload = _dump_response_object(part)
events.append(
openai_types.ContentPartAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part=openai_types.BaseLiteLLMOpenAIResponseObject(**part_payload),
)
)
_add_text_like_part_events(
events=events,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part_payload=part_payload,
chunk_size=chunk_size,
)
done_event = _build_content_part_done_event(
item_id=item_id,
output_index=output_index,
content_index=content_index,
part_payload=part_payload,
)
if done_event is not None:
events.append(done_event)
elif item_type == "function_call":
arguments = str(output_item_payload.get("arguments") or "")
for i in range(0, len(arguments), chunk_size):
events.append(
openai_types.FunctionCallArgumentsDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=item_id,
output_index=output_index,
delta=arguments[i : i + chunk_size],
)
)
events.append(
openai_types.FunctionCallArgumentsDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
item_id=item_id,
output_index=output_index,
arguments=arguments,
)
)
elif item_type == "reasoning":
summaries: Sequence[object] = _json_array_or_empty(output_item_payload.get("summary"))
for summary_index, summary in enumerate(summaries):
summary_payload = _dump_response_object(summary)
summary_text = str(summary_payload.get("text") or "")
for i in range(0, len(summary_text), chunk_size):
events.append(
openai_types.ReasoningSummaryTextDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
item_id=item_id,
output_index=output_index,
summary_index=summary_index,
delta=summary_text[i : i + chunk_size],
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
summary_index=summary_index,
text=summary_text,
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
summary_index=summary_index,
part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload),
)
)
sequence_number += 1
events.append(
openai_types.OutputItemDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=sequence_number,
item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload),
)
)
events.append(
openai_types.ResponseCompletedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=transformed,
)
)
return events
# ---------------------------------------------------------------------------
# WebSocket mode streaming (bidirectional forwarding)
# ---------------------------------------------------------------------------
from litellm._logging import verbose_logger
# Conservative per-frame output-token floor used when a response.create
# frame omits max_output_tokens, so a project OTPM quota can't be bypassed
# by simply never declaring an output cap.
_FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR: Final = 1024
# Rough chars-per-token ratio for estimating a frame's input tokens without
# resolving a real per-model tokenizer, matching the conservative estimate
# the proxy's own rate limiter uses for the same purpose.
_FRAME_CHARS_PER_TOKEN_ESTIMATE: Final = 4
def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple[int, int | None]:
"""Extract a rough input-token count and any explicit max_output_tokens
from a ``response.create`` frame, handling both wire shapes:
flat: {"type": "response.create", "input": ..., "max_output_tokens": ...}
nested: {"type": "response.create", "response": {"input": ..., "max_output_tokens": ...}}
"""
nested: Final = msg_obj.get("response")
params: Final[Mapping[str, object]] = (
nested
if _is_json_object(nested) and nested
else MappingProxyType( # mutable-ok: immediately frozen filtered frame
{k: v for k, v in msg_obj.items() if k != "type"}
)
)
text_parts: Final[list[str]] = [] # mutable-ok: local accumulator built in one pass, not shared
pending: Final[list[object]] = [ # mutable-ok: explicit worklist avoids recursion
params.get("input"),
params.get("instructions"),
]
while pending:
value = pending.pop()
if isinstance(value, str):
text_parts.append(value)
elif _is_json_array(value):
for item in value:
if isinstance(item, str):
text_parts.append(item)
elif _is_json_object(item):
pending.append(item.get("content"))
pending.append(item.get("text"))
total_chars: Final = sum(len(part) for part in text_parts)
estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0
max_output_tokens: Final = params.get("max_output_tokens")
return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None
async def _enforce_frame_project_quota(
quota_callbacks: Sequence[ProjectQuotaCallback],
user_api_key_dict: UserAPIKeyAuth | None,
model: str | None,
raw_message: str,
) -> None:
"""Charge one response.create frame's estimated tokens against every
registered project ITPM/OTPM quota callback, in isolation from PII
masking / logging so a malformed frame still reaches those callbacks."""
if not quota_callbacks:
return
try:
msg_obj: Final = _load_json_value(raw_message)
except (json.JSONDecodeError, TypeError):
return
if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create":
return
estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj)
estimated_output_tokens: Final = (
explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR
)
for callback in quota_callbacks:
await callback.enforce_project_io_token_quota_for_frame(
user_api_key_dict=user_api_key_dict,
requested_model=model,
estimated_input_tokens=estimated_input_tokens,
estimated_output_tokens=estimated_output_tokens,
)
RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
"response.created",
"response.completed",
"response.failed",
"response.incomplete",
"error",
]
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
class ResponsesWebSocketStreaming:
"""
Manages bidirectional WebSocket forwarding for the Responses API
WebSocket mode (wss://.../v1/responses).
Unlike the Realtime API, the Responses API WebSocket mode:
- Uses response.create as the client-to-server event
- Streams back the same events as the HTTP streaming Responses API
- Supports previous_response_id for incremental continuation
- Supports generate: false for warmup
- One response at a time per connection (sequential, no multiplexing)
"""
def __init__(
self,
websocket: ResponsesClientWebSocket,
backend_ws: ResponsesBackendWebSocket,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
first_message: str | None = None,
guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: str | None = authorized_model
def _should_store_event(self, event_obj: _MutableJsonObject) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: str | bytes | dict[str, object]) -> None:
if isinstance(event, bytes):
event = event.decode("utf-8")
if isinstance(event, str):
try:
event_obj = _load_json_object(event)
except (json.JSONDecodeError, TypeError):
return
else:
event_obj = event
if self._should_store_event(event_obj):
self.messages.append(event_obj)
def _collect_input_from_client_event(self, message: object) -> None:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = _load_json_object(message)
elif _is_json_object(message):
msg_obj = message
else:
return
if msg_obj.get("type") != "response.create":
return
input_items: Final = msg_obj.get("input", [])
if isinstance(input_items, str):
self.input_messages.append({"role": "user", "content": input_items})
return
if _is_json_array(input_items):
for item in input_items:
if not _is_json_object(item):
continue
if item.get("type") == "message" and item.get("role") == "user":
content = item.get("content", [])
if isinstance(content, str):
self.input_messages.append({"role": "user", "content": content})
elif _is_json_array(content):
for c in content:
if _is_json_object(c) and c.get("type") == "input_text":
text = c.get("text", "")
if text:
self.input_messages.append({"role": "user", "content": text})
except (json.JSONDecodeError, AttributeError, TypeError):
pass
def _store_input(self, message: object) -> None:
self._collect_input_from_client_event(message)
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
async def _log_messages(self) -> None:
if not self.logging_obj:
return
if self.input_messages:
self.logging_obj.model_call_details["messages"] = self.input_messages
if self.messages:
asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True))
async def backend_to_client(self) -> None:
"""Forward events from backend WebSocket to the client."""
import websockets
try:
while True:
try:
raw_response = await self.backend_ws.recv(decode=False)
except TypeError:
raw_response = await self.backend_ws.recv()
if isinstance(raw_response, bytes):
response_str = raw_response.decode("utf-8")
else:
response_str = raw_response
# When apply_to_output masking is active, suppress delta events
# and the text-bearing "done" events. Per-fragment Presidio
# cannot reliably catch PII spanning multiple delta chunks (e.g.
# "alice@" + "example.com"), and the done events carry the full
# output text that response.completed already delivers in
# fully-masked form; forwarding them would leak unmasked PII
# before response.completed arrives. The client receives only the
# masked response.completed.
if self.output_guardrail_callbacks:
try:
_evt_payload: Mapping[str, object] = _load_json_object(response_str)
_evt_type = _evt_payload.get("type")
except (json.JSONDecodeError, TypeError):
_evt_type = None
if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES:
continue
unmasked_str = self._unmask_response_event(response_str)
output_masked_str = await self._mask_response_completed(unmasked_str)
# Log the output-masked form so PII redacted by apply_to_output
# guardrails does not appear in success logs.
self._store_event(output_masked_str)
await self.websocket.send_text(output_masked_str)
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.debug("Responses WS backend connection closed: %s", e)
except Exception as e:
verbose_logger.exception("Error in responses WS backend_to_client: %s", e)
finally:
await self._log_messages()
def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool:
"""
Overwrite any ``model`` field in a ``response.create`` frame with the
connection-authorized model to prevent deployment-substitution attacks.
Handles both shapes:
flat: ``{"type": "response.create", "model": "...", ...}``
nested: ``{"type": "response.create", "response": {"model": "...", ...}}``
Returns True if the object was modified.
"""
if not self.authorized_model:
return False
modified = False
nested: Final = msg_obj.get("response")
if _is_json_object(nested):
if nested.get("model") != self.authorized_model:
nested["model"] = self.authorized_model
modified = True
if "model" in msg_obj and msg_obj["model"] != self.authorized_model:
msg_obj["model"] = self.authorized_model
modified = True
elif msg_obj.get("model") != self.authorized_model:
msg_obj["model"] = self.authorized_model
modified = True
return modified
async def _mask_response_create(self, message: str) -> str:
"""
Enforce the authorized model and apply Presidio PII masking to a
``response.create`` message before it is forwarded to the upstream
provider.
- Overwrites any ``model`` field with the connection-authorized model
to prevent deployment-substitution attacks (always applied).
- Walks the ``input`` and ``instructions`` fields, calls ``check_pii``
on every text block, and stores the resulting ``pii_tokens`` map in
``self.request_data["metadata"]`` for later unmasking.
Non-``response.create`` messages are returned unchanged.
"""
try:
msg_obj: Final = _load_json_object(message)
except (json.JSONDecodeError, TypeError):
return message
if msg_obj.get("type") != "response.create":
return message
# Always enforce the authorized model, even when PII masking is off.
model_modified: Final = self._enforce_authorized_model(msg_obj)
if not self.guardrail_callbacks:
return json.dumps(msg_obj) if model_modified else message
if "metadata" not in self.request_data:
self.request_data["metadata"] = {}
modified = model_modified
guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks)
for cb in guardrail_cbs:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
# response.create carries client text in two shapes:
# flat: {"type": "response.create", "input": ..., "instructions": ...}
# nested: {"type": "response.create", "response": {"input": ..., "instructions": ...}}
# Mask "input" and "instructions" in both shapes so PII is never
# forwarded unmasked regardless of where the client places it.
nested_candidate = msg_obj.get("response")
nested_response = nested_candidate if _is_json_object(nested_candidate) else None
text_containers: list[tuple[_MutableJsonObject, str]] = []
for container in (msg_obj, nested_response):
if container is None:
continue
if "input" in container:
text_containers.append((container, "input"))
if isinstance(container.get("instructions"), str):
text_containers.append((container, "instructions"))
for container, key in text_containers:
field_value = container[key]
if isinstance(field_value, str):
container[key] = await cb.check_pii(
text=field_value,
output_parse_pii=True,
presidio_config=presidio_config,
request_data=self.request_data,
)
modified = True
elif _is_json_array(field_value):
for item in field_value:
if not _is_json_object(item):
continue
for item_field in ("content", "output"):
value = item.get(item_field)
if isinstance(value, str):
item[item_field] = await cb.check_pii(
text=value,
output_parse_pii=True,
presidio_config=presidio_config,
request_data=self.request_data,
)
modified = True
elif _is_json_array(value):
for block in value:
if not _is_json_object(block):
continue
block_text = block.get("text")
if block.get("type") in RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES and isinstance(
block_text, str
):
block["text"] = await cb.check_pii(
text=block_text,
output_parse_pii=True,
presidio_config=presidio_config,
request_data=self.request_data,
)
modified = True
return json.dumps(msg_obj) if modified else message
# Delta event types whose ``delta`` field may contain PII tokens.
_DELTA_EVENT_TYPES = frozenset(
{
"response.output_text.delta",
"response.reasoning_summary_text.delta",
"response.refusal.delta",
"response.function_call_arguments.delta",
}
)
# Terminal events that carry the full output text or tool-call arguments
# already delivered by ``response.completed``. Suppressed when output masking
# is active so the unmasked copy never reaches the client before the masked
# completed event.
_OUTPUT_DONE_EVENT_TYPES = frozenset(
{
"response.output_text.done",
"response.content_part.done",
"response.output_item.done",
"response.function_call_arguments.done",
"response.reasoning_summary_text.done",
"response.reasoning_summary_part.done",
}
)
def _unmask_response_event(self, response_str: str) -> str:
"""
Apply Presidio PII unmasking to backend events before forwarding to
the client.
Handles two shapes:
- ``response.completed``: walks ``response.output[*].content[*].text``
- streaming delta events (``response.output_text.delta``, etc.):
replaces tokens in the ``delta`` field
Uses the ``pii_tokens`` map stored during ``_mask_response_create`` to
replace every token (e.g. ``<EMAIL_ADDRESS_1>``) with the original
value. Events with no stored tokens are returned unchanged.
"""
if not self.guardrail_callbacks:
return response_str
metadata: Final = self.request_data.get("metadata")
raw_pii_tokens: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None
pii_tokens: Final[Mapping[str, str]] = raw_pii_tokens if _is_str_mapping(raw_pii_tokens) else {}
if not pii_tokens:
return response_str
try:
evt_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
cb: Final = self.guardrail_callbacks[0]
unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR)
event_type: Final = evt_obj.get("type")
if event_type == "response.completed":
modified = False
response_obj: Final = evt_obj.get("response")
if not _is_json_object(response_obj):
return response_str
output_items: Final = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
return json.dumps(evt_obj) if modified else response_str
if event_type in self._DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
return response_str
async def _mask_response_completed(self, response_str: str) -> str:
"""
Apply Presidio output masking (apply_to_output=True) to the
``response.completed`` event before it is forwarded to the client.
Walks ``response.output[*].content[*].text`` and masks every text block,
as well as ``response.output[*].arguments`` on function-call items and
``response.output[*].summary[*].text`` on reasoning items. Delta and
``*.done`` events are suppressed upstream in ``backend_to_client`` when
output masking is active, so only the authoritative full-output view
reaches this method; events of other types are returned unchanged.
"""
if not self.output_guardrail_callbacks:
return response_str
try:
evt_obj: Final[Mapping[str, object]] = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
if evt_obj.get("type") != "response.completed":
return response_str
modified = False
for cb in self.output_guardrail_callbacks:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
response_obj = evt_obj.get("response")
if not _is_json_object(response_obj):
continue
output_items = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
arguments = output_item.get("arguments")
if isinstance(arguments, str):
masked_args = await cb.check_pii(
text=arguments,
output_parse_pii=False,
presidio_config=presidio_config,
request_data=self.request_data,
)
if masked_args != arguments:
output_item["arguments"] = masked_args
modified = True
summary = output_item.get("summary")
if _is_json_array(summary):
for summary_block in summary:
if not _is_json_object(summary_block):
continue
summary_text = summary_block.get("text")
if isinstance(summary_text, str):
masked_summary = await cb.check_pii(
text=summary_text,
output_parse_pii=False,
presidio_config=presidio_config,
request_data=self.request_data,
)
if masked_summary != summary_text:
summary_block["text"] = masked_summary
modified = True
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
masked = await cb.check_pii(
text=text,
output_parse_pii=False,
presidio_config=presidio_config,
request_data=self.request_data,
)
if masked != text:
content_block["text"] = masked
modified = True
return json.dumps(evt_obj) if modified else response_str
async def _enforce_or_reject_frame(self, message: str) -> bool:
"""Run the per-frame project quota check.
On rejection, sends an ``error`` event to the client and reports that
the frame must be dropped instead of forwarded, so the connection
stays open for the client to retry once the window resets.
"""
try:
await _enforce_frame_project_quota(
self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message
)
except RateLimitError as e:
try:
await self.websocket.send_text(
json.dumps( # mutable-ok: WebSocket wire payload requires JSON objects
{ # mutable-ok: WebSocket wire payload requires JSON objects
"type": "error",
"error": { # mutable-ok: nested WebSocket error object
"type": "rate_limit_exceeded",
"message": str(e),
},
}
)
)
except Exception: # noqa: BLE001, S110 # client may already be gone
pass
return False
return True
async def client_to_backend(self) -> None:
"""Forward response.create events from client to backend."""
try:
if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message):
masked_first: Final = await self._mask_response_create(self.first_message)
self._store_input(masked_first)
self._store_event(masked_first)
await self.backend_ws.send(masked_first)
while True:
message = await self.websocket.receive_text()
if not await self._enforce_or_reject_frame(message):
continue
masked = await self._mask_response_create(message)
self._store_input(masked)
self._store_event(masked)
await self.backend_ws.send(masked)
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
"""Run both forwarding directions concurrently."""
forward_task: Final = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
except Exception:
pass
finally:
if not forward_task.done():
forward_task.cancel()
try:
await forward_task
except asyncio.CancelledError:
pass
try:
await self.backend_ws.close()
except Exception:
pass
# ---------------------------------------------------------------------------
# Managed WebSocket mode (HTTP-backed, provider-agnostic)
# ---------------------------------------------------------------------------
_RESPONSE_CREATE_PARAMS: Final[frozenset[str]] = (
_get_openai_response_types().ResponsesAPIRequestParams.__required_keys__
| _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__
)
_MANAGED_WS_SKIP_KWARGS: Final[frozenset[str]] = frozenset(
{
"litellm_logging_obj",
"litellm_call_id",
"aresponses",
"_aresponses_websocket",
"user_api_key_dict",
}
)
_WARMUP_RESPONSE_ID_PREFIX: Final = "resp_warmup_"
class ManagedResponsesWebSocketHandler:
"""
Handles Responses API WebSocket mode for providers that do not expose a
native ``wss://`` responses endpoint.
Instead of proxying to a provider WebSocket, this handler:
- Listens for ``response.create`` events from the client
- Makes HTTP streaming calls via ``litellm.aresponses(stream=True)``
- Serialises and forwards every streaming event back over the WebSocket
- Supports ``previous_response_id`` for multi-turn conversations via
in-memory session tracking (avoids async DB-write timing issues)
- Supports sequential requests over a single persistent connection
This makes every provider that LiteLLM can reach over HTTP available on
the WebSocket transport without any provider-specific changes.
"""
def __init__(
self,
websocket: ResponsesClientWebSocket,
model: str,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: UserAPIKeyAuth | None = None,
litellm_metadata: Mapping[str, object] | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | None = None,
custom_llm_provider: str | None = None,
first_message: str | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
**kwargs: object,
) -> None:
self.websocket = websocket
self.model = model
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.litellm_metadata: Mapping[str, object] = litellm_metadata or {}
raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get(
"deployment_model_name"
)
self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None
self.api_key = api_key
self.api_base = api_base
self.timeout = timeout
self.custom_llm_provider = custom_llm_provider
self._connection_provider = self._resolve_provider(model) or custom_llm_provider
self.first_message = first_message
self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
# Carry through safe pass-through kwargs (e.g. extra_headers)
self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS}
# In-memory session history: response_id → full accumulated message list.
# Keyed by the DECODED (pre-encoding) response ID from response.completed.
# This avoids the async DB-write race condition where spend logs haven't
# been committed yet when the next response.create arrives.
self._session_history: dict[str, list[dict[str, object]]] = {}
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
@staticmethod
def _serialize_chunk(chunk: object) -> str | None:
"""Serialize a streaming chunk to a JSON string for WebSocket transmission."""
try:
if isinstance(chunk, _HasModelDumpJson):
return chunk.model_dump_json(exclude_none=True)
if isinstance(chunk, _HasModelDump):
return json.dumps(chunk.model_dump(exclude_none=True), default=str)
if _is_json_object(chunk):
return json.dumps(chunk, default=str)
return json.dumps(str(chunk))
except Exception as exc:
verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc)
return None
async def _send_error(self, message: str, error_type: str = "server_error") -> None:
try:
await self.websocket.send_text(
json.dumps({"type": "error", "error": {"type": error_type, "message": message}})
)
except Exception:
pass
def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]:
"""
Return accumulated message history for *previous_response_id*.
The key is the *decoded* response ID (the raw provider response ID before
LiteLLM base64-encodes it into the ``resp_...`` format).
"""
decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id)
raw_id: Final = decoded.get("response_id", previous_response_id)
return list(self._session_history.get(raw_id, []))
def _store_history(self, response_id: str, messages: list[dict[str, object]]) -> None:
"""
Store the complete accumulated message history for *response_id*.
Replaces any prior value — callers are responsible for passing the full
history (prior turns + current input + new output).
"""
self._session_history[response_id] = messages
@staticmethod
def _extract_response_id(completed_event: _MutableJsonObject) -> str | None:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Returns *None* if the event doesn't contain a usable ID.
"""
resp_obj: Final = completed_event.get("response", {})
raw_id: Final = resp_obj.get("id") if _is_json_object(resp_obj) else None
encoded_id: Final[str | None] = raw_id if isinstance(raw_id, str) else None
if not encoded_id:
return None
decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id)
return decoded.get("response_id", encoded_id)
@staticmethod
def _extract_output_messages(
completed_event: _MutableJsonObject,
) -> list[dict[str, object]]:
"""
Convert the output items in a ``response.completed`` event into
Responses API message dicts suitable for the next turn's ``input``.
"""
resp_obj: Final = completed_event.get("response", {})
if not isinstance(resp_obj, dict):
return []
messages: Final[list[dict[str, object]]] = []
for item in resp_obj.get("output", []) or []:
if not isinstance(item, dict):
continue
item_type = item.get("type")
role = item.get("role", "assistant")
if item_type == "message":
content_parts = item.get("content") or []
text_parts = [
p.get("text", "")
for p in content_parts
if isinstance(p, dict) and p.get("type") in ("output_text", "text")
]
text = "".join(text_parts)
if text:
messages.append(
{
"type": "message",
"role": role,
"content": [{"type": "output_text", "text": text}],
}
)
elif item_type == "function_call":
messages.append(item)
return messages
@staticmethod
def _input_to_messages(input_val: object) -> list[dict[str, object]]:
"""
Normalise the ``input`` field of a ``response.create`` event to a list
of Responses API message dicts.
"""
if isinstance(input_val, str):
return [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": input_val}],
}
]
if _is_json_array(input_val):
return [item for item in input_val if _is_json_object(item)]
return []
# ------------------------------------------------------------------
# _process_response_create sub-methods
# ------------------------------------------------------------------
async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj: Final = _load_json_object(raw_message)
except json.JSONDecodeError:
await self._send_error("Invalid JSON in response.create event", "invalid_request_error")
return None
if msg_obj.get("type") != "response.create":
# Silently ignore non-response.create messages (e.g. warmup pings)
return None
return msg_obj
@staticmethod
def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool:
"""Return True for a response.create whose generate flag is false."""
nested: Final = msg_obj.get("response")
source: Final = nested if _is_json_object(nested) and nested else msg_obj
return source.get("generate") is False
@staticmethod
def _is_warmup_response_id(response_id: str | None) -> bool:
"""Return True for synthetic warmup IDs that only exist on this connection."""
if not response_id:
return False
decoded: Final = ResponsesAPIRequestUtils._decode_responses_api_response_id(response_id)
raw_id: Final = decoded.get("response_id", response_id)
return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX)
@staticmethod
def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]:
nested: Final = msg_obj.get("response")
if _is_json_object(nested) and nested:
return nested
return {k: v for k, v in msg_obj.items() if k != "type"}
def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]:
"""Build a minimal completed Responses API object for a warmup ack."""
source: Final = self._warmup_source_params(msg_obj)
wire_model: Final = source.get("model") or self.model_group or self.model
return {
"id": f"{_WARMUP_RESPONSE_ID_PREFIX}{uuid.uuid4().hex}",
"object": "response",
"created_at": int(time.time()),
"status": "completed",
"model": wire_model,
"output": [],
"usage": {
"input_tokens": 0,
"output_tokens": 0,
"total_tokens": 0,
},
}
async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None:
"""
Acknowledge a generate=false prewarm without calling the provider.
Codex blocks on the warmup turn until it receives response.created and
response.completed over the WebSocket. Managed HTTP providers cannot
honor an empty-input warmup, so we synthesize the completion locally.
"""
response: Final = self._build_warmup_response(msg_obj)
for event_type, status in (
("response.created", "in_progress"),
("response.completed", "completed"),
):
event = {
"type": event_type,
"response": {**response, "status": status},
}
serialized = self._serialize_chunk(event)
if serialized is None:
continue
await self.websocket.send_text(serialized)
@staticmethod
def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]:
"""
Extract Responses API params from the event, handling both wire formats:
Nested: {"type": "response.create", "response": {"input": [...], ...}}
Flat: {"type": "response.create", "input": [...], "model": "...", ...}
"""
nested: Final = msg_obj.get("response")
response_params: Final[dict[str, object]] = (
nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"}
)
return {
param: response_params[param]
for param in _RESPONSE_CREATE_PARAMS
if param in response_params and response_params[param] is not None
}
def _apply_history(
self,
call_kwargs: dict[str, object],
previous_response_id: str | None,
current_messages: list[dict[str, object]],
prior_history: list[dict[str, object]],
) -> None:
"""Prepend in-memory turn history, or fall back to DB-based reconstruction."""
if not previous_response_id:
return
if self._is_warmup_response_id(previous_response_id):
verbose_logger.debug(
"ManagedResponsesWS: ignoring synthetic warmup previous_response_id=%s",
previous_response_id,
)
return
if prior_history:
call_kwargs["input"] = prior_history + current_messages
verbose_logger.debug(
"ManagedResponsesWS: prepended %d history messages for previous_response_id=%s",
len(prior_history),
previous_response_id,
)
else:
verbose_logger.debug(
"ManagedResponsesWS: no in-memory history for previous_response_id=%s; "
"falling back to DB-based session reconstruction",
previous_response_id,
)
# Fall back to DB-based session reconstruction (may work for
# cross-connection multi-turn when spend logs are committed)
call_kwargs["previous_response_id"] = previous_response_id
@staticmethod
def _resolve_provider(model: str | None) -> str | None:
"""Resolve the LLM provider for a model string, or None if unresolvable."""
if not model:
return None
try:
from litellm import get_llm_provider
_, provider, _, _ = get_llm_provider(model=model)
return provider
except Exception:
return None
def _same_provider(self, model: str | None) -> bool:
"""Return True if model uses the same LLM provider as the connection model."""
if model is None or model == self.model:
return True
event_provider: Final = self._resolve_provider(model)
if event_provider is None:
return False
return event_provider == self._connection_provider
def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None:
"""Inject connection-level credentials and metadata into call_kwargs."""
if self.api_key is not None:
call_kwargs["api_key"] = self.api_key
if self.api_base is not None:
call_kwargs["api_base"] = self.api_base
if self.timeout is not None:
call_kwargs["timeout"] = self.timeout
# Only force connection-level custom_llm_provider when the per-event model
# uses the same provider as the connection model. If the provider differs
# (e.g., connection is vertex_ai but event says openai/gpt-4), let litellm
# re-resolve from the model string. Same-provider model variants (e.g.,
# vertex_ai/gemini-2.0 -> vertex_ai/gemini-1.5) still inherit the provider.
if self.custom_llm_provider is not None and self._same_provider(model):
call_kwargs["custom_llm_provider"] = self.custom_llm_provider
if self.litellm_metadata:
call_kwargs["litellm_metadata"] = dict(self.litellm_metadata)
@staticmethod
def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None:
"""Update proxy_server_request body so spend logs record the full request."""
proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {}
if not isinstance(proxy_server_request, dict):
return
body: Final = dict(proxy_server_request.get("body") or {})
body["input"] = call_kwargs.get("input")
body["store"] = call_kwargs.get("store")
body["model"] = model
for k in ("tools", "tool_choice", "instructions", "metadata"):
if k in call_kwargs and call_kwargs[k] is not None:
body[k] = call_kwargs[k]
proxy_server_request = {**proxy_server_request, "body": body}
if "litellm_metadata" not in call_kwargs:
call_kwargs["litellm_metadata"] = {}
call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
Captures the ``response.completed`` event type from the chunk object
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
"""
completed_event: _MutableJsonObject | None = (
None # rebind-ok: captures the completed event once the stream yields it
)
stream_response: Final = await litellm.aresponses(model=model, **call_kwargs)
async for chunk in stream_response:
if chunk is None:
continue
# Read type from the object before serializing to avoid double JSON parse
chunk_type = getattr(chunk, "type", None) or (chunk.get("type") if isinstance(chunk, dict) else None)
serialized = self._serialize_chunk(chunk)
if serialized is None:
continue
if chunk_type == "response.completed" and completed_event is None:
try:
completed_event = _load_json_object(serialized)
except Exception:
pass
try:
await self.websocket.send_text(serialized)
except Exception as send_exc:
verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc)
return completed_event # Client disconnected
return completed_event
def _save_turn_history(
self,
completed_event: _MutableJsonObject | None,
prior_history: list[dict[str, object]],
current_messages: list[dict[str, object]],
) -> None:
"""Store this turn in in-memory history for future previous_response_id lookups."""
if completed_event is None:
return
new_response_id: Final = self._extract_response_id(completed_event)
if not new_response_id:
return
output_msgs: Final = self._extract_output_messages(completed_event)
all_messages: Final = prior_history + current_messages + output_msgs
self._store_history(new_response_id, all_messages)
verbose_logger.debug(
"ManagedResponsesWS: stored %d messages for response_id=%s",
len(all_messages),
new_response_id,
)
# ------------------------------------------------------------------
# Core request handler
# ------------------------------------------------------------------
async def _process_response_create(self, raw_message: str) -> None:
"""
Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``,
and forward every streaming event to the client.
Multi-turn support via in-memory session history
------------------------------------------------
When ``previous_response_id`` is present in the event:
1. Look up the accumulated message history in ``self._session_history``
(keyed by the decoded provider response ID).
2. Prepend those messages to the current ``input`` so the model has full
conversation context.
3. After the stream completes, extract the new response ID and output
messages from ``response.completed`` and store them in
``self._session_history`` for the next turn.
This in-memory approach avoids the async DB-write race condition that
occurs when spend logs haven't been committed by the time the second
``response.create`` arrives over the same WebSocket connection.
"""
msg_obj: Final = await self._parse_message(raw_message)
if msg_obj is None:
return
# generate=false is a prompt-cache warmup hint (sent by codex prewarm).
# Native provider sockets handle it server-side, but there is no HTTP
# equivalent and the frame carries empty input. Managed providers must
# synthesize a completion so clients like Codex can proceed.
if self._is_warmup_frame(msg_obj):
try:
await self._send_warmup_ack(msg_obj)
except Exception as exc:
verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc)
return
try:
await _enforce_frame_project_quota(
self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message
)
except RateLimitError as e:
await self._send_error(str(e), error_type="rate_limit_exceeded")
return
call_kwargs: Final = self._build_base_call_kwargs(msg_obj)
call_kwargs["stream"] = True
# A frame that repeats the connection's public alias (model_group) must
# reuse the router-resolved self.model; passing the alias raw to
# litellm.aresponses fails in get_llm_provider. A genuinely different
# provider-prefixed per-frame model is still honored.
requested_model: Final[str | None] = _optional_str(call_kwargs.pop("model", None))
model: Final[str] = (
self.model if requested_model is None or requested_model == self.model_group else requested_model
)
previous_response_id: Final[str | None] = _optional_str(call_kwargs.pop("previous_response_id", None))
current_messages: Final = self._input_to_messages(call_kwargs.get("input"))
# Fetch history once; reused in both _apply_history and _save_turn_history
prior_history: Final = self._get_history_messages(previous_response_id) if previous_response_id else []
self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history)
self._inject_credentials(call_kwargs, model=model)
self._update_proxy_request(call_kwargs, requested_model or self.model_group or model)
call_kwargs.update(self.extra_kwargs)
try:
completed_event: Final = await self._stream_and_forward(model, call_kwargs)
except Exception as exc:
verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc)
await self._send_error(str(exc))
return
self._save_turn_history(completed_event, prior_history, current_messages)
# ------------------------------------------------------------------
# Main entry point
# ------------------------------------------------------------------
async def run(self) -> None:
"""
Main loop: accept ``response.create`` events sequentially and handle
each one before waiting for the next message.
"""
try:
if self.first_message is not None:
await self._process_response_create(self.first_message)
while True:
try:
message = await self.websocket.receive_text()
except Exception as exc:
verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc)
break
await self._process_response_create(message)
except Exception as exc:
verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc)
await self._send_error(f"Internal server error: {exc}")
Reference in a new issue View git blame Copy permalink