scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

Merge pull request #38722 from BerriAI/litellm_bedrock_guardrail_stream_audit

Browse source 
feat(bedrock): honor streaming buffer/sampling config for unbuffered post_call scans
This commit is contained in:
Mateo Wang 2026-08-30 10:17:11 -07:00 • committed by GitHub
commit 4ba8517134
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 1287 additions and 162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -58,6 +58,8 @@ from litellm.types.utils import (
)
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -207,6 +209,19 @@ class AnthropicMessagesHandler(BaseTranslation):
return list(self._block_continuation_chunks(exc, responses_so_far or []))
return self._standalone_block_chunks(exc)
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
from litellm.proxy.common_request_processing import (
serialize_http_exception_detail,
)
from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames
message, _ = serialize_http_exception_detail(exc.detail)
return tuple(anthropic_sse_error_frames(message))
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
import uuid

48
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -1,8 +1,11 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, Optional
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -73,6 +76,31 @@ class BaseTranslation(ABC):
return transformed
@staticmethod
def merge_user_api_key_metadata_into_request(
request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place
user_api_key_dict: Optional["UserAPIKeyAuth"],
) -> None:
"""
Add the prefixed ``user_api_key_*`` metadata to the request's resolved
metadata bucket without overwriting existing keys.
Writes must go through ``get_or_create_metadata_bucket``: creating a
``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat
completions) flips the bucket for every later metadata write, and spend
logging never sees those writes (e.g. guardrail_information).
"""
from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
)
user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if not user_metadata:
return
_, metadata_bucket = get_or_create_metadata_bucket(request_data)
for key, value in user_metadata.items():
metadata_bucket.setdefault(key, value)
@abstractmethod
async def process_input_messages(
self,
@ -147,6 +175,26 @@ class BaseTranslation(ABC):
"""
return None
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
"""
Build the stream items that surface a guardrail HTTPException (a block
with the default exception-on-block config, or a failed scan) after the
response has already started streaming, in this endpoint's wire format.
Called only once chunks have been sent: the HTTP status is gone, so the
failure must travel as an in-stream error frame. ``responses_so_far``
holds the chunks the client has already received, for formats whose
error frame continues the stream (e.g. sequence numbers).
Returns None when the format has no in-stream error frame; the caller
then re-raises ``exc``. Override in endpoint subclasses.
"""
return None
def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None:
"""
Convert request data to OpenAI-spec structured messages.

32
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -14,6 +14,7 @@ Pattern Overview:
This pattern can be replicated for other message formats (e.g., Anthropic).
"""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Union, cast
import litellm
@ -46,6 +47,8 @@ from litellm.types.utils import (
)
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -382,11 +385,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "response" not in request_data:
request_data["response"] = response
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
@ -555,11 +554,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "responses" not in request_data:
request_data["responses"] = responses_so_far
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
@ -591,6 +586,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return responses_so_far
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
import json
from litellm.proxy.common_request_processing import sse_error_payload
_, error_obj = sse_error_payload(exc)
return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
@staticmethod
def _accumulate_string_content_by_choice_index(
responses_so_far: list["ModelResponseStream"],
@ -653,10 +660,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
request_data = {"responses": responses_so_far}
elif "responses" not in request_data:
request_data["responses"] = responses_so_far
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if responses_so_far and getattr(responses_so_far[0], "model", None):

35
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -48,6 +48,8 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
ErrorEvent,
ErrorEventError,
OpenAIMcpServerTool,
ResponsesAPIStreamEvents,
)
@ -59,6 +61,8 @@ from litellm.types.responses.main import (
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False):
text: ReadOnly[str]
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
sequence_numbers: Final = (
item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None)
for item in reversed(responses_so_far or ())
)
return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0)
class OpenAIResponsesHandler(BaseTranslation):
"""
Handler for processing OpenAI Responses API with guardrails.
@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation):
}
return responses_so_far[-1].get("type") in terminal_types
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
from litellm.proxy.common_request_processing import (
serialize_http_exception_detail,
)
message, _ = serialize_http_exception_detail(exc.detail)
return (
ErrorEvent(
type=ResponsesAPIStreamEvents.ERROR,
sequence_number=_next_stream_sequence_number(responses_so_far),
error=ErrorEventError(
type="guardrail_error",
code=str(exc.status_code),
message=message,
param=None,
),
),
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
"""
Get the string so far from the responses so far.

151
litellm/proxy/common_request_processing.py
View file

@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch
return logging_obj
def _serialize_http_exception_detail(
def serialize_http_exception_detail(
detail: object,
) -> tuple[str, dict | None]:
"""
@ -535,7 +535,7 @@ def _serialize_http_exception_detail(
def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException:
raw_detail: Final = _getattr_object(exc, "detail", str(exc))
message, structured_fields = _serialize_http_exception_detail(raw_detail)
message, structured_fields = serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
return ProxyException(
@ -818,7 +818,7 @@ async def _buffer_first_chunk_honoring_disconnect(
raise _ClientDisconnectedBeforeFirstChunk()
def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
"""Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames.
Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames
@ -827,7 +827,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
# Preserve status code from HTTPException (e.g. guardrail blocks)
error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start")
message, structured_fields = _serialize_http_exception_detail(raw_detail)
message, structured_fields = serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
@ -942,7 +942,7 @@ async def create_response(
# Unexpected error consuming first chunk.
verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e)
error_status, error_obj = _sse_error_payload(e)
error_status, error_obj = sse_error_payload(e)
async def error_gen_message() -> AsyncGenerator[str, None]:
for frame in _sse_error_frames(error_obj):
@ -1119,7 +1119,7 @@ async def open_sse_before_first_byte(
# would never fire and the failure would go unaudited. The hook
# also gets to sanitize what reaches the client, by returning or
# raising a replacement, so its answer decides the frame.
_, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
_, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
for frame in _sse_error_frames(error_obj):
yield frame.encode()
return
@ -2380,53 +2380,14 @@ class ProxyBaseLLMRequestProcessing:
if requested_model_from_client:
self.data["_litellm_client_requested_model"] = requested_model_from_client
# Streaming: attach a closure that fires after all guardrail
# end-of-stream blocks complete. CSW.__anext__ stores the
# assembled response on logging_obj; the outer consumer
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
# assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
CustomStreamWrapper,
)
if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper):
# Intentionally a live reference (not a copy) — mirrors
# ProxyLogging.post_call_success_hook which also mutates
# data["guardrail_to_apply"] during iteration.
_captured_data: Final = self.data
_captured_user_api_key_dict: Final = user_api_key_dict
_captured_logging_obj: Final = logging_obj
async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None:
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data=_captured_data,
captured_user_api_key_dict=_captured_user_api_key_dict,
captured_logging_obj=_captured_logging_obj,
assembled_response=assembled_response,
cache_hit=cache_hit,
)
logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete
elif (
_post_call_guardrails_active
and route_type == "anthropic_messages"
and self._is_streaming_response(response)
):
from litellm.litellm_core_utils.logging_worker import (
GLOBAL_LOGGING_WORKER,
if _post_call_guardrails_active:
self._arm_deferred_stream_dispatch(
response=response,
route_type=route_type,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
)
async def _on_deferred_native_stream_complete(
logging_coroutine: Coroutine[object, object, object],
) -> None:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete
if route_type == "allm_passthrough_route":
# Check if response is an async generator
if self._is_streaming_response(response):
@ -3096,6 +3057,94 @@ class ProxyBaseLLMRequestProcessing:
except Exception as e:
verbose_proxy_logger.exception("Error firing deferred logging: %s", e)
def _arm_deferred_stream_dispatch(
self,
response: object,
route_type: str,
user_api_key_dict: "UserAPIKeyAuth",
logging_obj: LiteLLMLoggingObj,
) -> None:
"""
Streaming with post-call guardrails active: attach a closure that
ProxyLogging._fire_deferred_stream_logging fires after all guardrail
end-of-stream blocks complete, so the spend log sees
guardrail_information.
Three closure shapes, matching who owns logging for the stream:
- CustomStreamWrapper (chat completions) stores
(assembled_response, cache_hit); the closure also runs
non-apply_guardrail post-call hooks via
_run_deferred_stream_guardrails.
- Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares
its inner CustomStreamWrapper's logging_obj, so it stores the same
(assembled_response, cache_hit) shape; the closure only dispatches
success logging, matching the route's pre-existing hook surface.
- Native anthropic_messages/aresponses iterators store a single
ready-made logging coroutine to enqueue.
Raw async generators from passthrough routes bypass all three and
would orphan the closure, so they are not armed here.
The router wraps iterators that cannot carry _hidden_params in
HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the
unwrapped inner iterator.
"""
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper
unwrapped: Final = response._inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response
if isinstance(unwrapped, CustomStreamWrapper):
# Intentionally a live reference (not a copy) — mirrors
# ProxyLogging.post_call_success_hook which also mutates
# data["guardrail_to_apply"] during iteration.
_captured_data: Final = self.data
_captured_user_api_key_dict: Final = user_api_key_dict
_captured_logging_obj: Final = logging_obj
async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None:
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data=_captured_data,
captured_user_api_key_dict=_captured_user_api_key_dict,
captured_logging_obj=_captured_logging_obj,
assembled_response=assembled_response,
cache_hit=cache_hit,
)
logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete
return
if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response):
return
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
if isinstance(unwrapped, LiteLLMCompletionStreamingIterator):
_captured_bridge_logging_obj: Final = logging_obj
async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None:
await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers(
assembled_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
async def _on_deferred_native_stream_complete(
logging_coroutine: Coroutine[object, object, object],
) -> None:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete
@staticmethod
async def _run_deferred_stream_guardrails(
captured_data: dict,

53
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -18,6 +18,7 @@ import time
from collections.abc import AsyncGenerator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import accumulate, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
import httpx
@ -42,7 +43,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import _serialize_http_exception_detail
from litellm.proxy.common_request_processing import serialize_http_exception_detail
from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired
from litellm.proxy.guardrails.anthropic_sse import (
anthropic_sse_chunks_from_response,
@ -52,7 +53,12 @@ from litellm.proxy.guardrails.anthropic_sse import (
model_response_text,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks
from litellm.types.guardrails import (
BedrockChecksConfigModel,
BedrockGuardrailStreamingParams,
GuardrailEventHooks,
LitellmParams,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockChecksMessage,
@ -221,9 +227,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
streaming_buffer_until_moderated: bool | None = None,
streaming_sampling_rate: int | None = None,
streaming_end_of_stream_only: bool | None = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self._set_streaming_params(
BedrockGuardrailStreamingParams.from_extras(
MappingProxyType(
{
"streaming_buffer_until_moderated": streaming_buffer_until_moderated,
"streaming_sampling_rate": streaming_sampling_rate,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
}
)
)
)
self.guardrailIdentifier = guardrailIdentifier
self.guardrailVersion = guardrailVersion
self.guardrail_provider = "bedrock"
@ -278,6 +298,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
list(self.checks.keys()) if self.checks else None,
)
def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None:
self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated
self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
super().update_in_memory_litellm_params(litellm_params)
self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
def _streams_incrementally(self) -> bool:
return not self.streaming_buffer_until_moderated and not self.mask_response_content
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
return [
@ -2660,6 +2692,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
Collect content from the stream and run the bedrock OUTPUT scan
(post_call only validates the response).
"""
if self._streams_incrementally():
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
guardrail_to_apply=self,
buffer_until_moderated_default=False,
):
yield streamed_chunk
return
# Import here to avoid circular imports
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
@ -2716,7 +2763,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
if not raw_sse or (not is_block and not headers_flushed):
raise
block_message, _ = _serialize_http_exception_detail(block_detail)
block_message, _ = serialize_http_exception_detail(block_detail)
for error_frame in anthropic_sse_error_frames(
block_message if is_block else f"{block_exc.status_code}: {block_message}"
):

67
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol):
@property
def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
@property
def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ...
def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation:
return translation
@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type: str | None,
responses_so_far: Sequence[object],
request_data: dict,
endpoint_translation: _EndpointTranslation | None = None,
stream_started: bool = False,
responses_yielded: Sequence[object] | None = None,
) -> AsyncGenerator[object, None]:
"""Surface a mid-stream HTTPException. For A2A call types the response has
already started, so emit an in-stream JSON-RPC error chunk; otherwise
re-raise so the proxy can report it.
"""Surface a mid-stream HTTPException (a guardrail block with the default
exception-on-block config, or a failed scan).
A2A call types emit an in-stream JSON-RPC error chunk. For other call
types, once chunks have already reached the client the HTTP status is
gone, so the failure is delegated to the endpoint translation's
``build_stream_error_items`` and travels as an in-stream error frame in
that endpoint's wire format. Before the first chunk (or when the format
has no in-stream error frame) the exception is re-raised so the proxy
can report it with a real HTTP status.
"""
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data))
return
if stream_started and endpoint_translation is not None:
error_items: Final = endpoint_translation.build_stream_error_items(
exc, responses_so_far=tuple(responses_yielded) if responses_yielded is not None else None
)
if error_items is not None:
for error_item in error_items:
yield error_item
return
raise exc
def _build_transform_chunk(
@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
raise _StreamTerminated()
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data):
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=bool(responses_yielded),
responses_yielded=responses_yielded,
):
yield error_item
raise _StreamTerminated()
@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger):
return
except HTTPException as e:
# Response already started (we already yielded chunks); cannot send 400.
# For A2A, yield an in-stream JSON-RPC error so the client sees it.
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data))
return
raise
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=chunks_yielded,
responses_yielded=responses_yielded,
):
yield error_item
return
chunks_yielded = True
responses_yielded.append(original_item)
yield original_item
@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
return
except HTTPException as e:
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data))
else:
raise
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=bool(responses_yielded),
responses_yielded=responses_yielded,
):
yield error_item

4
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -11,6 +11,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
BedrockGuardrail,
)
streaming_params: Final = BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)
_bedrock_callback: Final = BedrockGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
@ -38,6 +39,9 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
only_scan_new_messages=litellm_params.only_scan_new_messages or False,
streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
streaming_sampling_rate=streaming_params.streaming_sampling_rate,
streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
)
litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)
return _bedrock_callback

41
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -65,6 +65,19 @@ class PassThroughStreamingHandler:
route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
)
raw_bytes: Final[list[bytes]] = []
def _build_logging_coroutine() -> Coroutine[None, None, None]:
return resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(),
)
logging_scheduled = False
model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
request_body=request_body,
@ -114,6 +127,21 @@ class PassThroughStreamingHandler:
)
if pending:
yield pending
# Stream completed cleanly. When the proxy armed deferred
# dispatch (post-call guardrails active), park the logging
# coroutine on logging_obj instead of enqueueing now, so
# ProxyLogging._fire_deferred_stream_logging fires it after
# guardrail end-of-stream blocks populate guardrail_information.
# Disconnect/exception paths skip this and fall through to the
# immediate enqueue in ``finally`` to keep partial billing
# (LIT-2642).
if (
getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None
and raw_bytes
and response.status_code < 400
):
logging_scheduled = True
litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
except Exception as e:
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
raise
@ -128,18 +156,7 @@ class PassThroughStreamingHandler:
if not logging_scheduled and raw_bytes and response.status_code < 400:
logging_scheduled = True
try:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(),
)
)
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine())
except Exception as e:
verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)

21
litellm/responses/streaming_iterator.py
View file

@ -463,15 +463,20 @@ class BaseResponsesAPIStreamingIterator:
end_time: Final = datetime.now()
if is_async:
asyncio.create_task(
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,
)
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,

36
litellm/types/guardrails.py
View file

@ -1,5 +1,7 @@
from collections.abc import Mapping
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Any, Final, Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
@ -550,6 +552,40 @@ class BedrockGuardrailConfigModel(BaseModel):
)
class BedrockGuardrailStreamingParams(BaseModel):
streaming_buffer_until_moderated: bool = Field(
default=True,
description="If True (default), withhold every streamed chunk until the end-of-stream "
"ApplyGuardrail scan passes, so no flagged content reaches the client before a block. "
"If False, chunks stream through unbuffered, so flagged content can reach the client "
"before the scan finishes; a flagged scan still ends the stream, with a block message "
"when disable_exception_on_block is true and an in-stream error frame otherwise.",
)
streaming_sampling_rate: int = Field(
default=5,
ge=1,
description="When not buffering and not end-of-stream-only, scan the accumulated response "
"every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays "
"that chunk, so lower values add latency and AWS text-unit cost.",
)
streaming_end_of_stream_only: bool = Field(
default=False,
description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan "
"on the assembled response at end of stream. Combined with "
"streaming_buffer_until_moderated=false the full response streams live before the scan "
"and the scan result lands in guardrail_information; a flagged response still ends the "
"stream with a block message (disable_exception_on_block=true) or an error frame.",
)
@classmethod
def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":
if not extras:
return cls()
return cls.model_validate(
MappingProxyType({name: extras[name] for name in cls.model_fields if extras.get(name) is not None})
)
class LakeraV2GuardrailConfigModel(BaseModel):
"""Configuration parameters for the Lakera AI v2 guardrail"""

32
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py
View file

@ -482,23 +482,25 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
"metadata": {"guardrails": ["test-openai-moderation"]},
}
# Should raise HTTPException when processing streaming harmful content
from fastapi import HTTPException
# Chunks have already been flushed by end-of-stream moderation, so
# the block surfaces as the in-stream error frame, not a raise.
import json as _json
async def _drain():
result_chunks = []
async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=mock_stream(),
request_data=request_data,
):
result_chunks.append(chunk)
result_chunks = []
async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=mock_stream(),
request_data=request_data,
):
result_chunks.append(chunk)
with pytest.raises(HTTPException) as exc_info:
await _drain()
assert exc_info.value.status_code == 400
assert "Violated OpenAI moderation policy" in str(exc_info.value.detail)
frame = result_chunks[-1]
assert isinstance(frame, bytes)
text = frame.decode()
assert text.startswith("data: ")
assert "Violated OpenAI moderation policy" in text
payload = _json.loads(text[len("data: ") :])
assert payload["error"]["code"] == "400"
@pytest.mark.asyncio

32
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -161,19 +161,27 @@ async def test_openai_moderation_guardrail_streaming_harmful_content():
"metadata": {"guardrails": ["test-openai-moderation"]},
}
# Should raise HTTPException
with pytest.raises(HTTPException) as exc_info:
async for (
_
) in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=mock_stream(),
request_data=request_data,
):
pass
# Chunks have already been flushed by end-of-stream moderation, so
# the block surfaces as the in-stream error frame, not a raise.
import json as _json
assert exc_info.value.status_code == 400
assert "Violated OpenAI moderation policy" in str(exc_info.value.detail)
collected = []
async for (
chunk
) in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=mock_stream(),
request_data=request_data,
):
collected.append(chunk)
frame = collected[-1]
assert isinstance(frame, bytes)
text = frame.decode()
assert text.startswith("data: ")
assert "Violated OpenAI moderation policy" in text
payload = _json.loads(text[len("data: ") :])
assert payload["error"]["code"] == "400"
@pytest.mark.asyncio

245
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -5345,3 +5345,248 @@ def test_initialize_bedrock_forwards_aws_external_id():
assert guardrail.optional_params["aws_external_id"] == "external-id-123"
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, guardrail)
def _chat_chunk(content: str, finish_reason: str | None) -> litellm.ModelResponseStream:
return litellm.ModelResponseStream(
id="tid",
choices=[
litellm.types.utils.StreamingChoices(
delta=litellm.types.utils.Delta(content=content, role="assistant"),
finish_reason=finish_reason,
index=0,
)
],
created=1,
model="gpt-4o-mini",
object="chat.completion.chunk",
)
def _streaming_litellm_params(**extras):
from litellm.types.guardrails import LitellmParams
return LitellmParams(
guardrail="bedrock",
mode="post_call",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
**extras,
)
def test_initialize_bedrock_wires_streaming_flags():
from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock
configured = initialize_bedrock(
_streaming_litellm_params(
streaming_buffer_until_moderated=False,
streaming_sampling_rate=3,
streaming_end_of_stream_only=True,
),
{"guardrail_name": "bedrock-streaming"},
)
defaulted = initialize_bedrock(
_streaming_litellm_params(),
{"guardrail_name": "bedrock-defaults"},
)
for registered in (configured, defaulted):
litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, registered)
assert configured.streaming_buffer_until_moderated is False
assert configured.streaming_sampling_rate == 3
assert configured.streaming_end_of_stream_only is True
assert defaulted.streaming_buffer_until_moderated is True
assert defaulted.streaming_sampling_rate == 5
assert defaulted.streaming_end_of_stream_only is False
def test_initialize_bedrock_rejects_non_positive_sampling_rate():
from pydantic import ValidationError
from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock
with pytest.raises(ValidationError):
initialize_bedrock(
_streaming_litellm_params(streaming_sampling_rate=0),
{"guardrail_name": "bedrock-bad-rate"},
)
def test_update_in_memory_litellm_params_round_trips_streaming_flags():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-update",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
)
guardrail.update_in_memory_litellm_params(
_streaming_litellm_params(
streaming_buffer_until_moderated=False,
streaming_sampling_rate=7,
streaming_end_of_stream_only=True,
)
)
assert guardrail.streaming_buffer_until_moderated is False
assert guardrail.streaming_sampling_rate == 7
assert guardrail.streaming_end_of_stream_only is True
guardrail.update_in_memory_litellm_params(_streaming_litellm_params())
assert guardrail.streaming_buffer_until_moderated is True
assert guardrail.streaming_sampling_rate == 5
assert guardrail.streaming_end_of_stream_only is False
async def _run_streaming_hook_recording_order(guardrail: BedrockGuardrail) -> list:
events = []
minimal = {"action": "NONE", "assessments": [], "outputs": []}
async def record_scan(*args, **kwargs):
events.append("scan")
return minimal
async def mock_stream():
yield _chat_chunk("Hello", None)
yield _chat_chunk(" world", None)
yield _chat_chunk("", "stop")
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)):
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(),
response=mock_stream(),
request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]},
):
content = chunk.choices[0].delta.content if chunk.choices else None
events.append(("chunk", content))
return events
@pytest.mark.asyncio
async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-audit-mode",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
streaming_buffer_until_moderated=False,
streaming_end_of_stream_only=True,
)
events = await _run_streaming_hook_recording_order(guardrail)
scan_index = events.index("scan")
chunk_events = [e for e in events if e != "scan"]
assert events.count("scan") == 1
assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index]
assert ("chunk", "Hello") in events[:scan_index]
assert ("chunk", " world") in events[:scan_index]
assert len(chunk_events) == 3
@pytest.mark.asyncio
async def test_buffered_default_hook_scans_before_any_chunk():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-buffered-default",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
events = await _run_streaming_hook_recording_order(guardrail)
assert events[0] == "scan"
assert all(e == "scan" or e[0] == "chunk" for e in events)
assert len([e for e in events if e != "scan"]) >= 1
@pytest.mark.asyncio
async def test_masking_keeps_buffered_path_even_when_unbuffered_configured():
guardrail = BedrockGuardrail(
guardrail_name="bedrock-mask-buffered",
guardrailIdentifier="test-id",
guardrailVersion="DRAFT",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
mask_response_content=True,
streaming_buffer_until_moderated=False,
streaming_end_of_stream_only=True,
)
assert guardrail._streams_incrementally() is False
events = await _run_streaming_hook_recording_order(guardrail)
assert events[0] == "scan"
@pytest.mark.asyncio
async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_truncating():
"""Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream
scan used to raise after SSE headers were flushed, so the client saw a
silently truncated stream. The unified hook must emit the chat in-stream
error frame instead."""
from litellm.llms import load_guardrail_translation_mappings
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
unified_guardrail as unified_module,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
guardrail = BedrockGuardrail(
guardrailIdentifier="test-guardrail",
guardrailVersion="DRAFT",
streaming_end_of_stream_only=True,
streaming_buffer_until_moderated=False,
guardrail_name="bedrock-eos",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
blocked_response = {
"action": "GUARDRAIL_INTERVENED",
"actionReason": "Guardrail blocked.",
"outputs": [{"text": "Sorry, the model cannot answer this question."}],
"assessments": [
{"topicPolicy": {"topics": [{"name": "Forbidden topic", "type": "DENY", "action": "BLOCKED"}]}}
],
}
def _chunk(content, finish_reason=None):
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta={"content": content, "role": "assistant"},
finish_reason=finish_reason,
)
],
)
async def _mock_stream():
yield _chunk("the forbidden ")
yield _chunk("topic answer", finish_reason="stop")
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
try:
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response)
out = []
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"),
response=_mock_stream(),
request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"},
):
out.append(item)
finally:
unified_module.endpoint_guardrail_translation_mappings = None
assert len(out) == 3
assert isinstance(out[0], ModelResponseStream)
frame = out[-1]
assert isinstance(frame, bytes)
payload = json.loads(frame.decode()[len("data: ") :])
assert payload["error"]["message"] == "Violated guardrail policy"
assert payload["error"]["code"] == "400"
assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail"

235
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -948,19 +948,24 @@ class TestStreamingTransform:
assert streamed == "ABCDEFGHIJ"
@pytest.mark.asyncio
async def test_incremental_diff_underflow_raises(self):
async def test_incremental_diff_underflow_emits_error_frame(self):
"""A transform shorter than what was already streamed cannot retract
bytes: it raises HTTPException(stream_transform_underflow)."""
bytes. Chunks have already been flushed by then, so the underflow
surfaces as the in-stream error frame, not an unraisable HTTPException."""
import json as _json
# First sample emits "ABCDEF" (6 chars); second sample shrinks to 3.
guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1)
chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")]
with pytest.raises(unified_module.HTTPException) as exc_info:
await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["error"] == "stream_transform_underflow"
frame = out[-1]
assert isinstance(frame, bytes)
payload = _json.loads(frame.decode()[len("data: ") :])
assert payload["error"]["message"] == "stream_transform_underflow"
assert payload["error"]["code"] == "400"
@pytest.mark.asyncio
async def test_incremental_diff_final_chunk_preserves_finish_reason(self):
@ -1747,3 +1752,221 @@ class TestAppliedGuardrailsReflectsExecution:
async def test_ordinary_guardrail_is_auto_marked_applied(self):
data = await self._run(_AutoLoggingGuardrail())
assert "auto-logging" in _applied_guardrails(data)
class _EosHttpBlockingGuardrail(CustomGuardrail):
"""Raises the bedrock-shaped block HTTPException at end-of-stream scan time."""
def __init__(self):
super().__init__(guardrail_name="eos-http-block")
self.streaming_end_of_stream_only = True
def should_run_guardrail(self, data, event_type): # type: ignore[override]
return True
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
raise unified_module.HTTPException(
status_code=400,
detail={
"error": "Violated guardrail policy",
"bedrock_guardrail_response": "BLOCKED_TOPIC",
},
)
def _anthropic_sse_event(event_type, data):
import json as _json
return f"event: {event_type}\ndata: {_json.dumps(data)}\n\n".encode()
def _anthropic_message_chunks(texts):
head = [
_anthropic_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_test",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
),
_anthropic_sse_event(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
]
deltas = [
_anthropic_sse_event(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
)
for text in texts
]
tail = [
_anthropic_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}),
_anthropic_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 5},
},
),
_anthropic_sse_event("message_stop", {"type": "message_stop"}),
]
return head + deltas + tail
class TestStreamingHttpErrorFrames:
"""A post-flush end-of-stream guardrail block (HTTPException) must surface as
the endpoint's in-stream error frame instead of an unhandled raise that
silently truncates the SSE stream (PR #38722 defect 1)."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self):
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
yield
unified_module.endpoint_guardrail_translation_mappings = None
@pytest.mark.asyncio
async def test_chat_eos_block_emits_data_error_frame(self):
import json as _json
guardrail = _EosHttpBlockingGuardrail()
chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert out[:2] == chunks
frame = out[-1]
assert isinstance(frame, bytes)
text = frame.decode()
assert text.startswith("data: ")
payload = _json.loads(text[len("data: ") :])
assert payload["error"]["message"] == "Violated guardrail policy"
assert payload["error"]["code"] == "400"
@pytest.mark.asyncio
async def test_messages_eos_block_emits_anthropic_error_event(self):
guardrail = _EosHttpBlockingGuardrail()
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(
UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages"
)
raw = b"".join(c for c in out if isinstance(c, bytes)).decode()
assert "hello " in raw
assert "event: error" in raw
assert "Violated guardrail policy" in raw
assert "guardrail_error" in raw
@pytest.mark.asyncio
async def test_responses_eos_block_emits_error_event_with_next_sequence(self):
guardrail = _EosHttpBlockingGuardrail()
chunks = [
{"type": "response.created", "sequence_number": 0},
{"type": "response.output_text.delta", "sequence_number": 1, "delta": "hello"},
{
"type": "response.completed",
"sequence_number": 2,
"response": {
"model": "gpt-4",
"output": [{"type": "message", "content": [{"type": "output_text", "text": "hello"}]}],
},
},
]
out = await _drive_stream(
UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses"
)
assert chunks[0] in out and chunks[1] in out
assert chunks[2] not in out
error_event = out[-1]
assert error_event.type == "error"
assert error_event.sequence_number == 2
assert error_event.error.message == "Violated guardrail policy"
assert error_event.error.code == "400"
assert error_event.error.type == "guardrail_error"
@pytest.mark.asyncio
async def test_pre_flush_block_still_raises_http_exception(self):
guardrail = _EosHttpBlockingGuardrail()
guardrail.streaming_buffer_until_moderated = True
chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")]
with pytest.raises(unified_module.HTTPException) as exc_info:
await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["error"] == "Violated guardrail policy"
class _AuditRecordingGuardrail(CustomGuardrail):
"""Successful scan that records guardrail_information, like a flags-on audit."""
def __init__(self):
super().__init__(guardrail_name="audit-recorder")
self.streaming_end_of_stream_only = True
def should_run_guardrail(self, data, event_type): # type: ignore[override]
return True
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"action": "NONE"},
request_data=request_data,
guardrail_status="success",
)
return inputs
class TestStreamingGuardrailInformationBucket:
"""guardrail_information written during a chat streaming end-of-stream scan
must land in the request's ``metadata`` bucket that spend logging snapshots.
Regression for PR #38722 defect 2: the chat handler used to plant a
``litellm_metadata`` key first, flipping the bucket so every later
guardrail_information write was diverted and /spend/logs showed null."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self):
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
yield
unified_module.endpoint_guardrail_translation_mappings = None
@pytest.mark.asyncio
async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self):
guardrail = _AuditRecordingGuardrail()
chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")]
async def _mock_stream():
for chunk in chunks:
yield chunk
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key", user_id="user-1", request_route="/v1/chat/completions"
)
request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}}
out = []
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=_mock_stream(),
request_data=request_data,
):
out.append(item)
assert "litellm_metadata" not in request_data
recorded = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(recorded) == 1
assert recorded[0]["guardrail_name"] == "audit-recorder"
assert recorded[0]["guardrail_status"] == "success"
assert request_data["metadata"]["user_api_key_user_id"] == "user-1"

