mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge pull request #40284 from BerriAI/litellm_legacy_hook_streaming_pipeline_step
feat(guardrails): run legacy post-call hooks as streaming pipeline steps
This commit is contained in:
Mateo Wang
GitHub
commit
21ae29b759
11 changed files with 915 additions and 103 deletions
|
|
@ -213,11 +213,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
"""
|
||||
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
|
||||
def post_call_hook_response(self, response: object) -> object:
|
||||
if not isinstance(response, ModelResponse):
|
||||
return response
|
||||
return self.adapter.translate_openai_response_to_anthropic(response)
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_usage_response(
|
||||
responses_so_far: Sequence[object],
|
||||
|
|
|
|||
|
|
@ -61,6 +61,20 @@ class BaseTranslation(ABC):
|
|||
on every other translation are undeliverable: the pipeline executor
|
||||
discards them and releases the original chunks."""
|
||||
|
||||
assembles_streamed_response: ClassVar[bool] = False
|
||||
"""Whether ``process_output_streaming_response`` stores the assembled response of an
|
||||
ended stream under ``request_data["response"]`` before scanning it, the way the chat,
|
||||
Responses, and Messages translations do. A streaming pipeline runs a guardrail that only
|
||||
has the legacy post-call hook against that response, so on a translation without it such
|
||||
a guardrail keeps running on its own."""
|
||||
|
||||
def post_call_hook_response(self, response: object) -> object:
|
||||
"""The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
|
||||
the object the translation stores under ``request_data["response"]`` while scanning an
|
||||
ended stream. Chat and Responses scan that shape already; a translation that scans a
|
||||
different one (Messages scans an OpenAI-shaped ModelResponse) overrides this."""
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict: Any | None,
|
||||
|
|
|
|||
|
|
@ -81,6 +81,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
"""
|
||||
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -432,6 +432,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
"""
|
||||
|
||||
delivers_ended_stream_rewrites = True
|
||||
assembles_streamed_response = True
|
||||
|
||||
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3300,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
has completed.
|
||||
|
||||
Guardrails routed through unified_guardrail are skipped, since they already ran
|
||||
via its streaming iterator. Guardrails that override
|
||||
async_post_call_success_hook directly run here, including those that implement
|
||||
apply_guardrail but keep their native lifecycle hooks.
|
||||
via its streaming iterator, and so are guardrails a post_call policy pipeline
|
||||
manages, since the pipeline ran them against the buffered stream. Guardrails
|
||||
that override async_post_call_success_hook directly run here, including those
|
||||
that implement apply_guardrail but keep their native lifecycle hooks.
|
||||
|
||||
This is audit-only — content has already been delivered to the client.
|
||||
|
||||
|
|
@ -3312,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing:
|
|||
_response = assembled_response
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router as _global_llm_router
|
||||
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
|
||||
from litellm.proxy.utils import (
|
||||
_check_and_merge_model_level_guardrails,
|
||||
stream_gated_guardrail_names,
|
||||
)
|
||||
|
||||
guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router)
|
||||
stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict)
|
||||
for cb in litellm.callbacks:
|
||||
if not isinstance(cb, CustomGuardrail):
|
||||
continue
|
||||
if cb.guardrail_name in stream_gated:
|
||||
continue
|
||||
if not cb.should_run_guardrail(
|
||||
data=guardrail_data,
|
||||
event_type=GuardrailEventHooks.post_call,
|
||||
|
|
|
|||
|
|
@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
|
|||
return None if texts is None else tuple(texts)
|
||||
|
||||
|
||||
def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]:
|
||||
return tuple(texts or ())
|
||||
|
||||
|
||||
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
|
||||
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
|
||||
|
||||
|
|
@ -133,6 +137,94 @@ class _StreamRewriteObserver(CustomGuardrail):
|
|||
return outputs
|
||||
|
||||
|
||||
class _ScannedTextRecorder(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str) -> None:
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.inputs: GenericGuardrailAPIInputs | None = None
|
||||
|
||||
@_logged_by_inner_guardrail
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.inputs = inputs
|
||||
return inputs
|
||||
|
||||
|
||||
class _LegacyHookStreamAdapter(CustomGuardrail):
|
||||
"""Runs a guardrail that only implements the legacy post-call hook (no unified
|
||||
``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The
|
||||
endpoint translation hands it the texts it scanned plus the assembled response under
|
||||
``request_data["response"]``; the hook gets that response in the shape its route gives
|
||||
non-streaming hooks, an exception it raises ends the stream through the executor's
|
||||
fail/error classification, and the response it hands back, or the one it changed in place
|
||||
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
|
||||
client through the translation's ended-stream write-back. A
|
||||
replacement whose scanned texts do not line up with the originals, or whose tool calls
|
||||
differ from them, is undeliverable, so the executor releases the original chunks. A stream
|
||||
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
|
||||
long as the hook left the tool calls alone."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: CustomGuardrail,
|
||||
endpoint_translation: "BaseTranslation",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> None:
|
||||
super().__init__(guardrail_name=inner.guardrail_name)
|
||||
self.inner: Final = inner
|
||||
self.endpoint_translation: Final = endpoint_translation
|
||||
self.user_api_key_dict: Final = user_api_key_dict
|
||||
|
||||
def structured_messages_cover_full_request(self) -> bool:
|
||||
return self.inner.structured_messages_cover_full_request()
|
||||
|
||||
@_logged_by_inner_guardrail
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response"))
|
||||
replacement: Final = await self.inner.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=self.user_api_key_dict,
|
||||
response=hooked,
|
||||
)
|
||||
rewrite: Final = hooked if replacement is None else replacement
|
||||
if rewrite is None:
|
||||
return inputs
|
||||
rescanned: Final = await self._rescan(rewrite, logging_obj)
|
||||
if rescanned is None:
|
||||
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
|
||||
rewritten: Final = rescanned.get("texts")
|
||||
if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))):
|
||||
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
|
||||
if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")):
|
||||
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
|
||||
if not rewritten:
|
||||
return inputs
|
||||
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
|
||||
return rewritten_inputs
|
||||
|
||||
async def _rescan(
|
||||
self, response: object, logging_obj: "LiteLLMLoggingObj | None"
|
||||
) -> GenericGuardrailAPIInputs | None:
|
||||
recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown")
|
||||
await self.endpoint_translation.process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=recorder,
|
||||
litellm_logging_obj=logging_obj,
|
||||
user_api_key_dict=self.user_api_key_dict,
|
||||
)
|
||||
return recorder.inputs
|
||||
|
||||
|
||||
def _prepare_hook_input(
|
||||
step: PipelineStep,
|
||||
callback: CustomGuardrail,
|
||||
|
|
@ -300,18 +392,29 @@ class PipelineExecutor:
|
|||
endpoint_translation: "BaseTranslation",
|
||||
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
|
||||
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
|
||||
user_api_key_dict: "UserAPIKeyAuth | None",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
) -> None:
|
||||
"""Run one streaming post_call step through the endpoint translation, delivering
|
||||
text and tool-call rewrites on translations that support ended-stream write-back. A
|
||||
rewrite that cannot reach the client yet (one on a translation without write-back, or
|
||||
one the translation refused with ``UndeliverableStreamRewrite``) is discarded: the
|
||||
buffered chunks go back to the originals and the step passes, so the client gets the
|
||||
stream the merge base sent."""
|
||||
observer: Final = _StreamRewriteObserver(callback)
|
||||
guardrail without the unified interface runs its legacy post-call hook against the
|
||||
assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the
|
||||
client yet (one on a translation without write-back, one that drops or adds a tool call,
|
||||
or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is
|
||||
discarded: the buffered chunks go back to the originals and the step passes, so the
|
||||
client gets the stream the merge base sent, and the guardrail stays out of the
|
||||
applied-guardrails header since its output never reached the client. The response an
|
||||
earlier step's translation stored under ``request_data["response"]`` is dropped first,
|
||||
so this step's hook sees the stream as the steps before it left it."""
|
||||
scanner: Final = (
|
||||
callback
|
||||
if PipelineExecutor.supports_unified_execution(callback)
|
||||
else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict)
|
||||
)
|
||||
observer: Final = _StreamRewriteObserver(scanner)
|
||||
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
|
||||
originals: Final = copy.deepcopy(streaming_chunks)
|
||||
hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored
|
||||
try:
|
||||
if deliver_rewrites:
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
|
|
@ -332,11 +435,12 @@ class PipelineExecutor:
|
|||
)
|
||||
except UndeliverableStreamRewrite:
|
||||
_release_original_chunks(step.guardrail, streaming_chunks, originals)
|
||||
else:
|
||||
if observer.changed_tool_call_count or (
|
||||
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
|
||||
):
|
||||
_release_original_chunks(step.guardrail, streaming_chunks, originals)
|
||||
return
|
||||
if observer.changed_tool_call_count or (
|
||||
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
|
||||
):
|
||||
_release_original_chunks(step.guardrail, streaming_chunks, originals)
|
||||
return
|
||||
if not callback.records_own_guardrail_information:
|
||||
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
|
||||
|
||||
|
|
@ -396,11 +500,11 @@ class PipelineExecutor:
|
|||
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:
|
||||
if endpoint_translation is None:
|
||||
return (
|
||||
"error",
|
||||
None,
|
||||
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
|
||||
f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation",
|
||||
None,
|
||||
)
|
||||
await PipelineExecutor._run_streaming_step(
|
||||
|
|
@ -456,10 +560,22 @@ class PipelineExecutor:
|
|||
|
||||
@staticmethod
|
||||
def supports_unified_execution(callback: CustomGuardrail) -> bool:
|
||||
"""Whether this guardrail runs through the unified apply_guardrail path,
|
||||
the interface streaming pipeline execution requires."""
|
||||
"""Whether this guardrail runs through the unified apply_guardrail path."""
|
||||
return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
|
||||
|
||||
@staticmethod
|
||||
def supports_streaming_execution(callback: CustomGuardrail) -> bool:
|
||||
"""Whether a streaming pipeline step can run this guardrail against the buffered
|
||||
stream: through the unified path, or through its post-call hook on the assembled
|
||||
response when that hook is its only streaming path. A guardrail with its own
|
||||
streaming iterator hook, or with neither hook, keeps running on its own."""
|
||||
callback_type: Final = type(callback)
|
||||
return PipelineExecutor.supports_unified_execution(callback) or (
|
||||
callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook
|
||||
and callback_type.async_post_call_streaming_iterator_hook
|
||||
is CustomLogger.async_post_call_streaming_iterator_hook
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
|
||||
"""Look up an initialized guardrail callback by name from litellm.callbacks."""
|
||||
|
|
|
|||
|
|
@ -460,7 +460,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe
|
|||
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
|
||||
|
||||
|
||||
def _pipeline_managed_guardrail_names(
|
||||
def pipeline_managed_guardrail_names(
|
||||
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
|
||||
) -> frozenset[str]:
|
||||
return _pipeline_step_guardrail_names(
|
||||
|
|
@ -523,9 +523,17 @@ def _merge_pipeline_metadata_writes(
|
|||
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
|
||||
|
||||
|
||||
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
|
||||
def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool:
|
||||
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
|
||||
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
|
||||
if callback is None:
|
||||
return False
|
||||
if PipelineExecutor.supports_unified_execution(callback):
|
||||
return True
|
||||
return (
|
||||
translation is not None
|
||||
and type(translation).assembles_streamed_response
|
||||
and PipelineExecutor.supports_streaming_execution(callback)
|
||||
)
|
||||
|
||||
|
||||
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
|
||||
|
|
@ -582,7 +590,7 @@ def _withdraw_deferred_claims(
|
|||
outside_by_policy: Final = MappingProxyType(
|
||||
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
|
||||
)
|
||||
running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(
|
||||
running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
|
||||
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
|
||||
)
|
||||
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
|
||||
|
|
@ -657,37 +665,51 @@ def _body_selected_deferrals(
|
|||
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
|
||||
|
||||
|
||||
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
|
||||
unsupported: Final = tuple(
|
||||
def _pipeline_unsupported_streaming_guardrails(
|
||||
pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
|
||||
) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
dict.fromkeys(
|
||||
step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
|
||||
step.guardrail
|
||||
for step in pipeline.steps
|
||||
if not _pipeline_step_supports_streaming(step.guardrail, translation)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _pipeline_is_streamable(
|
||||
policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
|
||||
) -> bool:
|
||||
unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
||||
if not unsupported:
|
||||
return True
|
||||
verbose_proxy_logger.warning(
|
||||
"Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
|
||||
"which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
|
||||
"Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they "
|
||||
"need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a "
|
||||
"route whose translation assembles the streamed response. The stream skips the pipeline and its "
|
||||
"guardrails run on their own: %s",
|
||||
policy_name,
|
||||
", ".join(unsupported),
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
return resolve_endpoint_translation(user_api_key_dict, None) is not None
|
||||
def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None":
|
||||
resolved: Final = resolve_endpoint_translation(user_api_key_dict, None)
|
||||
return None if resolved is None else resolved[1]
|
||||
|
||||
|
||||
def _stream_gated_guardrail_names(
|
||||
def stream_gated_guardrail_names(
|
||||
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> frozenset[str]:
|
||||
if not _route_supports_streaming_pipelines(user_api_key_dict):
|
||||
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
|
||||
if translation is None:
|
||||
return frozenset()
|
||||
return _pipeline_step_guardrail_names(
|
||||
tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in _post_call_pipelines(request_data)
|
||||
if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
|
||||
if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -699,16 +721,19 @@ def _streamable_post_call_pipelines(
|
|||
The post_call pipelines a streaming response can be gated through.
|
||||
|
||||
Streaming pipelines scan the buffered stream through the endpoint guardrail
|
||||
translation of the request route, so every step's guardrail needs the
|
||||
unified apply_guardrail interface and the route needs a translation. A
|
||||
pipeline that cannot be run that way yet is left out and its guardrails
|
||||
run on the stream on their own, the way they did before pipelines ran on
|
||||
streams at all, with a warning naming the pipeline.
|
||||
translation of the request route, so every step's guardrail needs either the
|
||||
unified apply_guardrail interface or, on a route whose translation assembles
|
||||
the streamed response, a post-call hook that is its only streaming path, and
|
||||
the route needs a translation. A pipeline that
|
||||
cannot be run that way yet is left out and its guardrails run on the stream
|
||||
on their own, the way they did before pipelines ran on streams at all, with
|
||||
a warning naming the pipeline.
|
||||
"""
|
||||
post_call_pipelines: Final = _post_call_pipelines(request_data)
|
||||
if not post_call_pipelines:
|
||||
return ()
|
||||
if not _route_supports_streaming_pipelines(user_api_key_dict):
|
||||
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
|
||||
if translation is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
|
||||
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
|
||||
|
|
@ -720,7 +745,7 @@ def _streamable_post_call_pipelines(
|
|||
return tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in post_call_pipelines
|
||||
if _pipeline_is_streamable(policy_name, pipeline)
|
||||
if _pipeline_is_streamable(policy_name, pipeline, translation)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2110,7 +2135,7 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
# Get pipeline-managed guardrails to skip in normal loop
|
||||
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
|
||||
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
|
||||
caps: Final = ProxyLogging._callback_capabilities()
|
||||
# Skip the per-request callback walk entirely when nothing in
|
||||
|
|
@ -3114,7 +3139,7 @@ class ProxyLogging:
|
|||
if pipeline_response is not None:
|
||||
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
|
||||
|
||||
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
|
||||
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
|
||||
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
|
||||
try:
|
||||
# Merge model-level guardrails before checking which guardrails to run
|
||||
|
|
@ -3430,7 +3455,7 @@ class ProxyLogging:
|
|||
_cached_guardrail_data: dict | None = None
|
||||
_guardrail_data_computed = False
|
||||
pipeline_gated: Final = (
|
||||
_stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
|
||||
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
|
||||
)
|
||||
|
||||
for callback in litellm.callbacks:
|
||||
|
|
|
|||
|
|
@ -2270,3 +2270,29 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
|
|||
assert open_key == StreamingScanKey(texts=("hi",))
|
||||
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
|
||||
assert ended_key != open_key
|
||||
|
||||
|
||||
class TestAnthropicMessagesHandlerPostCallHookResponse:
|
||||
def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
|
||||
from litellm.types.utils import Choices, Message, ModelResponse, Usage
|
||||
|
||||
assembled = ModelResponse(
|
||||
id="msg_1",
|
||||
model="claude",
|
||||
choices=[Choices(message=Message(role="assistant", content="hello world"), finish_reason="stop")],
|
||||
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
|
||||
)
|
||||
|
||||
hook_response = AnthropicMessagesHandler().post_call_hook_response(assembled)
|
||||
|
||||
assert hook_response["type"] == "message"
|
||||
assert hook_response["role"] == "assistant"
|
||||
assert hook_response["content"] == [{"type": "text", "text": "hello world"}]
|
||||
assert hook_response["stop_reason"] == "end_turn"
|
||||
assert hook_response["usage"]["input_tokens"] == 1
|
||||
assert hook_response["usage"]["output_tokens"] == 2
|
||||
|
||||
def test_anything_else_reaches_the_hook_untouched(self):
|
||||
native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]}
|
||||
|
||||
assert AnthropicMessagesHandler().post_call_hook_response(native) is native
|
||||
|
|
|
|||
|
|
@ -675,8 +675,6 @@ async def test_guardrail_not_found_uses_on_fail(monkeypatch):
|
|||
],
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode=pipeline.mode,
|
||||
|
|
@ -1149,6 +1147,7 @@ def _assert_passed_with_discard_warning(result, caplog):
|
|||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
|
||||
assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1329,3 +1328,350 @@ async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewri
|
|||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
class _LegacyHookGuardrail(CustomGuardrail):
|
||||
"""A guardrail with only the legacy post-call hook: it never defines apply_guardrail."""
|
||||
|
||||
def __init__(self, replacement=None, raises=None, guardrail_name="masker", rewrite_in_place=None):
|
||||
super().__init__(guardrail_name=guardrail_name, event_hook="post_call", default_on=True)
|
||||
self.replacement = replacement
|
||||
self.raises = raises
|
||||
self.rewrite_in_place = rewrite_in_place
|
||||
self.calls = []
|
||||
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
self.calls.append({"data": data, "user_api_key_dict": user_api_key_dict, "response": response})
|
||||
if self.raises is not None:
|
||||
raise self.raises
|
||||
if self.rewrite_in_place is not None:
|
||||
response["text"] = self.rewrite_in_place
|
||||
return self.replacement
|
||||
|
||||
|
||||
class _NativeHooksGuardrail(_LegacyHookGuardrail):
|
||||
use_native_lifecycle_hooks = True
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail")
|
||||
|
||||
|
||||
class _LegacyScanningTranslation:
|
||||
"""Stores the assembled response under request_data["response"] before scanning, like the
|
||||
chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one
|
||||
text per entry of a replacement's "texts"."""
|
||||
|
||||
delivers_ended_stream_rewrites = True
|
||||
|
||||
def post_call_hook_response(self, response):
|
||||
return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]}
|
||||
|
||||
async def process_output_streaming_response(
|
||||
self,
|
||||
responses_so_far,
|
||||
guardrail_to_apply,
|
||||
litellm_logging_obj=None,
|
||||
user_api_key_dict=None,
|
||||
request_data=None,
|
||||
deliver_ended_stream_rewrites=False,
|
||||
):
|
||||
request_data.setdefault(
|
||||
"response", {"text": responses_so_far[0]["text"], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}
|
||||
)
|
||||
outputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]},
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
responses_so_far[0]["text"] = outputs["texts"][0]
|
||||
return responses_so_far
|
||||
|
||||
async def process_output_response(
|
||||
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
|
||||
):
|
||||
inputs = {"texts": [response["text"]] if "text" in response else list(response["texts"])}
|
||||
if response.get("tool_calls"):
|
||||
inputs["tool_calls"] = list(response["tool_calls"])
|
||||
await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"response": response},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
class _ToolOnlyLegacyScanningTranslation(_LegacyScanningTranslation):
|
||||
"""Like the Messages handler on a tool-only message: the ended-stream scan omits "texts" from
|
||||
the inputs, while the non-streaming scan of the same response sends an empty list."""
|
||||
|
||||
async def process_output_streaming_response(
|
||||
self,
|
||||
responses_so_far,
|
||||
guardrail_to_apply,
|
||||
litellm_logging_obj=None,
|
||||
user_api_key_dict=None,
|
||||
request_data=None,
|
||||
deliver_ended_stream_rewrites=False,
|
||||
):
|
||||
request_data.setdefault("response", {"text": "", "tool_calls": [dict(responses_so_far[0]["tool_call"])]})
|
||||
await guardrail_to_apply.apply_guardrail(
|
||||
inputs={"tool_calls": [dict(responses_so_far[0]["tool_call"])]},
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
return responses_so_far
|
||||
|
||||
async def process_output_response(
|
||||
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
|
||||
):
|
||||
await guardrail_to_apply.apply_guardrail(
|
||||
inputs={"texts": [], "tool_calls": list(response.get("tool_calls") or [])},
|
||||
request_data={"response": response},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
def _tool_only_chunk():
|
||||
return {"text": "", "tool_call": _chunk()["tool_call"]}
|
||||
|
||||
|
||||
def _native(text):
|
||||
return {"native": True, "text": text, "tool_calls": [_chunk()["tool_call"]]}
|
||||
|
||||
|
||||
def _legacy_replacement(*texts, tool_calls=None):
|
||||
return {"texts": list(texts), "tool_calls": [_chunk()["tool_call"]] if tool_calls is None else tool_calls}
|
||||
|
||||
|
||||
async def _run_legacy_streaming_step(
|
||||
monkeypatch, guardrail, chunks, on_fail="block", on_error="next", translation=None
|
||||
):
|
||||
return await _run_legacy_streaming_steps(
|
||||
monkeypatch, [guardrail], chunks, on_fail=on_fail, on_error=on_error, translation=translation
|
||||
)
|
||||
|
||||
|
||||
async def _run_legacy_streaming_steps(
|
||||
monkeypatch, guardrails, chunks, on_fail="block", on_error="next", translation=None
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", list(guardrails))
|
||||
return await PipelineExecutor.execute_steps(
|
||||
steps=[
|
||||
PipelineStep(
|
||||
guardrail=guardrail.guardrail_name,
|
||||
on_pass="next" if position + 1 < len(guardrails) else "allow",
|
||||
on_fail=on_fail,
|
||||
on_error=on_error,
|
||||
)
|
||||
for position, guardrail in enumerate(guardrails)
|
||||
],
|
||||
mode="post_call",
|
||||
data={"model": "m"},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="p",
|
||||
streaming_chunks=chunks,
|
||||
endpoint_translation=_LegacyScanningTranslation() if translation is None else translation,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("guardrail_class", [_LegacyHookGuardrail, _NativeHooksGuardrail])
|
||||
async def test_streaming_step_runs_legacy_hook_and_delivers_its_rewrite(monkeypatch, caplog, guardrail_class):
|
||||
guardrail = guardrail_class(replacement=_legacy_replacement("[REWRITTEN] hello world"))
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert chunks[0]["text"] == "[REWRITTEN] hello world"
|
||||
assert [call["response"] for call in guardrail.calls] == [_native("hello world")]
|
||||
assert guardrail.calls[0]["data"]["model"] == "m"
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_delivers_a_legacy_rewrite_made_in_place(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(rewrite_in_place="[REWRITTEN] hello world")
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert chunks[0]["text"] == "[REWRITTEN] hello world"
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_passes_untouched_when_legacy_hook_returns_none(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=None)
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert len(guardrail.calls) == 1
|
||||
assert chunks == [_chunk()]
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_blocks_with_the_legacy_hook_exception(monkeypatch):
|
||||
exc = HTTPException(status_code=400, detail={"error": "output blocked"})
|
||||
chunks = [_chunk()]
|
||||
|
||||
result = await _run_legacy_streaming_step(monkeypatch, _LegacyHookGuardrail(raises=exc), chunks)
|
||||
|
||||
assert result.terminal_action == "block"
|
||||
assert [step.outcome for step in result.step_results] == ["fail"]
|
||||
assert result.original_exception is exc
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_takes_on_error_when_legacy_hook_crashes(monkeypatch):
|
||||
chunks = [_chunk()]
|
||||
|
||||
result = await _run_legacy_streaming_step(
|
||||
monkeypatch, _LegacyHookGuardrail(raises=ValueError("boom")), chunks, on_error="block"
|
||||
)
|
||||
|
||||
assert result.terminal_action == "block"
|
||||
assert [step.outcome for step in result.step_results] == ["error"]
|
||||
assert result.step_results[0].error_detail == "boom"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("split", "in two"))
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog):
|
||||
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
|
||||
guardrail = _LegacyHookGuardrail(
|
||||
replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call])
|
||||
)
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[]))
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alone(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=None)
|
||||
chunks = [_tool_only_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(
|
||||
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
|
||||
)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
assert chunks == [_tool_only_chunk()]
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog):
|
||||
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call]))
|
||||
chunks = [_tool_only_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(
|
||||
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
|
||||
)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_tool_only_chunk()]
|
||||
|
||||
|
||||
class _NoHooksGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
|
||||
class _IteratorAndLegacyHookGuardrail(_LegacyHookGuardrail):
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
async for item in response:
|
||||
yield item
|
||||
|
||||
|
||||
class _UnscannableRewriteTranslation(_LegacyScanningTranslation):
|
||||
"""Like the chat handler on a response whose choices are plain dicts: the non-streaming scan
|
||||
never hands anything to the guardrail."""
|
||||
|
||||
async def process_output_response(
|
||||
self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
|
||||
):
|
||||
return response
|
||||
|
||||
|
||||
def test_streaming_execution_runs_legacy_hooks_only_when_that_hook_is_their_only_streaming_path():
|
||||
assert PipelineExecutor.supports_streaming_execution(_LegacyHookGuardrail()) is True
|
||||
assert PipelineExecutor.supports_streaming_execution(_NativeHooksGuardrail()) is True
|
||||
assert PipelineExecutor.supports_streaming_execution(_IteratorAndLegacyHookGuardrail()) is False
|
||||
assert PipelineExecutor.supports_streaming_execution(_NoHooksGuardrail(guardrail_name="neither")) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_rescan(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("hello [MASKED]"))
|
||||
chunks = [_chunk()]
|
||||
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks, translation=_UnscannableRewriteTranslation())
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog)
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch):
|
||||
masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world"))
|
||||
auditor = _LegacyHookGuardrail(replacement=None, guardrail_name="auditor")
|
||||
chunks = [_chunk()]
|
||||
|
||||
result = await _run_legacy_streaming_steps(monkeypatch, [masker, auditor], chunks, on_fail="next")
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass", "pass"]
|
||||
assert chunks[0]["text"] == "[REWRITTEN] hello world"
|
||||
assert [call["response"] for call in masker.calls] == [_native("hello world")]
|
||||
assert [call["response"] for call in auditor.calls] == [_native("[REWRITTEN] hello world")]
|
||||
|
|
|
|||
|
|
@ -671,6 +671,95 @@ async def test_deferred_stream_guardrails_run_native_hook_when_opted_out(monkeyp
|
|||
assert routed.native_hooks_ran == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_stream_guardrails_skip_pipeline_managed_native_hook(monkeypatch):
|
||||
"""A post_call pipeline step already ran the opted-out guardrail's own hook against
|
||||
the buffered stream, so the deferred audit must not run it a second time."""
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
|
||||
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
||||
captured_data={
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
||||
},
|
||||
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
|
||||
captured_logging_obj=_streaming_logging_obj(),
|
||||
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
||||
cache_hit=False,
|
||||
)
|
||||
|
||||
assert pipeline_managed.native_hooks_ran == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_stream_guardrails_run_native_hook_whose_pipeline_could_not_stream(monkeypatch):
|
||||
"""A pipeline step with neither streaming interface keeps the whole pipeline off the
|
||||
stream, so the deferred audit is the only place the opted-out guardrail's own hook
|
||||
still runs, the way it did before pipelines ran on streams."""
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
class NeitherHookGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
neither = NeitherHookGuardrail(guardrail_name="gr-neither", event_hook=GuardrailEventHooks.post_call)
|
||||
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed, neither])
|
||||
pipeline = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="keeps_native", on_fail="next"),
|
||||
PipelineStep(guardrail="gr-neither", on_fail="block"),
|
||||
],
|
||||
)
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
||||
captured_data={
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
||||
},
|
||||
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"),
|
||||
captured_logging_obj=_streaming_logging_obj(),
|
||||
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
||||
cache_hit=False,
|
||||
)
|
||||
|
||||
assert pipeline_managed.native_hooks_ran == ["post_call"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_stream_guardrails_run_native_hook_on_route_without_translation(monkeypatch):
|
||||
"""A route with no endpoint guardrail translation cannot gate the stream through its
|
||||
pipelines, so the deferred audit still owes the opted-out guardrail its own hook."""
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
pipeline_managed = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
monkeypatch.setattr(litellm, "callbacks", [pipeline_managed])
|
||||
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="keeps_native", on_fail="block")])
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
||||
captured_data={
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"_guardrail_pipelines": [("response-governance", pipeline)]},
|
||||
},
|
||||
captured_user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/custom/stream"),
|
||||
captured_logging_obj=_streaming_logging_obj(),
|
||||
assembled_response=ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]),
|
||||
cache_hit=False,
|
||||
)
|
||||
|
||||
assert pipeline_managed.native_hooks_ran == ["post_call"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_guardrails_skip_opted_out_guardrail(monkeypatch):
|
||||
"""The realtime path calls apply_guardrail directly, so the opt-out has to be
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from copy import deepcopy
|
||||
import logging
|
||||
from collections.abc import Iterator
|
||||
from typing import Any, Callable, Dict, List
|
||||
|
|
@ -29,7 +30,7 @@ from litellm.integrations.prometheus import PrometheusLogger
|
|||
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
||||
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines
|
||||
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
|
@ -1723,21 +1724,85 @@ async def _async_chunk_iter(chunks: List[Any]):
|
|||
yield chunk
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
|
||||
def _legacy_hook_stream_guardrail(
|
||||
seen: Dict[str, Any],
|
||||
rewrite: Callable[[Any], Any] | None = None,
|
||||
raises: Exception | None = None,
|
||||
native_lifecycle: bool = False,
|
||||
) -> CustomGuardrail:
|
||||
class LegacyHookGuardrail(CustomGuardrail):
|
||||
use_native_lifecycle_hooks = native_lifecycle
|
||||
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
seen["data"] = data
|
||||
seen["user_api_key_dict"] = user_api_key_dict
|
||||
seen["response"] = deepcopy(response)
|
||||
if raises is not None:
|
||||
raise raises
|
||||
return None if rewrite is None else rewrite(response)
|
||||
|
||||
if native_lifecycle:
|
||||
|
||||
class NativeLifecycleGuardrail(LegacyHookGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
raise AssertionError("a guardrail that keeps its native hooks never runs apply_guardrail")
|
||||
|
||||
return NativeLifecycleGuardrail(
|
||||
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
|
||||
)
|
||||
return LegacyHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
|
||||
|
||||
|
||||
def _iterator_hook_only_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail:
|
||||
class IteratorHookGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
async for item in response:
|
||||
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
|
||||
yield item
|
||||
|
||||
return IteratorHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
|
||||
|
||||
def _iterator_and_legacy_hook_guardrail(name: str, seen: Dict[str, Any]) -> CustomGuardrail:
|
||||
class IteratorAndLegacyHookGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["success_hook_calls"] = seen.get("success_hook_calls", 0) + 1
|
||||
return None
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
seen["iterator_hook_calls"] = seen.get("iterator_hook_calls", 0) + 1
|
||||
async for item in response:
|
||||
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
|
||||
yield item
|
||||
|
||||
return IteratorAndLegacyHookGuardrail(guardrail_name=name, event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
|
||||
|
||||
def _rewritten_model_response(response: Any) -> litellm.ModelResponse:
|
||||
payload = response.model_dump()
|
||||
payload["choices"][0]["message"]["content"] = "[REWRITTEN] " + payload["choices"][0]["message"]["content"]
|
||||
return litellm.ModelResponse(**payload)
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator_only(
|
||||
make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
supported = _unified_stream_guardrail({})
|
||||
native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call)
|
||||
monkeypatch.setattr(litellm, "callbacks", [supported, native_only])
|
||||
governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
legacy = _legacy_hook_stream_guardrail({})
|
||||
legacy.guardrail_name = "gr-legacy"
|
||||
iterator_only = _iterator_hook_only_guardrail("gr-iterator", {})
|
||||
monkeypatch.setattr(litellm, "callbacks", [supported, legacy, iterator_only])
|
||||
governed = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-legacy", on_fail="block")],
|
||||
)
|
||||
ungoverned = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")],
|
||||
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")],
|
||||
)
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")])
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")])
|
||||
data = {
|
||||
"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}
|
||||
}
|
||||
|
|
@ -1746,8 +1811,44 @@ def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
|
|||
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions"))
|
||||
|
||||
assert streamable == (("governed", governed),)
|
||||
assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog))
|
||||
assert not any("'governed'" in message for message in _warnings(caplog))
|
||||
assert any("'ungoverned'" in message and "gr-iterator" in message for message in _warnings(caplog))
|
||||
assert not any("'governed'" in message or "gr-legacy" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_route",
|
||||
["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
|
||||
)
|
||||
def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response(
|
||||
make_user_api_key_auth, monkeypatch, caplog, request_route
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail({})])
|
||||
legacy = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {"metadata": {"_guardrail_pipelines": [("legacy-governance", legacy)]}}
|
||||
auth = make_user_api_key_auth(request_route=request_route)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
streamable = _streamable_post_call_pipelines(data, auth)
|
||||
|
||||
assert streamable == ()
|
||||
assert stream_gated_guardrail_names(data, auth) == frozenset()
|
||||
assert any("'legacy-governance'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_keeps_guardrails_with_their_own_iterator_hook_on_their_own_path(
|
||||
make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", {})])
|
||||
both_hooks = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {"metadata": {"_guardrail_pipelines": [("both-hooks", both_hooks)]}}
|
||||
auth = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
streamable = _streamable_post_call_pipelines(data, auth)
|
||||
|
||||
assert streamable == ()
|
||||
assert stream_gated_guardrail_names(data, auth) == frozenset()
|
||||
assert any("'both-hooks'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def test_streamable_post_call_pipelines_is_empty_on_route_without_translation(
|
||||
|
|
@ -1797,55 +1898,123 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("native_lifecycle", [False, True])
|
||||
async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support(
|
||||
async def test_streaming_iterator_hook_runs_legacy_hook_and_delivers_its_rewrite(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
if native_lifecycle:
|
||||
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
use_native_lifecycle_hooks = True
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
return inputs
|
||||
|
||||
else:
|
||||
|
||||
class NativeOnlyGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
|
||||
)
|
||||
guardrail = _legacy_hook_stream_guardrail(seen, rewrite=_rewritten_model_response, native_lifecycle=native_lifecycle)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
chunks = _stream_chunks()
|
||||
delivered: List[Any] = []
|
||||
auth = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
guardrails_only=True,
|
||||
user_api_key_dict=auth, data=data, call_type="completion", guardrails_only=True
|
||||
)
|
||||
delivered = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
|
||||
)
|
||||
]
|
||||
|
||||
assert out is not None and out.get("stream") is True
|
||||
assert seen["count"] == 1
|
||||
assert isinstance(seen["response"], litellm.ModelResponse)
|
||||
assert seen["response"].choices[0].message.content == "hello world"
|
||||
assert seen["data"]["messages"] == data["messages"]
|
||||
assert seen["user_api_key_dict"] is auth
|
||||
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
|
||||
assert delivered[0].choices[0].delta.content == "[REWRITTEN] hello world"
|
||||
assert delivered[1].choices[0].delta.content in (None, "")
|
||||
assert delivered[1].choices[0].finish_reason == "stop"
|
||||
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
|
||||
assert _warnings(caplog) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_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 seen["count"] == 1
|
||||
assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks]
|
||||
assert [item.choices[0].delta.content for item in delivered] == ["hello ", "world"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_ends_stream_with_legacy_hook_exception(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
blocked = HTTPException(status_code=400, detail={"error": "output blocked"})
|
||||
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, raises=blocked)])
|
||||
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)
|
||||
|
||||
assert out is not None
|
||||
assert out.get("stream") is True
|
||||
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
|
||||
assert len(delivered) == 2
|
||||
assert seen.get("count") is None
|
||||
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await _drain()
|
||||
|
||||
assert seen["count"] == 1
|
||||
assert delivered == []
|
||||
assert info.value is blocked
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_anthropic_sse(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
def rewrite(response: Any) -> Dict[str, Any]:
|
||||
return {**response, "content": [{"type": "text", "text": "[REWRITTEN] " + response["content"][0]["text"]}]}
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=rewrite)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
||||
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/messages"),
|
||||
response=_async_chunk_iter(_anthropic_sse_chunks()),
|
||||
request_data=data,
|
||||
)
|
||||
]
|
||||
|
||||
assert seen["count"] == 1
|
||||
assert seen["response"]["content"][0]["text"] == "hello world"
|
||||
assert seen["response"]["role"] == "assistant"
|
||||
raw = b"".join(delivered).decode()
|
||||
assert "[REWRITTEN] hello world" in raw
|
||||
assert raw.count("event: content_block_delta") == 1
|
||||
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
|
||||
assert f"event: {expected_event}" in raw
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1853,19 +2022,7 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli
|
|||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
class IteratorHookGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
async for item in response:
|
||||
item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}"
|
||||
yield item
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [_iterator_hook_only_guardrail("gr-post", seen)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
||||
|
|
@ -1884,6 +2041,30 @@ async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeli
|
|||
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_iterator_hook_runs_the_iterator_hook_of_a_guardrail_that_also_has_a_post_call_hook(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_iterator_and_legacy_hook_guardrail("gr-post", seen)])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
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(_stream_chunks()),
|
||||
request_data=data,
|
||||
)
|
||||
]
|
||||
|
||||
assert seen == {"iterator_hook_calls": 1}
|
||||
assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"]
|
||||
assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"rewrite_attribute, value",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue