mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
feat(policy_engine): execute post_call guardrail pipelines on streaming responses
This commit is contained in:
mateo-berri
parent
c64318cfbd
commit
c5bcf3a735
5 changed files with 417 additions and 33 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"}],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue