diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..60b4444e1d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -62,6 +62,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran return translation +def resolve_endpoint_translation( + user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None +) -> "tuple[str, BaseTranslation] | None": + """ + Resolve the endpoint guardrail translation for a streamed response: the + request route wins, falling back to inferring the call type from the first + response chunk (the same resolution order the streaming iterator hook uses). + Returns None when the call type is unresolvable or has no translation. + """ + route_call_types: Final = ( + get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None + ) + call_type: Final = ( + route_call_types[0].value + if route_call_types + else ( + _infer_call_type(call_type=None, completion_response=first_response_item) + if first_response_item is not None + else None + ) + ) + if call_type is None: + return None + try: + handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type)) + except ValueError: + return None + return call_type, handler_cls() + + def _chunk_choices(item: object) -> Sequence[object]: choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] return choices @@ -346,7 +376,7 @@ class UnifiedLLMGuardrails(CustomLogger): return response - async def _handle_streaming_block( + async def handle_streaming_block( self, exc: "ModifyResponseException", endpoint_translation: _EndpointTranslation, @@ -402,7 +432,7 @@ class UnifiedLLMGuardrails(CustomLogger): return None return call_type - async def _emit_streaming_http_error( + async def emit_streaming_http_error( self, exc: HTTPException, call_type: str | None, @@ -577,7 +607,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -586,7 +616,7 @@ 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): yield error_item raise _StreamTerminated() @@ -758,7 +788,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -1060,7 +1090,7 @@ class UnifiedLLMGuardrails(CustomLogger): # The current chunk was appended to responses_so_far but not # yet yielded, so exclude it: the continuation must reflect # only what the client has actually received. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=chunks_yielded, @@ -1124,7 +1154,7 @@ class UnifiedLLMGuardrails(CustomLogger): # terminating SSE sequence with the block message rather than # propagating into a bare error blob that truncates the stream. # The withheld original chunks are never released. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 190784a2a60..c422a7c0964 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal import litellm from litellm._logging import verbose_proxy_logger @@ -25,6 +25,11 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStepResult, ) +if TYPE_CHECKING: + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + try: from fastapi.exceptions import HTTPException except ImportError: @@ -43,6 +48,8 @@ class PipelineExecutor: call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ Execute pipeline steps sequentially with conditional actions. @@ -59,6 +66,12 @@ class PipelineExecutor: step whose guardrail opted into ``scan_raw_request`` evaluates the original request instead of whatever an earlier ``pass_data`` step in this same pipeline already rewrote. + streaming_chunks: buffered chunks of a completed stream. When set + (with ``endpoint_translation``), post_call steps scan the + assembled streamed output through the endpoint translation + instead of calling ``async_post_call_success_hook``. + endpoint_translation: the guardrail translation for the streamed + endpoint, resolved by the caller. Returns: PipelineExecutionResult with terminal action and step results @@ -83,6 +96,8 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, call_type=call_type, raw_request_snapshot=raw_request_snapshot, + streaming_chunks=streaming_chunks, + endpoint_translation=endpoint_translation, ) duration = time.perf_counter() - start_time @@ -154,6 +169,8 @@ class PipelineExecutor: user_api_key_dict: Any, call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], dict | None, @@ -198,10 +215,8 @@ class PipelineExecutor: # Use unified_guardrail path if callback implements apply_guardrail target: CustomLogger = callback - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: + use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + if use_unified and streaming_chunks is None: hook_input["guardrail_to_apply"] = callback target = UnifiedLLMGuardrails() @@ -216,6 +231,22 @@ class PipelineExecutor: callback.mark_pre_call_hook_ran(data) if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) + elif mode == "post_call" and streaming_chunks is not None: + if not use_unified or endpoint_translation is None: + return ( + "error", + None, + f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + None, + ) + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=callback, + litellm_logging_obj=data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + response = None elif mode == "post_call": response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, @@ -246,6 +277,12 @@ class PipelineExecutor: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + @staticmethod + def supports_unified_execution(callback: CustomGuardrail) -> bool: + """Whether this guardrail runs through the unified apply_guardrail path, + the interface streaming pipeline execution requires.""" + return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 80e991640a9..5642eb5383e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,7 +19,7 @@ from email.mime.text import MIMEText from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -486,29 +486,80 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) +def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + return callback is not None and PipelineExecutor.supports_unified_execution(callback) + + +class _PipelineErrorBody(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + policies: ReadOnly[tuple[str, ...]] + guardrails: NotRequired[ReadOnly[tuple[str, ...]]] + + +class _PipelineErrorDetail(TypedDict): + error: ReadOnly[_PipelineErrorBody] + + def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("stream") is not True and data.get("background") is not True: + """ + Reject up front the requests whose post_call pipelines could never run. + + Background responses skip the post_call hooks entirely, so a pipeline + governing one would silently never execute. Streaming responses execute + pipelines against the buffered stream through the endpoint guardrail + translations, which requires every step's guardrail to support the unified + apply_guardrail interface; steps that cannot (native-lifecycle guardrails, + or guardrails not registered at all) keep the 400 rather than letting + ungoverned output stream through. + """ + is_stream: Final = data.get("stream") is True + is_background: Final = data.get("background") is True + if not is_stream and not is_background: return - post_call_policies: Final = tuple( - policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + post_call_pipelines: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" ) - if not post_call_policies: + if not post_call_pipelines: return - raise HTTPException( - status_code=400, - detail={ + post_call_policies: Final = tuple(policy_name for policy_name, _pipeline in post_call_pipelines) + if is_background: + background_detail: Final[_PipelineErrorDetail] = { "error": { "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming or background " - f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and " - "background=false, or move these policies' output guardrails from pipeline steps to " - "guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern background " + f"responses: {', '.join(post_call_policies)}. Retry with background=false." ), "type": "guardrail_pipeline_error", - "policies": list(post_call_policies), + "policies": post_call_policies, } - }, + } + raise HTTPException(status_code=400, detail=background_detail) + unsupported_guardrails: Final = tuple( + dict.fromkeys( + step.guardrail + for _policy_name, pipeline in post_call_pipelines + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail) + ) ) + if not unsupported_guardrails: + return + unsupported_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses " + "because these pipeline guardrails do not support the unified apply_guardrail " + f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or move " + "them from pipeline steps to guardrails.add, which scans streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": post_call_policies, + "guardrails": unsupported_guardrails, + } + } + raise HTTPException(status_code=400, detail=unsupported_detail) def _prompt_block_text(block: object) -> str: @@ -1689,7 +1740,7 @@ class ProxyLogging: result: PipelineExecutionResult, data: dict, policy_name: str, - original_response: LLMResponseTypes | None = None, + original_response: "LLMResponseTypes | Sequence[object] | None" = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. @@ -1699,7 +1750,9 @@ class ProxyLogging: payload (already sent upstream) must stay untouched; a replacement response carried in ``modified_data`` is adopted by the caller, and metadata-bucket writes (applied guardrails, guardrail logging info) - are merged back so headers and spend logs still see them. + are merged back so headers and spend logs still see them. On the + streaming path it is the buffered chunk list, carried into + ``ModifyResponseException.original_response`` for usage reporting. """ if result.terminal_action == "allow": if result.modified_data is not None: @@ -3195,11 +3248,16 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False + pipeline_managed: Final = ( + _pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): + if callback.guardrail_name in pipeline_managed: + continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks @@ -3256,12 +3314,17 @@ class ProxyLogging: 1. /chat/completions """ caps: Final = ProxyLogging._callback_capabilities() + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in _policy_pipelines(request_data) + if pipeline.mode == "post_call" + ) # Fast path: no real overrides. Internal proxy CustomLogger callbacks # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. - if not caps.iterator_overrides: + if not caps.iterator_overrides and not post_call_pipelines: try: async for chunk in response: yield chunk @@ -3281,8 +3344,11 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) + pipeline_managed_names: Final = _pipeline_managed_guardrail_names(request_data, "post_call") for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): + if resolved_callback.guardrail_name in pipeline_managed_names: + continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True @@ -3322,6 +3388,17 @@ class ProxyLogging: ), ) + # Policy pipelines run last, over the fully buffered stream, so a + # pipeline verdict covers whatever the flat guardrail chain above + # already let through. + if post_call_pipelines: + current_response = self._pipeline_gated_stream( + response=current_response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + pipelines=post_call_pipelines, + ) + try: async for chunk in current_response: yield chunk @@ -3337,6 +3414,82 @@ class ProxyLogging: # we reach this point the metadata is fully populated. ProxyLogging._fire_deferred_stream_logging(request_data) + async def _pipeline_gated_stream( + self, + response: "AsyncGenerator[object, None]", + user_api_key_dict: UserAPIKeyAuth, + request_data: dict, + pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + ) -> "AsyncGenerator[Any, None]": + """ + Execute post_call policy pipelines against a streamed response. + + Buffers the whole stream (nothing reaches the client until every + pipeline allows it), then runs each pipeline's steps against the + assembled output through the endpoint guardrail translation, the same + machinery flat post_call guardrails use at end of stream. An allow + releases the buffered chunks verbatim; a block or modify_response + terminates with the translation's block chunks or the raised error. + """ + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + resolve_endpoint_translation, + ) + + buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict + async for item in response: + buffered.append(item) + if not buffered: + return + + resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) + if resolved is None: + policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines) + unresolvable_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policy pipelines could not govern this streaming response shape; " + f"the response was withheld: {', '.join(policy_names)}." + ), + "type": "guardrail_pipeline_error", + "policies": policy_names, + } + } + raise HTTPException(status_code=500, detail=unresolvable_detail) + call_type, endpoint_translation = resolved + + for policy_name, pipeline in pipelines: + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode="post_call", + data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + policy_name=policy_name, + streaming_chunks=buffered, + endpoint_translation=endpoint_translation, + ) + try: + ProxyLogging._handle_pipeline_result( + result, data=request_data, policy_name=policy_name, original_response=buffered + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = buffered + async for block_chunk in unified_guardrail.handle_streaming_block( + e, endpoint_translation, stream_started=False, responses_so_far=() + ): + yield block_chunk + return + except HTTPException as e: + async for error_chunk in unified_guardrail.emit_streaming_http_error( + e, call_type, buffered, request_data + ): + yield error_chunk + return + + for buffered_item in buffered: + yield buffered_item + @staticmethod def _fire_deferred_stream_logging(request_data: dict) -> None: """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..fe4acbf8277 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1026,7 +1026,7 @@ class TestStreamingTransform: ) emitted = [] - async for item in handler._emit_streaming_http_error( + async for item in handler.emit_streaming_http_error( exc, call_type=CallTypes.asend_message.value, responses_so_far=[{"id": "req-1"}], diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 34a25d75bc0..2b36ae111dc 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1284,7 +1284,8 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + assert info.value.detail["error"]["policies"] == ("response-governance",) + assert info.value.detail["error"]["guardrails"] == ("gr-post",) assert "stream=false" in info.value.detail["error"]["message"] @@ -1304,7 +1305,7 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + assert info.value.detail["error"]["policies"] == ("response-governance",) assert "background=false" in info.value.detail["error"]["message"] @@ -1333,3 +1334,166 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c ) assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None assert _raise_for_streaming_post_call_pipelines({"background": True}) is None + + +# --------------------------------------------------------------------------- +# post_call pipelines on streaming responses +# --------------------------------------------------------------------------- + + +def _unified_stream_guardrail(seen: Dict[str, Any], block: bool = False) -> CustomGuardrail: + class UnifiedStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + seen["count"] = seen.get("count", 0) + 1 + seen["input_type"] = input_type + if block: + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + return inputs + + return UnifiedStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _stream_chunks() -> List[Any]: + return [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]), + ] + + +async def _async_chunk_iter(chunks: List[Any]): + for chunk in chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert out is not None + assert out.get("stream") is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_lifecycle", [False, True]) +async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_unified_support( + proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle +): + if native_lifecycle: + + class NativeOnlyGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + else: + + class NativeOnlyGuardrail(CustomGuardrail): + pass + + monkeypatch.setattr( + litellm, + "callbacks", + [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "apply_guardrail" in info.value.detail["error"]["message"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert seen["count"] == 1 + assert seen["input_type"] == "response" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 400 + assert "output blocked" in str(info.value.detail) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter([object(), object()]), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 500 + assert "withheld" in info.value.detail["error"]["message"] + assert seen.get("count") is None