Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py: 10%
515 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
41. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_pre_call_hook
52. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_success_hook
63. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook
7"""
9import copy
10import json
11from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence
12from typing import TYPE_CHECKING, Any, Final, Protocol
14from fastapi import HTTPException
16from litellm._logging import verbose_proxy_logger
17from litellm.caching.caching import DualCache
18from litellm.cost_calculator import _infer_call_type
19from litellm.integrations.custom_guardrail import CustomGuardrail
20from litellm.integrations.custom_logger import CustomLogger
21from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
22from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings
23from litellm.proxy._types import UserAPIKeyAuth
24from litellm.types.guardrails import GuardrailEventHooks
25from litellm.types.utils import (
26 CallTypes,
27 CallTypesLiteral,
28 Delta,
29 ModelResponseStream,
30 StreamingChoices,
31)
33if TYPE_CHECKING: 33 ↛ 36line 33 didn't jump to line 36 because the condition on line 33 was never true
34 # Imported lazily at runtime (inside the streaming hook) to avoid a
35 # module-level cyclic import with litellm.integrations.custom_guardrail.
36 from litellm.integrations.custom_guardrail import ModifyResponseException
37 from litellm.llms.base_llm.guardrail_translation.base_translation import (
38 BaseTranslation,
39 StreamingScanKey,
40 )
42# Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error
43A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message)
45GUARDRAIL_NAME: Final = "unified_llm_guardrails"
48class _EndpointTranslation(Protocol):
49 @property
50 def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ... 50 ↛ exitline 50 didn't return from function 'process_input_messages' because
52 @property
53 def process_output_response(self) -> "Callable[..., Awaitable[object]]": ... 53 ↛ exitline 53 didn't return from function 'process_output_response' because
55 @property
56 def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... 56 ↛ exitline 56 didn't return from function 'process_output_streaming_response' because
58 @property
59 def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... 59 ↛ exitline 59 didn't return from function 'get_streaming_scan_key' because
61 @property
62 def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... 62 ↛ exitline 62 didn't return from function 'build_block_sse_chunks' because
64 @property
65 def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ... 65 ↛ exitline 65 didn't return from function 'build_stream_error_items' because
68def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation:
69 return translation
72def resolve_endpoint_translation(
73 user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None
74) -> "tuple[str, BaseTranslation] | None":
75 """
76 Resolve the endpoint guardrail translation for a streamed response: the
77 request route wins, falling back to inferring the call type from the first
78 response chunk (the same resolution order the streaming iterator hook uses).
79 Returns None when the call type is unresolvable or has no translation.
80 """
81 route_call_types: Final = (
82 get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None
83 )
84 call_type: Final = (
85 route_call_types[0].value
86 if route_call_types
87 else (
88 _infer_call_type(call_type=None, completion_response=first_response_item)
89 if first_response_item is not None
90 else None
91 )
92 )
93 if call_type is None:
94 return None
95 try:
96 handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type))
97 except ValueError:
98 return None
99 return call_type, handler_cls()
102def _chunk_choices(item: object) -> Sequence[object]:
103 choices: Final[Sequence[object]] = getattr(item, "choices", None) or []
104 return choices
107def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
108 return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
111def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
112 if scan_key is None:
113 return False
114 return scan_key == last_scan_key or scan_key.has_nothing_to_scan
117class _StreamTerminated(Exception):
118 """Internal signal that the incremental transform stream has already emitted
119 its terminal chunks (block message or in-stream error) and must stop."""
122def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None:
123 """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting."""
124 for item in responses_so_far:
125 if isinstance(item, dict) and "id" in item:
126 return item.get("id")
127 if isinstance(item, str):
128 try:
129 obj = json.loads(item.strip())
130 if isinstance(obj, dict) and "id" in obj:
131 return obj.get("id")
132 except (json.JSONDecodeError, TypeError):
133 continue
134 body: Final = request_data.get("body") or request_data.get("data") or {}
135 if isinstance(body, dict):
136 return body.get("id")
137 return None
140def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapping[str, object]:
141 """Build the in-stream JSON-RPC error object for a mid-stream A2A failure.
143 Returned as an object, not a serialized string: the A2A endpoint owns wire
144 framing and serializes whatever the stream yields.
145 """
146 detail: Final = exc.detail if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
147 return {
148 "jsonrpc": "2.0",
149 "id": request_id,
150 "error": {
151 "code": -32603,
152 "message": detail.get("error", detail.get("message", str(exc.detail))),
153 "data": {k: v for k, v in detail.items() if k not in ("error", "message")},
154 },
155 }
158def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None:
159 """Populate data['litellm_metadata'] from user_api_key_dict if absent."""
160 if "litellm_metadata" not in data:
161 from litellm.llms.base_llm.guardrail_translation.base_translation import (
162 BaseTranslation,
163 )
165 user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
166 if user_metadata:
167 data["litellm_metadata"] = user_metadata
170class UnifiedLLMGuardrails(CustomLogger):
171 def __init__(
172 self,
173 **kwargs,
174 ):
175 # store kwargs as optional_params
176 self.optional_params = kwargs
178 super().__init__(**kwargs)
180 verbose_proxy_logger.debug(
181 "UnifiedLLMGuardrails initialized with optional_params: %s",
182 self.optional_params,
183 )
185 async def async_pre_call_hook(
186 self,
187 user_api_key_dict: UserAPIKeyAuth,
188 cache: DualCache,
189 data: dict,
190 call_type: CallTypesLiteral,
191 ) -> Exception | str | dict | None:
192 """
193 Runs before the LLM API call
194 Runs on only Input
195 Use this if you want to MODIFY the input
196 """
198 from litellm.proxy.common_utils.callback_utils import (
199 add_guardrail_to_applied_guardrails_header,
200 )
202 verbose_proxy_logger.debug("Running UnifiedLLMGuardrails pre-call hook")
204 guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
205 if guardrail_to_apply is None:
206 return data
208 event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call
209 if call_type == CallTypes.call_mcp_tool.value:
210 event_type = GuardrailEventHooks.pre_mcp_call
212 if guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True:
213 verbose_proxy_logger.debug(
214 "UnifiedLLMGuardrails: Pre-call scanning disabled for %s",
215 guardrail_to_apply.guardrail_name,
216 )
217 return data
219 mappings: Final = load_guardrail_translation_mappings()
221 try:
222 if CallTypes(call_type) not in mappings:
223 return data
224 except ValueError:
225 return data # handle unmapped call types
227 endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
229 _ensure_litellm_metadata(data, user_api_key_dict)
231 data = await endpoint_translation.process_input_messages(
232 data=data,
233 guardrail_to_apply=guardrail_to_apply,
234 litellm_logging_obj=data.get("litellm_logging_obj"),
235 )
237 if not guardrail_to_apply.records_own_guardrail_information:
238 add_guardrail_to_applied_guardrails_header(
239 request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
240 )
241 return data
243 async def async_moderation_hook(
244 self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral
245 ) -> object:
246 """
247 Runs in parallel to LLM API call
248 Runs on only Input
250 This can NOT modify the input, only used to reject or accept a call before going to LLM API
251 """
252 verbose_proxy_logger.debug("Running UnifiedLLMGuardrails moderation hook")
254 guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
255 if guardrail_to_apply is None:
256 return data
258 event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
259 if call_type == CallTypes.call_mcp_tool.value:
260 event_type = GuardrailEventHooks.during_mcp_call
262 if guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True:
263 verbose_proxy_logger.debug(
264 "UnifiedLLMGuardrails: Pre-call scanning disabled for %s",
265 guardrail_to_apply.guardrail_name,
266 )
267 return data
269 mappings: Final = load_guardrail_translation_mappings()
270 if call_type is not None and CallTypes(call_type) not in mappings:
271 return data
273 endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
275 _ensure_litellm_metadata(data, user_api_key_dict)
277 return await endpoint_translation.process_input_messages(
278 data=data,
279 guardrail_to_apply=guardrail_to_apply,
280 litellm_logging_obj=data.get("litellm_logging_obj"),
281 )
283 async def async_post_call_success_hook(
284 self,
285 data: dict,
286 user_api_key_dict: UserAPIKeyAuth,
287 response,
288 ) -> object:
289 """
290 Runs on response from LLM API call
292 It can be used to reject a response
294 Uses Enkrypt AI guardrails to check the response for policy violations, PII, and injection attacks
295 """
296 # Local import avoids a module-level cyclic import with
297 # litellm.integrations.custom_guardrail.
298 from litellm.integrations.custom_guardrail import ModifyResponseException
299 from litellm.proxy.common_utils.callback_utils import (
300 add_guardrail_to_applied_guardrails_header,
301 )
302 from litellm.types.guardrails import GuardrailEventHooks
304 guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
306 if guardrail_to_apply is None:
307 return
309 if guardrail_to_apply.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True:
310 return
312 verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response)
314 call_type: CallTypesLiteral | None = None
315 if user_api_key_dict.request_route is not None:
316 call_types: Final = get_call_types_for_route(user_api_key_dict.request_route)
317 if call_types is not None and len(call_types) > 0:
318 call_type = call_types[0]
319 if call_type is None:
320 call_type = _infer_call_type(call_type=None, completion_response=response)
322 if call_type is None:
323 litellm_logging_obj: Final = data.get("litellm_logging_obj")
324 logging_call_type: Final = (
325 getattr(litellm_logging_obj, "call_type", None) if litellm_logging_obj is not None else None
326 )
327 if logging_call_type in (
328 CallTypes.pass_through.value,
329 CallTypes.allm_passthrough_route.value,
330 CallTypes.ocr.value,
331 CallTypes.aocr.value,
332 ):
333 call_type = logging_call_type
335 if call_type is None:
336 verbose_proxy_logger.warning(
337 "Guardrail '%s' selected for route '%s' but its call type could not be resolved; "
338 "skipping post-call scanning. Add the route to API_ROUTE_TO_CALL_TYPES.",
339 guardrail_to_apply.guardrail_name,
340 user_api_key_dict.request_route,
341 )
342 return response
344 mappings: Final = load_guardrail_translation_mappings()
346 if CallTypes(call_type) not in mappings:
347 verbose_proxy_logger.warning(
348 "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; "
349 "skipping post-call scanning.",
350 guardrail_to_apply.guardrail_name,
351 user_api_key_dict.request_route,
352 call_type,
353 )
354 return response
356 endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
358 try:
359 response = await endpoint_translation.process_output_response(
360 response=response,
361 guardrail_to_apply=guardrail_to_apply,
362 litellm_logging_obj=data.get("litellm_logging_obj"),
363 user_api_key_dict=user_api_key_dict,
364 request_data=data,
365 )
366 except ModifyResponseException as e:
367 # The guardrail blocked the response. Attach the original LLM
368 # response so the endpoint handler can report its real token usage
369 # instead of discarding it (the block replaces the content, but the
370 # upstream call already consumed those tokens).
371 if e.original_response is None:
372 e.original_response = response
373 raise
374 if not guardrail_to_apply.records_own_guardrail_information:
375 add_guardrail_to_applied_guardrails_header(
376 request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
377 )
379 return response
381 async def handle_streaming_block(
382 self,
383 exc: "ModifyResponseException",
384 endpoint_translation: _EndpointTranslation,
385 stream_started: bool,
386 responses_so_far: Sequence[object],
387 ) -> AsyncGenerator[object, None]:
388 """
389 Terminate a streamed response cleanly when a guardrail blocks it.
391 Format-agnostic routing: delegates to the provider translation handler's
392 ``build_block_sse_chunks`` (see ``BaseTranslation.build_block_sse_chunks``
393 for the ``stream_started`` / ``responses_so_far`` contract). When the
394 format has no safe terminator the handler returns None and we re-raise
395 ``exc`` so the proxy can surface a clean error.
396 """
397 block_chunks: Final = endpoint_translation.build_block_sse_chunks(
398 exc, stream_started=stream_started, responses_so_far=responses_so_far
399 )
400 if block_chunks is None:
401 raise exc
402 for chunk in block_chunks:
403 yield chunk
405 @staticmethod
406 def _resolve_transform_call_type(
407 user_api_key_dict: UserAPIKeyAuth,
408 mappings: Mapping[CallTypes, type["BaseTranslation"]],
409 ) -> str | None:
410 """Resolve the call type for the incremental_diff path, or None if the
411 route is unresolvable / unsupported.
413 Incremental transformation needs a route we can resolve before the first
414 chunk and a handler that supports the streaming text-diff protocol (v1:
415 the OpenAI chat completions handler only). Returning None makes the caller
416 fall back to block_only.
417 """
418 from litellm.llms.openai.chat.guardrail_translation.handler import (
419 OpenAIChatCompletionsHandler,
420 )
422 if user_api_key_dict.request_route is None:
423 return None
424 call_types: Final = get_call_types_for_route(user_api_key_dict.request_route)
425 if not call_types:
426 return None
427 call_type: Final = call_types[0].value
428 try:
429 mapped: Final = CallTypes(call_type)
430 except ValueError:
431 return None
432 handler_cls: Final = mappings.get(mapped)
433 if handler_cls is None or not issubclass(handler_cls, OpenAIChatCompletionsHandler):
434 return None
435 return call_type
437 async def emit_streaming_http_error(
438 self,
439 exc: HTTPException,
440 call_type: str | None,
441 responses_so_far: Sequence[object],
442 request_data: dict,
443 endpoint_translation: _EndpointTranslation | None = None,
444 stream_started: bool = False,
445 responses_yielded: Sequence[object] | None = None,
446 ) -> AsyncGenerator[object, None]:
447 """Surface a mid-stream HTTPException (a guardrail block with the default
448 exception-on-block config, or a failed scan).
450 A2A call types emit an in-stream JSON-RPC error chunk. For other call
451 types, once chunks have already reached the client the HTTP status is
452 gone, so the failure is delegated to the endpoint translation's
453 ``build_stream_error_items`` and travels as an in-stream error frame in
454 that endpoint's wire format. Before the first chunk (or when the format
455 has no in-stream error frame) the exception is re-raised so the proxy
456 can report it with a real HTTP status.
457 """
458 if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
459 yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data))
460 return
461 if stream_started and endpoint_translation is not None:
462 error_items: Final = endpoint_translation.build_stream_error_items(
463 exc, responses_so_far=tuple(responses_yielded) if responses_yielded is not None else None
464 )
465 if error_items is not None:
466 for error_item in error_items:
467 yield error_item
468 return
469 raise exc
471 def _build_transform_chunk(
472 self,
473 *,
474 reference_chunk: object,
475 mutated_text_per_choice: dict[int, str],
476 emitted_text_per_choice: dict[int, str],
477 holdback_per_choice: dict[int, int],
478 finish_reason_per_choice: dict[int, str | None],
479 held_chars_per_choice: dict[int, int],
480 is_final: bool,
481 ) -> ModelResponseStream | None:
482 """Build the synthetic chunk carrying the newly-guardrailed deltas.
484 For each choice, the new delta is the mutated accumulated text past what
485 has already been emitted, minus a trailing holdback (forced to 0 on the
486 final flush). ``emitted_text_per_choice`` holds the exact bytes already
487 sent per choice and is extended in place; ``held_chars_per_choice`` is
488 updated in place with how many mutated chars per choice are still withheld
489 after this round. Returns None when there is no
490 text to emit (e.g. a tool-call-only turn) or nothing new and this is not
491 the final chunk.
493 Raises HTTPException(400, stream_transform_underflow) when the guardrail's
494 transform is not a forward extension of what has already been streamed
495 (shorter than, or rewrites, the already-sent prefix), since emitted bytes
496 cannot be retracted. This makes the framework fail closed rather than
497 silently leave un-transformed text on the wire; a guardrail that needs to
498 rewrite recent output must withhold it first via ``stream_holdback_chars``.
499 """
500 if not mutated_text_per_choice:
501 # Fix #4 — on the final flush a deferred finish_reason (from a mixed
502 # content+tool_calls chunk whose passthrough suppressed it) still
503 # needs to reach the client, even if the guardrail returned no text
504 # to emit. Build a terminator chunk carrying finish_reason per choice.
505 if is_final and finish_reason_per_choice:
506 terminator_choices: Final[list[StreamingChoices]] = []
507 for choice_idx, finish_reason in finish_reason_per_choice.items():
508 if finish_reason is None:
509 continue
510 terminator_choices.append(
511 StreamingChoices(
512 index=choice_idx,
513 delta=Delta(content="", role=None, tool_calls=None),
514 finish_reason=finish_reason,
515 )
516 )
517 if terminator_choices:
518 return ModelResponseStream(
519 id=getattr(reference_chunk, "id", None),
520 created=getattr(reference_chunk, "created", None),
521 model=getattr(reference_chunk, "model", None),
522 choices=terminator_choices,
523 )
524 return None
526 deltas: Final[dict[int, str]] = {}
527 for choice_idx, text in mutated_text_per_choice.items():
528 already = emitted_text_per_choice.get(choice_idx, "")
529 if not text.startswith(already):
530 raise HTTPException(
531 status_code=400,
532 detail={
533 "error": "stream_transform_underflow",
534 "message": (
535 f"Guardrail streaming transform for choice {choice_idx} is not a forward "
536 f"extension of the {len(already)} chars already streamed to the client "
537 "(it is shorter than, or rewrites, the emitted prefix); emitted bytes "
538 "cannot be retracted. Withhold recent output via stream_holdback_chars "
539 "before rewriting it."
540 ),
541 },
542 )
543 holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
544 end = max(len(already), len(text) - holdback)
545 deltas[choice_idx] = text[len(already) : end]
546 held_chars_per_choice[choice_idx] = len(text) - end
548 # Iterate the mutated choices (not just those in reference_chunk) so a
549 # choice with pending text is never dropped for n > 1. finish_reason is
550 # taken per choice from the accumulated map (a choice can finish in an
551 # earlier chunk than the stream's last one); tool_calls are dropped since
552 # v1 does not transform streamed tool calls (they pass through raw).
553 synthetic_choices: Final[list[StreamingChoices]] = []
554 for choice_idx in mutated_text_per_choice:
555 delta_text = deltas.get(choice_idx, "")
556 finish_reason = finish_reason_per_choice.get(choice_idx) if is_final else None
557 # Skip a choice with nothing to say: no new content and no
558 # finish_reason to deliver. This avoids emitting an empty delta for an
559 # already-finished choice (e.g. one that terminated via a passed-through
560 # tool-call chunk, which already carried its own finish_reason).
561 if not delta_text and finish_reason is None:
562 continue
563 # role="assistant" on this choice's first emitted delta only.
564 role = "assistant" if not emitted_text_per_choice.get(choice_idx) else None
565 synthetic_choices.append(
566 StreamingChoices(
567 index=choice_idx,
568 delta=Delta(content=delta_text, role=role, tool_calls=None),
569 finish_reason=finish_reason,
570 )
571 )
573 if not synthetic_choices:
574 return None
576 for choice_idx in mutated_text_per_choice:
577 emitted_text_per_choice[choice_idx] = emitted_text_per_choice.get(choice_idx, "") + deltas.get(
578 choice_idx, ""
579 )
581 return ModelResponseStream(
582 id=getattr(reference_chunk, "id", None),
583 created=getattr(reference_chunk, "created", None),
584 model=getattr(reference_chunk, "model", None),
585 choices=synthetic_choices,
586 )
588 async def _emit_transform_round(
589 self,
590 *,
591 endpoint_translation: _EndpointTranslation,
592 guardrail_to_apply: CustomGuardrail,
593 request_data: dict,
594 user_api_key_dict: UserAPIKeyAuth,
595 call_type: str,
596 reference_chunk: object,
597 responses_so_far: Sequence[object],
598 responses_yielded: list[object],
599 emitted_text_per_choice: dict[int, str],
600 finish_reason_per_choice: dict[int, str | None],
601 held_chars_per_choice: dict[int, int],
602 is_final: bool,
603 ) -> AsyncGenerator[object, None]:
604 """Run one guardrail processing round and emit the resulting diff chunk.
606 Raises ``_StreamTerminated`` (after emitting the terminal block message or
607 in-stream error) when the guardrail blocks or an underflow occurs.
608 """
609 from litellm.integrations.custom_guardrail import ModifyResponseException
610 from litellm.llms.base_llm.guardrail_translation.base_translation import (
611 StreamTransformSink,
612 )
614 sink: Final = StreamTransformSink()
615 try:
616 await endpoint_translation.process_output_streaming_response(
617 responses_so_far=responses_so_far,
618 guardrail_to_apply=guardrail_to_apply,
619 litellm_logging_obj=request_data.get("litellm_logging_obj"),
620 user_api_key_dict=user_api_key_dict,
621 request_data=request_data,
622 stream_transform_sink=sink,
623 )
624 synthetic: Final = self._build_transform_chunk(
625 reference_chunk=reference_chunk,
626 mutated_text_per_choice=sink.mutated_text_per_choice,
627 emitted_text_per_choice=emitted_text_per_choice,
628 holdback_per_choice=sink.holdback_per_choice,
629 finish_reason_per_choice=finish_reason_per_choice,
630 held_chars_per_choice=held_chars_per_choice,
631 is_final=is_final,
632 )
633 except ModifyResponseException as e:
634 if e.original_response is None:
635 e.original_response = responses_so_far
636 async for block_chunk in self.handle_streaming_block(
637 e,
638 endpoint_translation,
639 stream_started=bool(responses_yielded),
640 responses_so_far=responses_yielded,
641 ):
642 yield block_chunk
643 raise _StreamTerminated()
644 except HTTPException as e:
645 async for error_item in self.emit_streaming_http_error(
646 e,
647 call_type,
648 responses_so_far,
649 request_data,
650 endpoint_translation=endpoint_translation,
651 stream_started=bool(responses_yielded),
652 responses_yielded=responses_yielded,
653 ):
654 yield error_item
655 raise _StreamTerminated()
657 if synthetic is not None:
658 responses_yielded.append(synthetic)
659 yield synthetic
661 async def _run_incremental_transform_stream(
662 self,
663 *,
664 guardrail_to_apply: CustomGuardrail,
665 response: AsyncIterable[object],
666 request_data: dict,
667 user_api_key_dict: UserAPIKeyAuth,
668 call_type: str,
669 sampling_rate: int,
670 end_of_stream_only: bool,
671 mappings: Mapping[CallTypes, type["BaseTranslation"]],
672 ) -> AsyncGenerator[object, None]:
673 """Emit guardrail text transformations as new deltas on the stream.
675 Raw chunks are withheld and accumulated; on each sampled processing round
676 (and once at end of stream) the guardrailed accumulated text is diffed
677 against what has already been emitted and the new portion is sent as a
678 synthetic chunk. A BLOCK terminates the stream via the shared block
679 handler; an underflow surfaces as an HTTPException.
680 """
681 endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
682 responses_so_far: Final[list[object]] = []
683 responses_yielded: Final[list[object]] = []
684 emitted_text_per_choice: Final[dict[int, str]] = {}
685 finish_reason_per_choice: Final[dict[int, str | None]] = {}
686 held_chars_per_choice: Final[dict[int, int]] = {}
687 chunk_counter = 0
688 last_chunk: object | None = None
690 def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]:
691 return self._emit_transform_round(
692 endpoint_translation=endpoint_translation,
693 guardrail_to_apply=guardrail_to_apply,
694 request_data=request_data,
695 user_api_key_dict=user_api_key_dict,
696 call_type=call_type,
697 reference_chunk=reference_chunk,
698 responses_so_far=responses_so_far,
699 responses_yielded=responses_yielded,
700 emitted_text_per_choice=emitted_text_per_choice,
701 finish_reason_per_choice=finish_reason_per_choice,
702 held_chars_per_choice=held_chars_per_choice,
703 is_final=is_final,
704 )
706 saw_tool_calls = False
707 saw_text_content = False
709 try:
710 async for item in response:
711 # v1 transforms only text. A chunk carrying tool_calls is passed
712 # through raw so function-calling turns are not dropped, but ONLY
713 # its tool-call fields are forwarded: content is stripped so any
714 # response text (in the same delta, or in another choice of an n>1
715 # chunk) can never bypass the transform. The original chunk is kept
716 # in responses_so_far so its text is still accumulated + redacted +
717 # emitted as synthetic deltas, and so the guardrail inspects the
718 # assembled tool calls at end of stream (see the block inspection
719 # below), matching block_only. finish_reason rides on the raw
720 # tool-only chunk, so it is not recorded for the text flush.
721 if self._chunk_has_tool_calls(item):
722 saw_tool_calls = True
723 responses_so_far.append(item)
724 last_chunk = item
725 # Fix #3 — flush accumulated text BEFORE the tool-call
726 # passthrough. Without this, a stream of text chunks that
727 # hasn't yet hit a sampled round can be trailed by a
728 # tool-call chunk carrying finish_reason="tool_calls"; an
729 # SSE-compliant client stops reading at that finish_reason
730 # and drops the end-of-stream text flush that would follow.
731 if saw_text_content:
732 async for out in _round(item, is_final=False):
733 yield out
734 # Fix #1 — pass finish_reason_per_choice into the
735 # passthrough so a mixed content+tool_call chunk defers its
736 # finish_reason to the final text terminator (see the
737 # _tool_call_passthrough_chunk docstring).
738 tool_only = self._tool_call_passthrough_chunk(
739 item,
740 finish_reason_per_choice=finish_reason_per_choice,
741 held_choices=_held_choices(held_chars_per_choice),
742 )
743 responses_yielded.append(tool_only)
744 yield tool_only
745 continue
747 if self._is_trailing_metadata_chunk(item):
748 responses_so_far.append(item)
749 continue
751 chunk_counter += 1
752 responses_so_far.append(item)
753 last_chunk = item
754 self._record_finish_reasons(item, finish_reason_per_choice)
755 if self._chunk_carries_text(item):
756 saw_text_content = True
757 # Skip the sampled round for a terminal chunk: the end-of-stream
758 # flush below processes it once with holdback forced to 0, so a
759 # sampled round here would guardrail the same content twice.
760 if (
761 not end_of_stream_only
762 and not self._chunk_has_finish_reason(item)
763 and chunk_counter % sampling_rate == 0
764 ):
765 async for out in _round(item, is_final=False):
766 yield out
768 # v1 does not transform streamed tool calls, but they must still go
769 # through the guardrail's block decision. Run the block_only inspection
770 # over the full assembled response so tool calls cannot bypass it.
771 #
772 # Pass a deep copy of responses_so_far — the block path routes through
773 # ``_process_streaming_block_only`` which mutates ``delta.content``
774 # in-place on the chunk objects it receives. For an n>1 chunk carrying
775 # text on one choice and tool_calls (with finish_reason) on another,
776 # ``has_stream_ended`` reads ``choices[0]`` alone and can miss the
777 # terminal signal, letting the block path rewrite the raw accumulator.
778 # The subsequent final ``_round`` would then re-read the already-mutated
779 # text, producing double-application for a non-idempotent guardrail or a
780 # ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow
781 # list copy wouldn't help — the mutation is on the chunk objects
782 # themselves — so we deepcopy.
783 if saw_tool_calls:
784 async for out in self._inspect_full_response_for_block(
785 endpoint_translation=endpoint_translation,
786 guardrail_to_apply=guardrail_to_apply,
787 request_data=request_data,
788 user_api_key_dict=user_api_key_dict,
789 responses_so_far=copy.deepcopy(responses_so_far),
790 responses_yielded=responses_yielded,
791 ):
792 yield out
794 async for out in self._emit_stream_tail(
795 last_chunk=last_chunk,
796 final_round=_round,
797 responses_so_far=responses_so_far,
798 responses_yielded=responses_yielded,
799 ):
800 yield out
801 except _StreamTerminated:
802 return
804 async def _emit_stream_tail(
805 self,
806 *,
807 last_chunk: object | None,
808 final_round: Callable[[object, bool], AsyncGenerator[object, None]],
809 responses_so_far: Sequence[object],
810 responses_yielded: list[object],
811 ) -> AsyncGenerator[object, None]:
812 """Flush the held text with holdback 0, then replay metadata-only chunks
813 (usage) so they land after the text and its finish_reason, as upstream sent them."""
814 if last_chunk is not None:
815 async for out in final_round(last_chunk, True):
816 yield out
817 for trailing in self._trailing_metadata_chunks(responses_so_far):
818 responses_yielded.append(trailing)
819 yield trailing
821 async def _inspect_full_response_for_block(
822 self,
823 *,
824 endpoint_translation: _EndpointTranslation,
825 guardrail_to_apply: CustomGuardrail,
826 request_data: dict,
827 user_api_key_dict: UserAPIKeyAuth,
828 responses_so_far: Sequence[object],
829 responses_yielded: Sequence[object],
830 ) -> AsyncGenerator[object, None]:
831 """Run the block-only guardrail inspection over the full assembled
832 response (text + tool calls) so nothing bypasses the block decision.
834 The guardrail's returned transforms are discarded here (v1 does not
835 transform tool calls); only its block decision matters. A block is
836 surfaced the same way as elsewhere: ModifyResponseException terminates the
837 stream via the shared block handler; a GenericGuardrailAPI block raises and
838 propagates, matching block_only.
839 """
840 from litellm.integrations.custom_guardrail import ModifyResponseException
842 try:
843 await endpoint_translation.process_output_streaming_response(
844 responses_so_far=responses_so_far,
845 guardrail_to_apply=guardrail_to_apply,
846 litellm_logging_obj=request_data.get("litellm_logging_obj"),
847 user_api_key_dict=user_api_key_dict,
848 request_data=request_data,
849 stream_transform_sink=None,
850 )
851 except ModifyResponseException as e:
852 if e.original_response is None:
853 e.original_response = responses_so_far
854 async for block_chunk in self.handle_streaming_block(
855 e,
856 endpoint_translation,
857 stream_started=bool(responses_yielded),
858 responses_so_far=responses_yielded,
859 ):
860 yield block_chunk
861 raise _StreamTerminated()
863 @staticmethod
864 def _chunk_has_tool_calls(item: object) -> bool:
865 for choice in _chunk_choices(item):
866 delta = getattr(choice, "delta", None)
867 if getattr(delta, "tool_calls", None):
868 return True
869 return False
871 @classmethod
872 def _is_trailing_metadata_chunk(cls, item: object) -> bool:
873 """True for a chunk that carries only stream metadata (no choices, or a
874 ``usage`` chunk whose deltas are empty); such chunks are replayed after
875 the final text flush instead of being folded into the transform."""
876 if not _chunk_choices(item):
877 return True
878 return (
879 getattr(item, "usage", None) is not None
880 and not cls._chunk_carries_text(item)
881 and not cls._chunk_has_finish_reason(item)
882 )
884 @classmethod
885 def _trailing_metadata_chunks(cls, items: Sequence[object]) -> tuple[object, ...]:
886 return tuple(item for item in items if cls._is_trailing_metadata_chunk(item))
888 @staticmethod
889 def _chunk_carries_text(item: object) -> bool:
890 """True if any choice in this chunk has non-empty string ``delta.content``."""
891 for choice in _chunk_choices(item):
892 delta = getattr(choice, "delta", None)
893 content = getattr(delta, "content", None)
894 if isinstance(content, str) and content != "":
895 return True
896 return False
898 @staticmethod
899 def _tool_call_passthrough_chunk(
900 item: object,
901 finish_reason_per_choice: "dict[int, str | None] | None" = None,
902 held_choices: frozenset[int] = frozenset(),
903 ) -> ModelResponseStream:
904 """Copy of a chunk carrying tool calls with all text content stripped.
906 Only tool_calls, role and finish_reason are forwarded; content is set to
907 None so response text can never be delivered raw (it flows through the
908 transform instead). Applies per choice so an n>1 chunk mixing a text
909 choice and a tool-call choice does not leak the text choice.
911 For a choice that carries BOTH text content AND tool_calls, or whose earlier
912 text is still withheld (``held_choices``), ``finish_reason`` is suppressed on
913 the passthrough and recorded on
914 ``finish_reason_per_choice`` (when provided) so the final synthetic text
915 chunk delivers it. Emitting the passthrough's ``finish_reason`` before the
916 text flush would let a spec-compliant SSE client stop reading at
917 ``finish_reason`` and silently drop the guardrailed text, defeating the
918 redaction purpose.
919 """
920 synthetic_choices: Final[list[StreamingChoices]] = []
921 for choice in _chunk_choices(item):
922 delta = getattr(choice, "delta", None)
923 idx = getattr(choice, "index", 0) or 0
924 original_finish = getattr(choice, "finish_reason", None)
925 has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != ""
926 text_pending = has_text or idx in held_choices
927 if text_pending and original_finish is not None and finish_reason_per_choice is not None:
928 finish_reason_per_choice[idx] = original_finish
929 passthrough_finish: str | None = None
930 else:
931 passthrough_finish = original_finish
932 synthetic_choices.append(
933 StreamingChoices(
934 index=idx,
935 delta=Delta(
936 content=None,
937 role=getattr(delta, "role", None),
938 tool_calls=getattr(delta, "tool_calls", None),
939 ),
940 finish_reason=passthrough_finish,
941 )
942 )
943 return ModelResponseStream(
944 id=getattr(item, "id", None),
945 created=getattr(item, "created", None),
946 model=getattr(item, "model", None),
947 choices=synthetic_choices,
948 )
950 @staticmethod
951 def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None:
952 for choice in _chunk_choices(item):
953 finish_reason = getattr(choice, "finish_reason", None)
954 if finish_reason is not None:
955 finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason
957 @staticmethod
958 def _chunk_has_finish_reason(item: object) -> bool:
959 choices: Final = _chunk_choices(item)
960 return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
962 def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object:
963 """Streaming flag resolution order (later wins): default < guardrail
964 attribute < guardrail_config dict < this callback's optional_params."""
965 attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default)
966 config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None)
967 config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value
968 return self.optional_params.get(name, config_value)
970 async def async_post_call_streaming_iterator_hook(
971 self,
972 user_api_key_dict: UserAPIKeyAuth,
973 response: AsyncIterable[object],
974 request_data: dict,
975 guardrail_to_apply: CustomGuardrail | None = None,
976 buffer_until_moderated_default: bool = False,
977 ) -> AsyncGenerator[Any, None]:
978 """
979 Passes the entire stream to the guardrail
981 This is useful for guardrails that need to see the entire response, such as PII masking.
983 See Aim guardrail implementation for an example - https://github.com/BerriAI/litellm/blob/d0e022cfacb8e9ebc5409bb652059b6fd97b45c0/litellm/proxy/guardrails/guardrail_hooks/aim.py#L168
985 Triggered by mode: 'post_call'
987 Supports sampling_rate parameter to control how often chunks are processed.
988 sampling_rate=1 means every chunk, sampling_rate=5 means every 5th chunk, etc.
989 """
991 # Local import avoids a module-level cyclic import with
992 # litellm.integrations.custom_guardrail.
993 from litellm.integrations.custom_guardrail import ModifyResponseException
995 if guardrail_to_apply is None:
996 guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
998 def _streaming_flag(name: str, default: object) -> Any:
999 return self.resolve_streaming_flag(guardrail_to_apply, name, default)
1001 sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
1002 # Only apply the guardrail at end of stream (not per chunk).
1003 end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False)
1004 # "block_only" (default) drops guardrail text rewrites on the streaming
1005 # path; "incremental_diff" emits them as synthetic deltas (see
1006 # _run_incremental_transform_stream).
1007 streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only")
1008 # Withhold every chunk until end-of-stream moderation passes, then
1009 # release the original chunks (clean) or only the block message
1010 # (blocked) -- moderating the whole response *before* any content
1011 # reaches the client. Only safe for allow/block guardrails: on
1012 # release the original chunks are replayed as-is, so a
1013 # content-rewriting guardrail (e.g. PII masking) would leak
1014 # unredacted content. Guarded below via mask_response_content.
1015 buffer_until_moderated: bool = _streaming_flag(
1016 "streaming_buffer_until_moderated", buffer_until_moderated_default
1017 )
1018 release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False)
1020 if (
1021 buffer_until_moderated
1022 and guardrail_to_apply is not None
1023 and getattr(guardrail_to_apply, "mask_response_content", False)
1024 ):
1025 verbose_proxy_logger.warning(
1026 "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s "
1027 "because mask_response_content=True -- buffered replay would release "
1028 "unredacted original chunks instead of the moderated output.",
1029 guardrail_to_apply.guardrail_name,
1030 )
1031 buffer_until_moderated = False
1033 if buffer_until_moderated and not release_on_scan:
1034 end_of_stream_only = True
1036 if guardrail_to_apply is None:
1037 async for item in response:
1038 yield item
1039 return
1041 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call
1042 if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True:
1043 verbose_proxy_logger.debug(
1044 "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s",
1045 guardrail_to_apply.guardrail_name,
1046 )
1047 async for item in response:
1048 yield item
1049 return
1051 mappings: Final = load_guardrail_translation_mappings()
1053 # Streaming text transformation (incremental_diff) diverges enough from the
1054 # block_only path that it runs as its own iterator. It requires a route we
1055 # can resolve up front to an OpenAI-chat handler (the only supported v1
1056 # surface); anything else falls back to the block_only behavior below.
1057 if streaming_transform_mode == "incremental_diff":
1058 transform_call_type: Final = self._resolve_transform_call_type(
1059 user_api_key_dict=user_api_key_dict,
1060 mappings=mappings,
1061 )
1062 if transform_call_type is not None:
1063 async for transformed_item in self._run_incremental_transform_stream(
1064 guardrail_to_apply=guardrail_to_apply,
1065 response=response,
1066 request_data=request_data,
1067 user_api_key_dict=user_api_key_dict,
1068 call_type=transform_call_type,
1069 sampling_rate=sampling_rate,
1070 end_of_stream_only=end_of_stream_only,
1071 mappings=mappings,
1072 ):
1073 yield transformed_item
1074 return
1075 verbose_proxy_logger.warning(
1076 "UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported "
1077 "for the OpenAI chat completions streaming path with a resolvable request route; "
1078 "falling back to block_only for %s",
1079 getattr(guardrail_to_apply, "guardrail_name", None),
1080 )
1082 # Infer call type from first chunk
1083 call_type = None
1084 chunk_counter = 0
1085 responses_so_far: Final[list[object]] = []
1086 responses_yielded: Final[list[object]] = []
1087 withheld_items: Final[list[object]] = [] # mutable-ok: streaming window must be released incrementally
1088 pending_end_of_stream_items: Final[list[object]] = []
1089 # Whether any real response chunk has been forwarded to the client.
1090 # Drives how a block terminates the stream: continue the in-progress
1091 # message (True) vs emit a standalone block message (False, buffered).
1092 chunks_yielded = False
1093 last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
1094 tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls
1096 async for item in response:
1097 chunk_counter += 1
1098 responses_so_far.append(item)
1100 # Infer call type from first chunk if not already done
1101 if call_type is None and user_api_key_dict.request_route is not None:
1102 call_types = get_call_types_for_route(user_api_key_dict.request_route)
1103 if call_types is not None:
1104 call_type = call_types[0].value
1106 if call_type is None:
1107 call_type = _infer_call_type(call_type=None, completion_response=item)
1109 # If call type not supported, just pass through all chunks
1110 if call_type is None or CallTypes(call_type) not in mappings:
1111 yield item
1112 async for remaining_item in response:
1113 yield remaining_item
1114 return
1116 # If end_of_stream_only mode, yield chunks without processing.
1117 # When buffering, withhold them instead -- they are released (or
1118 # replaced by the block message) only after end-of-stream
1119 # moderation runs below.
1120 if end_of_stream_only:
1121 if not buffer_until_moderated:
1122 endpoint_translation = mappings[CallTypes(call_type)]()
1123 stream_has_ended = hasattr(
1124 endpoint_translation, "_check_streaming_has_ended"
1125 ) and endpoint_translation._check_streaming_has_ended(responses_so_far)
1126 if pending_end_of_stream_items or stream_has_ended:
1127 pending_end_of_stream_items.append(item)
1128 else:
1129 chunks_yielded = True
1130 responses_yielded.append(item)
1131 yield item
1132 else:
1133 withheld_items.append(item)
1134 continue
1136 # Process chunk based on sampling rate
1137 if buffer_until_moderated:
1138 withheld_items.append(item)
1139 if chunk_counter % sampling_rate == 0:
1140 endpoint_translation = mappings[CallTypes(call_type)]()
1141 scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
1142 if scan_key is not None:
1143 tool_calls_in_flight = scan_key.tool_calls_in_flight
1144 hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight)
1145 if _is_redundant_scan(scan_key, last_scan_key):
1146 verbose_proxy_logger.debug(
1147 "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
1148 chunk_counter,
1149 guardrail_to_apply.guardrail_name,
1150 )
1151 if buffer_until_moderated:
1152 if hold_window:
1153 continue
1154 for withheld_item in withheld_items:
1155 chunks_yielded = True
1156 responses_yielded.append(withheld_item)
1157 yield withheld_item
1158 withheld_items.clear()
1159 else:
1160 chunks_yielded = True
1161 responses_yielded.append(item)
1162 yield item
1163 continue
1165 verbose_proxy_logger.debug(
1166 "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
1167 chunk_counter,
1168 sampling_rate,
1169 guardrail_to_apply.guardrail_name,
1170 )
1172 original_items = (
1173 tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),)
1174 )
1176 try:
1177 await endpoint_translation.process_output_streaming_response(
1178 responses_so_far=responses_so_far,
1179 guardrail_to_apply=guardrail_to_apply,
1180 litellm_logging_obj=request_data.get("litellm_logging_obj"),
1181 user_api_key_dict=user_api_key_dict,
1182 request_data=request_data,
1183 )
1184 except ModifyResponseException as e:
1185 if e.original_response is None:
1186 e.original_response = responses_so_far
1187 # Guardrail blocked the response mid-stream. Emit a clean
1188 # terminating SSE sequence delivering the block message
1189 # instead of letting the exception propagate into a bare
1190 # `data: {"error": ...}` blob (which truncates the stream).
1191 # Chunks have already been forwarded here, so the block
1192 # continues the in-progress message (stream_started=True).
1193 # The current chunk was appended to responses_so_far but not
1194 # yet yielded, so exclude it: the continuation must reflect
1195 # only what the client has actually received.
1196 async for block_chunk in self.handle_streaming_block(
1197 e,
1198 endpoint_translation,
1199 stream_started=chunks_yielded,
1200 responses_so_far=responses_yielded,
1201 ):
1202 yield block_chunk
1203 return
1204 except HTTPException as e:
1205 # Response already started (we already yielded chunks); cannot send 400.
1206 async for error_item in self.emit_streaming_http_error(
1207 e,
1208 call_type,
1209 responses_so_far,
1210 request_data,
1211 endpoint_translation=endpoint_translation,
1212 stream_started=chunks_yielded,
1213 responses_yielded=responses_yielded,
1214 ):
1215 yield error_item
1216 return
1217 if scan_key is not None:
1218 last_scan_key = scan_key
1219 if hold_window:
1220 verbose_proxy_logger.debug(
1221 "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window",
1222 len(withheld_items),
1223 guardrail_to_apply.guardrail_name,
1224 )
1225 withheld_items[:] = original_items
1226 continue
1227 for original_item in original_items:
1228 chunks_yielded = True
1229 responses_yielded.append(original_item)
1230 yield original_item
1231 withheld_items.clear()
1232 else:
1233 if not buffer_until_moderated:
1234 chunks_yielded = True
1235 responses_yielded.append(item)
1236 yield item
1238 # Stream has ended - do final processing with all collected chunks
1239 if call_type is not None and CallTypes(call_type) in mappings:
1240 verbose_proxy_logger.debug(
1241 "Processing final streaming response with all %s chunks for guardrail %s",
1242 len(responses_so_far),
1243 guardrail_to_apply.guardrail_name,
1244 )
1246 endpoint_translation = mappings[CallTypes(call_type)]()
1248 buffered_items: Final = (
1249 tuple(copy.deepcopy(withheld_items))
1250 if buffer_until_moderated and release_on_scan and not end_of_stream_only
1251 else tuple(withheld_items)
1252 if buffer_until_moderated
1253 else None
1254 )
1255 end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
1256 if _is_redundant_scan(end_scan_key, last_scan_key):
1257 verbose_proxy_logger.debug(
1258 "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
1259 guardrail_to_apply.guardrail_name,
1260 )
1261 for buffered_item in buffered_items or ():
1262 yield buffered_item
1263 for pending_item in pending_end_of_stream_items:
1264 responses_yielded.append(pending_item)
1265 yield pending_item
1266 return
1268 try:
1269 await endpoint_translation.process_output_streaming_response(
1270 responses_so_far=responses_so_far,
1271 guardrail_to_apply=guardrail_to_apply,
1272 litellm_logging_obj=request_data.get("litellm_logging_obj"),
1273 user_api_key_dict=user_api_key_dict,
1274 request_data=request_data,
1275 )
1276 # Moderation passed: release the withheld original chunks.
1277 if buffered_items is not None:
1278 for buffered_item in buffered_items:
1279 yield buffered_item
1280 for pending_item in pending_end_of_stream_items:
1281 responses_yielded.append(pending_item)
1282 yield pending_item
1283 except ModifyResponseException as e:
1284 if e.original_response is None:
1285 e.original_response = responses_so_far
1286 # Block detected during end-of-stream processing. Emit a clean
1287 # terminating SSE sequence with the block message rather than
1288 # propagating into a bare error blob that truncates the stream.
1289 # The withheld original chunks are never released.
1290 async for block_chunk in self.handle_streaming_block(
1291 e,
1292 endpoint_translation,
1293 stream_started=bool(responses_yielded),
1294 responses_so_far=responses_yielded,
1295 ):
1296 yield block_chunk
1297 return
1298 except HTTPException as e:
1299 async for error_item in self.emit_streaming_http_error(
1300 e,
1301 call_type,
1302 responses_so_far,
1303 request_data,
1304 endpoint_translation=endpoint_translation,
1305 stream_started=bool(responses_yielded),
1306 responses_yielded=responses_yielded,
1307 ):
1308 yield error_item