226
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -1228,3 +1228,229 @@ class TestFireDeferredStreamLogging:
assert info is not None, "guardrail_information should be populated"
assert len(info) == 1
assert info[0]["guardrail_name"] == "info-writer"
class TestResponsesIteratorDeferredLogging:
"""Regression for PR #38722 defect 2 on /v1/responses streams: when the
proxy arms _on_deferred_stream_complete, the responses streaming iterator
must store the logging coroutine for ProxyLogging._fire_deferred_stream_logging
(which runs AFTER end-of-stream guardrail scans write guardrail_information)
instead of dispatching immediately with a premature metadata snapshot."""
def _iterator(self, logging_obj):
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
iterator = object.__new__(BaseResponsesAPIStreamingIterator)
iterator.logging_obj = logging_obj
iterator.start_time = None
iterator.completed_response = None
iterator._completed_response_logged = False
iterator._completed_response_cache_hit = None
iterator._persist_completed_response_before_logging = False
return iterator
def _logging_obj(self):
recorded = {}
async def dispatch_success_handlers(result=None, **kwargs):
recorded["dispatched"] = True
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = dispatch_success_handlers
return logging_obj, recorded
@pytest.mark.asyncio
async def test_armed_iterator_stores_deferred_coroutine(self):
logging_obj, recorded = self._logging_obj()
logging_obj._on_deferred_stream_complete = MagicMock()
iterator = self._iterator(logging_obj)
with patch("asyncio.create_task") as mock_create_task:
iterator._log_completed_response(is_async=True)
mock_create_task.assert_not_called()
args = logging_obj._deferred_stream_complete_args
assert isinstance(args, tuple) and len(args) == 1
assert "dispatched" not in recorded
await args[0]
assert recorded["dispatched"] is True
@pytest.mark.asyncio
async def test_unarmed_iterator_dispatches_immediately(self):
logging_obj, recorded = self._logging_obj()
logging_obj._on_deferred_stream_complete = None
iterator = self._iterator(logging_obj)
created = []
real_create_task = asyncio.create_task
def tracking_create_task(coro):
task = real_create_task(coro)
created.append(task)
return task
with patch("asyncio.create_task", side_effect=tracking_create_task):
iterator._log_completed_response(is_async=True)
assert len(created) == 1
await created[0]
assert recorded["dispatched"] is True
class TestArmDeferredStreamDispatch:
"""Regression for PR #38722: the closure shape armed on logging_obj must
match the args the stream's logging owner stores. Bridged /v1/responses
(LiteLLMCompletionStreamingIterator) shares its inner CustomStreamWrapper's
logging_obj, which stores (assembled_response, cache_hit); arming the
single-coroutine native closure there made _fire_deferred_stream_logging
raise TypeError inside the streaming hook, leaking an in-stream 500 error
frame on every streamed /v1/responses request."""
def _processor(self):
return ProxyBaseLLMRequestProcessing(data={"model": "gpt-test"})
def _dispatch_recording_logging_obj(self):
recorded = {}
async def dispatch_success_handlers(
result=None, start_time=None, end_time=None, cache_hit=None, prefer_async_handlers=False
):
recorded["result"] = result
recorded["cache_hit"] = cache_hit
recorded["prefer_async_handlers"] = prefer_async_handlers
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = dispatch_success_handlers
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
return logging_obj, recorded
@pytest.mark.asyncio
async def test_bridged_responses_iterator_gets_csw_arg_shape(self):
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
logging_obj, recorded = self._dispatch_recording_logging_obj()
bridged = object.__new__(LiteLLMCompletionStreamingIterator)
self._processor()._arm_deferred_stream_dispatch(
response=bridged,
route_type="aresponses",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
assert recorded["prefer_async_handlers"] is True
@pytest.mark.asyncio
async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self):
"""The router wraps iterators without _hidden_params in
HiddenParamsAsyncIteratorWrapper before the proxy arms deferral, so
every production streamed /v1/responses reaches arming wrapped;
sniffing the wrapper instead of the inner iterator armed the 1-arg
native closure against the CSW's 2-arg stored shape and leaked a
TypeError 500 frame into the stream."""
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
from litellm.router_utils.add_retry_fallback_headers import (
HiddenParamsAsyncIteratorWrapper,
)
logging_obj, recorded = self._dispatch_recording_logging_obj()
wrapped = HiddenParamsAsyncIteratorWrapper(object.__new__(LiteLLMCompletionStreamingIterator))
self._processor()._arm_deferred_stream_dispatch(
response=wrapped,
route_type="aresponses",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, False)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
assert recorded["prefer_async_handlers"] is True
@pytest.mark.asyncio
async def test_native_stream_closure_enqueues_single_coroutine(self):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, _ = self._dispatch_recording_logging_obj()
async def _agen():
yield b"x"
self._processor()._arm_deferred_stream_dispatch(
response=_agen(),
route_type="anthropic_messages",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
closure = logging_obj._on_deferred_stream_complete
assert closure is not None
async def _logging_coroutine():
return None
coro = _logging_coroutine()
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue:
await closure(coro)
mock_enqueue.assert_called_once_with(async_coroutine=coro)
coro.close()
@pytest.mark.asyncio
async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch):
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
logging_obj, recorded = self._dispatch_recording_logging_obj()
csw = object.__new__(CustomStreamWrapper)
processor = self._processor()
monkeypatch.setattr( # test-quality-ok: empty the process-global callback registry so no ambient guardrail runs
litellm, "callbacks", []
)
processor._arm_deferred_stream_dispatch(
response=csw,
route_type="acompletion",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assembled = object()
await logging_obj._on_deferred_stream_complete(assembled, False)
await asyncio.sleep(0)
assert recorded["result"] is assembled
assert recorded["cache_hit"] is False
assert recorded["prefer_async_handlers"] is True
def test_non_native_route_generator_not_armed(self):
logging_obj, _ = self._dispatch_recording_logging_obj()
async def _agen():
yield b"x"
self._processor()._arm_deferred_stream_dispatch(
response=_agen(),
route_type="acompletion",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assert logging_obj._on_deferred_stream_complete is None

133
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -28,12 +28,21 @@ def _make_streaming_response(chunks):
return mock
def _unarmed_logging_obj():
"""Real Logging objects only carry _on_deferred_stream_complete when the
proxy arms deferred dispatch; a bare MagicMock's auto-attribute is truthy
and would spuriously trigger the deferral branch."""
obj = MagicMock()
obj._on_deferred_stream_complete = None
return obj
@pytest.mark.asyncio
async def test_chunk_processor_logs_on_normal_completion():
chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
response = _make_streaming_response(chunks)
mock_logging_obj = MagicMock()
mock_logging_obj = _unarmed_logging_obj()
mock_passthrough_handler = MagicMock()
with patch.object(
@ -66,7 +75,7 @@ async def test_chunk_processor_logs_on_client_disconnect():
chunks = [b"event-1", b"event-2", b"event-3"]
response = _make_streaming_response(chunks)
mock_logging_obj = MagicMock()
mock_logging_obj = _unarmed_logging_obj()
mock_passthrough_handler = MagicMock()
with patch.object(
@ -104,7 +113,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er
response = _make_streaming_response(chunks)
response.status_code = 403
mock_logging_obj = MagicMock()
mock_logging_obj = _unarmed_logging_obj()
mock_passthrough_handler = MagicMock()
with patch.object(
@ -134,7 +143,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er
async def test_chunk_processor_does_not_schedule_logging_when_no_chunks():
response = _make_streaming_response([])
mock_logging_obj = MagicMock()
mock_logging_obj = _unarmed_logging_obj()
mock_passthrough_handler = MagicMock()
with patch.object(
@ -189,7 +198,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker():
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=MagicMock(),
litellm_logging_obj=_unarmed_logging_obj(),
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
@ -230,7 +239,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne
gen = PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=MagicMock(),
litellm_logging_obj=_unarmed_logging_obj(),
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
@ -246,7 +255,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne
def _logging_obj_with_write_once_cst():
"""Build a MagicMock that mirrors the real Logging behavior: _update_completion_start_time
latches self.completion_start_time so the write-once guard actually latches."""
obj = MagicMock()
obj = _unarmed_logging_obj()
obj.completion_start_time = None
def _update(*, completion_start_time):
@ -301,7 +310,7 @@ async def test_chunk_processor_does_not_reset_completion_start_time_on_later_chu
response = _make_streaming_response(chunks)
real_first = datetime(2020, 1, 1, 0, 0, 0)
mock_logging_obj = MagicMock()
mock_logging_obj = _unarmed_logging_obj()
# Simulate first-chunk stamp having already landed (e.g. under contention or a
# prior wrapper that already set it): later chunks must be no-ops.
mock_logging_obj.completion_start_time = real_first
@ -387,7 +396,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type):
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "gpt-4o-mini", "stream": True},
litellm_logging_obj=MagicMock(),
litellm_logging_obj=_unarmed_logging_obj(),
endpoint_type=endpoint_type,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
@ -517,3 +526,109 @@ def test_convert_raw_bytes_survives_truncated_multibyte_sequence():
lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes)
assert any('"type": "message_delta"' in line for line in lines)
@pytest.mark.asyncio
async def test_chunk_processor_defers_logging_until_fire_when_armed():
"""Regression for PR #38722: native /v1/messages streams route through
chunk_processor, which enqueued the spend log the moment the stream ended,
racing the guardrail end-of-stream scan and logging
guardrail_information as null. With deferred dispatch armed, the completed
stream must park the logging coroutine on logging_obj and only enqueue it
when ProxyLogging._fire_deferred_stream_logging fires after the scan."""
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.utils import ProxyLogging
chunks = [b"event-1", b"event-2"]
response = _make_streaming_response(chunks)
logging_obj = _unarmed_logging_obj()
logging_obj._deferred_stream_complete_args = None
enqueued = []
def _capture(async_coroutine):
enqueued.append(async_coroutine)
async_coroutine.close()
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER,
"ensure_initialized_and_enqueue",
side_effect=_capture,
) as mock_enqueue:
gen = PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=AsyncMock(),
)
ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch(
response=gen,
route_type="anthropic_messages",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
received = []
async for chunk in gen:
received.append(chunk)
await asyncio.sleep(0)
assert received == chunks
mock_enqueue.assert_not_called()
parked = logging_obj._deferred_stream_complete_args
assert isinstance(parked, tuple) and len(parked) == 1
assert asyncio.iscoroutine(parked[0])
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
mock_enqueue.assert_called_once()
@pytest.mark.asyncio
async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_armed():
"""Client disconnects never reach _fire_deferred_stream_logging, so parking
the coroutine there would lose the partial-usage spend log (LIT-2642); the
disconnect path must keep enqueueing immediately."""
chunks = [b"event-1", b"event-2", b"event-3"]
response = _make_streaming_response(chunks)
logging_obj = _unarmed_logging_obj()
async def _armed_closure(logging_coroutine):
raise AssertionError("deferred closure must not fire on disconnect")
logging_obj._on_deferred_stream_complete = _armed_closure
logging_obj._deferred_stream_complete_args = None
enqueued = []
def _capture(async_coroutine):
enqueued.append(async_coroutine)
async_coroutine.close()
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER,
"ensure_initialized_and_enqueue",
side_effect=_capture,
) as mock_enqueue:
gen = PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=AsyncMock(),
)
await gen.__anext__()
await gen.aclose()
mock_enqueue.assert_called_once()
assert logging_obj._deferred_stream_complete_args is None

14
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -1659,29 +1659,29 @@ class TestCommonRequestProcessingHelpers:
async def test_serialize_http_exception_detail_helper(self):
"""Direct unit coverage for the L1 helper across all branches."""
from litellm.proxy.common_request_processing import (
_serialize_http_exception_detail,
serialize_http_exception_detail,
)
import json as _json
assert _serialize_http_exception_detail("plain") == ("plain", None)
assert serialize_http_exception_detail("plain") == ("plain", None)
msg, fields = _serialize_http_exception_detail({"error": "Violated", "extra": "x"})
msg, fields = serialize_http_exception_detail({"error": "Violated", "extra": "x"})
assert msg == "Violated"
assert fields == {"error": "Violated", "extra": "x"}
msg, fields = _serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}})
msg, fields = serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}})
assert msg == "blocked"
assert fields == {"error": {"message": "blocked", "code": "x"}}
msg, fields = _serialize_http_exception_detail({"message": "top-level"})
msg, fields = serialize_http_exception_detail({"message": "top-level"})
assert msg == "top-level"
assert fields == {"message": "top-level"}
msg, fields = _serialize_http_exception_detail({"weird": ["a", "b"]})
msg, fields = serialize_http_exception_detail({"weird": ["a", "b"]})
assert msg == _json.dumps({"weird": ["a", "b"]})
assert fields == {"weird": ["a", "b"]}
assert _serialize_http_exception_detail(42) == ("42", None)
assert serialize_http_exception_detail(42) == ("42", None)
async def test_proxy_exception_from_http_exception_helper(self):
"""The shared HTTPException -> ProxyException conversion keeps a clean

29
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -346,14 +346,14 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions
@pytest.mark.asyncio
async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch):
async def test_unified_guardrail_iterator_accepts_explicit_guardrail():
"""
The dispatch passes each guardrail explicitly instead of through a shared
request_data key, so chaining two unified-routed guardrails cannot drop
all but the last one.
all but the last one. The block fires after the deltas were already
flushed to the client, so it surfaces as a trailing in-stream error frame
rather than a raised HTTPException.
"""
from fastapi import HTTPException
from litellm.proxy.utils import unified_guardrail
guardrail = _content_filter_guardrail("BLOCK")
@ -367,14 +367,19 @@ async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
yield chunk
with pytest.raises(HTTPException):
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
response=fake_stream(),
request_data=request_data,
guardrail_to_apply=guardrail,
):
pass
delivered = []
async for item in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
response=fake_stream(),
request_data=request_data,
guardrail_to_apply=guardrail,
):
delivered.append(item)
raw = b"".join(c for c in delivered if isinstance(c, bytes)).decode()
assert "event: error" in raw
assert "guardrail_error" in raw
assert raw.index("guardrail_error") > raw.index(" zebra runs")
@pytest.mark.asyncio