mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
Merge pull request #49482 from BerriAI/litellm_backport_49262_49264_rc_1_108_0
fix(guardrails): backport #49262 and #49264 to rc/1.108.0, name only the pipeline steps that ran in x-litellm-applied-guardrails
This commit is contained in:
yucheng-berri
GitHub
commit
57b539bf99
7 changed files with 2415 additions and 51 deletions
|
|
@ -3,7 +3,7 @@ import copy
|
|||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
|
@ -69,6 +69,8 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
|
|||
# field to suppress a guardrail on the direct-SDK path that never reaches the
|
||||
# proxy's metadata sanitizer.
|
||||
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
|
||||
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
|
||||
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
|
||||
|
||||
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
|
||||
|
||||
|
|
@ -218,6 +220,25 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
|
|||
)
|
||||
|
||||
|
||||
def _str_collection(value: object) -> frozenset[str]:
|
||||
if not isinstance(value, (list, tuple, set, frozenset)):
|
||||
return frozenset()
|
||||
items: Final = cast("Iterable[object]", value) # cast-ok: a metadata collection of guardrail names
|
||||
return frozenset(item for item in items if isinstance(item, str))
|
||||
|
||||
|
||||
def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> bool:
|
||||
if marker is None:
|
||||
return False
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list) and marker in executed:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
|
@ -824,13 +845,16 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
return False
|
||||
|
||||
def _pre_call_marker(self) -> str | None:
|
||||
def _marker(self, token: str) -> str | None:
|
||||
name: Final = self.guardrail_name
|
||||
if not name:
|
||||
return None
|
||||
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
return f"{token}:{name}"
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
|
||||
def _pre_call_marker(self, *, raised: bool = False) -> str | None:
|
||||
return self._marker(_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN)
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None:
|
||||
"""
|
||||
Record that this guardrail's ``async_pre_call_hook`` already ran for this
|
||||
request, so the deployment-level hook does not run it a second time.
|
||||
|
|
@ -839,8 +863,11 @@ class CustomGuardrail(CustomLogger):
|
|||
router later spreads a deployment's model-level ``guardrails`` into the
|
||||
top-level request kwargs, which would otherwise re-trigger the same hook
|
||||
from ``async_pre_call_deployment_hook``.
|
||||
|
||||
With ``raised``, records instead that the hook raised as a policy pipeline
|
||||
step the pipeline then moved past. The deployment-level hook still runs it.
|
||||
"""
|
||||
marker: Final = self._pre_call_marker()
|
||||
marker: Final = self._pre_call_marker(raised=raised)
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
|
|
@ -855,17 +882,57 @@ class CustomGuardrail(CustomLogger):
|
|||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool:
|
||||
marker: Final = self._pre_call_marker()
|
||||
if marker is None:
|
||||
return False
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list) and marker in executed:
|
||||
return True
|
||||
return False
|
||||
def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
|
||||
"""Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
|
||||
return self._pre_call_hook_already_ran(data) or self._pre_call_hook_already_ran(data, raised=True)
|
||||
|
||||
def _pre_call_hook_already_ran(self, data: Mapping[str, object], *, raised: bool = False) -> bool:
|
||||
return _request_carries_marker(data, self._pre_call_marker(raised=raised))
|
||||
|
||||
def deployment_hook_ran(self, data: Mapping[str, object]) -> bool:
|
||||
"""Whether a deployment-level hook ran this guardrail as a policy pipeline step of this request."""
|
||||
return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN))
|
||||
|
||||
def _record_deployment_hook_run(
|
||||
self,
|
||||
data: dict[str, object], # mutable-ok: the request payload the run is recorded on
|
||||
) -> None:
|
||||
"""
|
||||
Record that a deployment-level hook ran this guardrail when it is a policy pipeline step of
|
||||
this request, and name it in the applied guardrails when the request's metadata lists it.
|
||||
"""
|
||||
name: Final = self.guardrail_name
|
||||
marker: Final = self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)
|
||||
buckets: Final = tuple(
|
||||
cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed
|
||||
for bucket in (data.get("metadata"), data.get("litellm_metadata"))
|
||||
if isinstance(bucket, dict)
|
||||
)
|
||||
policy_state: Final = next(
|
||||
(bucket for bucket in buckets if name in _str_collection(bucket.get("_pipeline_managed_guardrails"))),
|
||||
None,
|
||||
)
|
||||
if marker is None or policy_state is None:
|
||||
return
|
||||
recorded: Final = policy_state.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if not isinstance(recorded, list):
|
||||
policy_state[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
elif marker not in recorded:
|
||||
cast("list[object]", recorded).append(marker) # cast-ok: markers are only compared as strings
|
||||
if not any(name in _str_collection(bucket.get("guardrails")) for bucket in buckets):
|
||||
return
|
||||
named_lists: Final = tuple(
|
||||
cast("list[object]", bucket["applied_guardrails"]) # cast-ok: entries are only compared with names
|
||||
for bucket in buckets
|
||||
if isinstance(bucket.get("applied_guardrails"), list)
|
||||
)
|
||||
for applied in named_lists:
|
||||
if name not in applied:
|
||||
applied.append(name)
|
||||
if named_lists:
|
||||
return
|
||||
metadata_bucket: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str keys
|
||||
metadata_bucket["applied_guardrails"] = [name]
|
||||
|
||||
@classmethod
|
||||
def uses_apply_guardrail_interface(cls) -> bool:
|
||||
|
|
@ -914,12 +981,15 @@ class CustomGuardrail(CustomLogger):
|
|||
target: Final = self._deployment_hook_target()
|
||||
if target is not self:
|
||||
kwargs["guardrail_to_apply"] = self
|
||||
result: Final = await target.async_pre_call_hook(
|
||||
user_api_key_dict=_user_api_key_auth_from_request(kwargs),
|
||||
cache=dc,
|
||||
data=kwargs,
|
||||
call_type="completion" if call_type == CallTypes.completion else "acompletion",
|
||||
)
|
||||
try:
|
||||
result: Final = await target.async_pre_call_hook(
|
||||
user_api_key_dict=_user_api_key_auth_from_request(kwargs),
|
||||
cache=dc,
|
||||
data=kwargs,
|
||||
call_type="completion" if call_type == CallTypes.completion else "acompletion",
|
||||
)
|
||||
finally:
|
||||
self._record_deployment_hook_run(kwargs)
|
||||
|
||||
if result is not None and isinstance(result, dict):
|
||||
result_messages: Final = result.get("messages")
|
||||
|
|
@ -1003,6 +1073,9 @@ class CustomGuardrail(CustomLogger):
|
|||
finally:
|
||||
if target is not self:
|
||||
request_data.pop("guardrail_to_apply", None)
|
||||
self._record_deployment_hook_run(
|
||||
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
|
||||
)
|
||||
|
||||
if not self._is_valid_response_type(result):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -325,6 +325,48 @@ def _prepare_hook_input(
|
|||
return hook_input, scans_raw_request
|
||||
|
||||
|
||||
async def _run_pre_call_hook(
|
||||
target: CustomLogger,
|
||||
callback: CustomGuardrail,
|
||||
data: dict[str, object], # mutable-ok: the request payload the run is recorded on
|
||||
hook_input: dict[str, object], # mutable-ok: same request-payload shape the hooks mutate
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
call_type: str,
|
||||
) -> dict[str, object] | None:
|
||||
try:
|
||||
response: Final = await target.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=None,
|
||||
data=hook_input,
|
||||
call_type=call_type,
|
||||
)
|
||||
except Exception:
|
||||
callback.mark_pre_call_hook_ran(data, raised=True)
|
||||
raise
|
||||
callback.mark_pre_call_hook_ran(data)
|
||||
if not isinstance(response, dict):
|
||||
return None
|
||||
payload: Final = cast(dict[str, object], response) # cast-ok: a pre_call hook returns the request payload
|
||||
callback.mark_pre_call_hook_ran(payload)
|
||||
return payload
|
||||
|
||||
|
||||
def _name_step_that_ran(
|
||||
request_data: dict[str, object], # mutable-ok: the request payload the step's name is recorded on
|
||||
metadata_key: Literal["metadata", "litellm_metadata"],
|
||||
guardrail_name: str,
|
||||
) -> None:
|
||||
"""Name a step that ran in the applied guardrails of the metadata bucket the pipeline's request started with."""
|
||||
existing: Final = request_data.get(metadata_key)
|
||||
bucket: Final[dict[str, object]] = ( # mutable-ok: the request's own metadata bucket, written in place
|
||||
cast("dict[str, object]", existing) # cast-ok: metadata buckets are str-keyed
|
||||
if isinstance(existing, dict)
|
||||
else {}
|
||||
)
|
||||
request_data[metadata_key] = bucket # rebind-ok: the step's name is recorded on the request payload
|
||||
add_guardrail_to_applied_guardrails_header(request_data={metadata_key: bucket}, guardrail_name=guardrail_name)
|
||||
|
||||
|
||||
def _release_original_chunks(
|
||||
guardrail_name: str,
|
||||
reason: str,
|
||||
|
|
@ -385,6 +427,9 @@ class PipelineExecutor:
|
|||
PipelineExecutionResult with terminal action and step results
|
||||
"""
|
||||
step_results: Final[list[PipelineStepResult]] = []
|
||||
metadata_key: Final = get_metadata_variable_name_from_kwargs(
|
||||
cast(Mapping[str, object], data) # cast-ok: the request payload is str-keyed
|
||||
)
|
||||
working_data = data.copy()
|
||||
if "metadata" in working_data:
|
||||
working_data["metadata"] = working_data["metadata"].copy()
|
||||
|
|
@ -403,6 +448,7 @@ class PipelineExecutor:
|
|||
data=working_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
metadata_key=metadata_key,
|
||||
raw_request_snapshot=raw_request_snapshot,
|
||||
streaming_chunks=streaming_chunks,
|
||||
endpoint_translation=endpoint_translation,
|
||||
|
|
@ -476,6 +522,7 @@ class PipelineExecutor:
|
|||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
litellm_logging_obj: "LiteLLMLoggingObj | None",
|
||||
stream_already_sent: bool,
|
||||
metadata_key: Literal["metadata", "litellm_metadata"],
|
||||
) -> 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
|
||||
|
|
@ -485,8 +532,8 @@ class PipelineExecutor:
|
|||
live, 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
|
||||
client gets the stream the merge base sent. The guardrail is named in the
|
||||
applied-guardrails header either way, since it ran. 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 = (
|
||||
|
|
@ -521,12 +568,11 @@ class PipelineExecutor:
|
|||
except UndeliverableStreamRewrite as undeliverable:
|
||||
_release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals)
|
||||
return
|
||||
finally:
|
||||
_name_step_that_ran(hook_input, metadata_key, step.guardrail)
|
||||
discard_reason: Final = observer.discard_reason(rewrite_undeliverable_reason)
|
||||
if discard_reason is not None:
|
||||
_release_original_chunks(step.guardrail, discard_reason, 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)
|
||||
|
||||
@staticmethod
|
||||
async def _run_step(
|
||||
|
|
@ -535,6 +581,7 @@ class PipelineExecutor:
|
|||
data: dict,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
call_type: str,
|
||||
metadata_key: Literal["metadata", "litellm_metadata"],
|
||||
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
|
||||
streaming_chunks: list[object] | None = None, # mutable-ok: shared buffered-stream chunks, read per step
|
||||
endpoint_translation: "BaseTranslation | None" = None,
|
||||
|
|
@ -578,16 +625,7 @@ class PipelineExecutor:
|
|||
|
||||
try:
|
||||
if mode == "pre_call":
|
||||
response = await target.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=None,
|
||||
data=hook_input,
|
||||
call_type=call_type,
|
||||
)
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
callback.mark_pre_call_hook_ran(data)
|
||||
if isinstance(response, dict):
|
||||
callback.mark_pre_call_hook_ran(response)
|
||||
response = await _run_pre_call_hook(target, callback, data, hook_input, user_api_key_dict, call_type)
|
||||
elif mode == "post_call" and streaming_chunks is not None:
|
||||
if endpoint_translation is None:
|
||||
return (
|
||||
|
|
@ -605,9 +643,15 @@ class PipelineExecutor:
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_logging_obj=data.get("litellm_logging_obj"),
|
||||
stream_already_sent=stream_already_sent,
|
||||
metadata_key=metadata_key,
|
||||
)
|
||||
response = None
|
||||
elif mode == "post_call":
|
||||
_name_step_that_ran(
|
||||
cast(dict[str, object], data), # cast-ok: the request payload is str-keyed
|
||||
metadata_key,
|
||||
step.guardrail,
|
||||
)
|
||||
response = await target.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -647,6 +647,120 @@ def pipeline_managed_guardrail_names(
|
|||
)
|
||||
|
||||
|
||||
def _ran_or_runs_outside_pipelines(
|
||||
callback: CustomGuardrail | None,
|
||||
data: Mapping[str, object],
|
||||
pipeline_modes: frozenset[str],
|
||||
) -> bool:
|
||||
return callback is not None and (
|
||||
callback.pre_call_hook_ran_or_raised(data)
|
||||
or callback.deployment_hook_ran(data)
|
||||
or any(
|
||||
callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # no public hook check
|
||||
for hook in GuardrailEventHooks
|
||||
if hook.value not in pipeline_modes
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _runs_only_in_pipelines(
|
||||
name: str, data: Mapping[str, object], pre_call: frozenset[str], post_call: frozenset[str]
|
||||
) -> bool:
|
||||
return not _ran_or_runs_outside_pipelines(
|
||||
PipelineExecutor.find_guardrail_callback(name),
|
||||
data,
|
||||
frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names),
|
||||
)
|
||||
|
||||
|
||||
def _request_steps_that_have_not_run(data: Mapping[str, object]) -> frozenset[str]:
|
||||
"""Pre_call policy pipeline steps a pipeline skipped, whose guardrail nothing else ran or runs for this request."""
|
||||
pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
return frozenset(name for name in pre_call if _runs_only_in_pipelines(name, data, pre_call, frozenset()))
|
||||
|
||||
|
||||
def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]:
|
||||
return frozenset(step.guardrail_name for step in result.step_results if step.outcome != "skip")
|
||||
|
||||
|
||||
def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool:
|
||||
"""
|
||||
The success log is parked until the response hooks finish, and the headers are built
|
||||
after them. The LLM pass-through route builds its headers first, and the generateContent
|
||||
adapter logs its nested completion before the response hooks run.
|
||||
"""
|
||||
logging_obj: Final = data.get("litellm_logging_obj")
|
||||
return (
|
||||
call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value)
|
||||
and isinstance(logging_obj, Logging)
|
||||
and logging_obj.enqueue_deferred_logging is not None
|
||||
)
|
||||
|
||||
|
||||
def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool:
|
||||
return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None
|
||||
|
||||
|
||||
def _pipelines_settled_by(
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
|
||||
) -> Sequence[tuple[str, "GuardrailPipeline"]]:
|
||||
"""This pipeline, plus every later response pipeline when its verdict ends the request before they run."""
|
||||
if result.terminal_action == "allow":
|
||||
return pipelines[index : index + 1]
|
||||
return (
|
||||
pipelines[index],
|
||||
*((policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"),
|
||||
)
|
||||
|
||||
|
||||
def _steps_later_pipelines_still_run(
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
|
||||
) -> frozenset[str]:
|
||||
"""The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on."""
|
||||
if result.terminal_action != "allow":
|
||||
return frozenset()
|
||||
later: Final = tuple(
|
||||
(policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"
|
||||
)
|
||||
return _pipeline_step_guardrail_names(later)
|
||||
|
||||
|
||||
def _drop_response_steps_that_did_not_run(
|
||||
data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]],
|
||||
steps_to_keep: frozenset[str],
|
||||
) -> None:
|
||||
"""
|
||||
Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``,
|
||||
unless their guardrail also runs outside pipelines. Every step is named up front, so a
|
||||
log built before the pipeline finished, such as a cache hit's, still lists the steps
|
||||
that ran. The list is trimmed in place, since a route's logging object may hold a
|
||||
shallow copy of the metadata.
|
||||
"""
|
||||
pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
post_call: Final = _pipeline_step_guardrail_names(pipelines)
|
||||
not_run: Final = frozenset(
|
||||
name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call)
|
||||
)
|
||||
if not not_run:
|
||||
return
|
||||
for bucket_key in ("metadata", "litellm_metadata"):
|
||||
_drop_applied_guardrails_in_place(data.get(bucket_key), not_run)
|
||||
|
||||
|
||||
def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None:
|
||||
if not isinstance(bucket, dict):
|
||||
return
|
||||
metadata: Final = cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed
|
||||
applied: Final = metadata.get("applied_guardrails")
|
||||
if not isinstance(applied, list):
|
||||
return
|
||||
entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names
|
||||
entries[:] = [name for name in entries if name not in names]
|
||||
if not entries:
|
||||
metadata.pop("applied_guardrails")
|
||||
|
||||
|
||||
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
|
||||
resolved: Final = tuple(
|
||||
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
|
|
@ -2571,8 +2685,12 @@ class ProxyLogging:
|
|||
data.pop("prompt_version", None)
|
||||
data.pop("prompt_environment", None)
|
||||
|
||||
def _process_guardrail_metadata(self, data: dict) -> None:
|
||||
"""Process guardrails from metadata and add to applied_guardrails."""
|
||||
def _process_guardrail_metadata(self, data: dict, *, skip_guardrails: bool = False) -> None:
|
||||
"""
|
||||
Process guardrails from metadata and add to applied_guardrails, leaving out the pre_call
|
||||
policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on
|
||||
another request, so every step is added.
|
||||
"""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
|
@ -2596,8 +2714,17 @@ class ProxyLogging:
|
|||
if not isinstance(applied_guardrails, list):
|
||||
applied_guardrails = []
|
||||
|
||||
not_run: Final = (
|
||||
frozenset[str]()
|
||||
if skip_guardrails
|
||||
else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload
|
||||
)
|
||||
for guardrail_name in guardrails_in_metadata:
|
||||
if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails:
|
||||
if (
|
||||
isinstance(guardrail_name, str)
|
||||
and guardrail_name not in applied_guardrails
|
||||
and guardrail_name not in not_run
|
||||
):
|
||||
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name)
|
||||
|
||||
async def _maybe_execute_pipelines(
|
||||
|
|
@ -2629,7 +2756,10 @@ class ProxyLogging:
|
|||
return data, None
|
||||
|
||||
current_response = response # rebind-ok: chains each pipeline's replacement response into the next
|
||||
for policy_name, pipeline in pipelines:
|
||||
steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran
|
||||
request: Final = cast(dict[str, object], data) # cast-ok: request payload, updated in place by every pipeline
|
||||
log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type)
|
||||
for index, (policy_name, pipeline) in enumerate(pipelines):
|
||||
if pipeline.mode != event_hook:
|
||||
continue
|
||||
|
||||
|
|
@ -2645,6 +2775,14 @@ class ProxyLogging:
|
|||
raw_request_snapshot=raw_request_snapshot,
|
||||
)
|
||||
|
||||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if log_waits:
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
|
||||
data = self._handle_pipeline_result(
|
||||
result=result,
|
||||
data=data,
|
||||
|
|
@ -2884,7 +3022,7 @@ class ProxyLogging:
|
|||
and (guardrails_only or not caps.has_pre_call_override)
|
||||
):
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
return data
|
||||
|
||||
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
|
||||
|
|
@ -2988,17 +3126,17 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
|
||||
return data
|
||||
except SensitiveDataRouteException as e:
|
||||
data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict)
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
return data
|
||||
except Exception:
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
raise
|
||||
|
||||
async def _run_parallel_pre_call_guardrails(
|
||||
|
|
@ -4586,7 +4724,8 @@ class ProxyLogging:
|
|||
|
||||
call_type, endpoint_translation = translation
|
||||
|
||||
for policy_name, pipeline in pipelines:
|
||||
steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran
|
||||
for index, (policy_name, pipeline) in enumerate(pipelines):
|
||||
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode="post_call",
|
||||
|
|
@ -4597,6 +4736,13 @@ class ProxyLogging:
|
|||
streaming_chunks=buffered,
|
||||
endpoint_translation=endpoint_translation,
|
||||
)
|
||||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if _stream_log_waits_for_its_pipelines(request_data):
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request_data,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
try:
|
||||
ProxyLogging._handle_pipeline_result(
|
||||
result, data=request_data, policy_name=policy_name, original_response=buffered
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -15,6 +15,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
|
||||
|
|
@ -30,7 +31,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
|
|||
GuardrailPipeline,
|
||||
PipelineStep,
|
||||
)
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs
|
||||
|
||||
try:
|
||||
from fastapi.exceptions import HTTPException
|
||||
|
|
@ -1157,6 +1158,7 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
|
|||
data=data,
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
metadata_key="metadata",
|
||||
)
|
||||
|
||||
assert outcome == "pass"
|
||||
|
|
@ -1270,7 +1272,7 @@ def _assert_passed_with_discard_warning(result, caplog, reason):
|
|||
assert len(discard_warnings) == 1
|
||||
assert reason in discard_warnings[0]
|
||||
assert "text rewrites included" in discard_warnings[0]
|
||||
assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", [])
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1433,6 +1435,60 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he
|
|||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
|
||||
|
||||
class _SelfRecordingGuardrail(_TextReturningGuardrail):
|
||||
records_own_guardrail_information = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])])
|
||||
|
||||
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
|
||||
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
|
||||
|
||||
class _LitellmMetadataOpeningGuardrail(CustomGuardrail):
|
||||
"""Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
request_data["litellm_metadata"] = {"user_api_key_hash": "hashed"}
|
||||
return inputs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_is_named_in_the_metadata_the_request_started_with(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [_LitellmMetadataOpeningGuardrail()])
|
||||
|
||||
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
|
||||
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
assert "applied_guardrails" not in result.modified_data["litellm_metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()])
|
||||
|
||||
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
|
||||
|
||||
assert [(step.outcome, step.action_taken) for step in result.step_results] == [("fail", "next")]
|
||||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()])
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Iterator
|
||||
from contextlib import AbstractContextManager, nullcontext
|
||||
from copy import deepcopy
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, Final, List
|
||||
|
|
@ -28,14 +29,20 @@ from litellm.integrations.custom_guardrail import (
|
|||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
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.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
get_logging_caching_headers,
|
||||
)
|
||||
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
|
||||
from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_information
|
||||
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.caching.dual_cache import DualCache
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import (
|
||||
GuardrailPipeline,
|
||||
|
|
@ -1160,6 +1167,831 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe
|
|||
assert seen["count"] == 1
|
||||
|
||||
|
||||
class _PipelineStepGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str, event_hook: str | list[str]) -> None:
|
||||
super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False)
|
||||
self.calls = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
|
||||
) -> dict[str, object]:
|
||||
self.calls += 1
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
) -> None:
|
||||
self.calls += 1
|
||||
|
||||
|
||||
def _two_step_pipeline_data(
|
||||
mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block"
|
||||
) -> dict[str, object]:
|
||||
pipeline: Final = GuardrailPipeline(
|
||||
mode=mode,
|
||||
steps=[
|
||||
PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail),
|
||||
PipelineStep(guardrail="second", on_pass="next", on_fail=second_on_fail),
|
||||
],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [("governance", pipeline)],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _logging_obj(*, log_parked: bool, call_type: str = "acompletion") -> LiteLLMLoggingObj:
|
||||
logging_obj: Final = LiteLLMLoggingObj(
|
||||
model="m",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type=call_type,
|
||||
start_time=None,
|
||||
litellm_call_id="call-1",
|
||||
function_id="fn-1",
|
||||
)
|
||||
logging_obj.defer_async_logging = True
|
||||
logging_obj.enqueue_deferred_logging = (lambda: None) if log_parked else None
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("first_on_pass", "expected_calls", "expected_header"),
|
||||
[("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])],
|
||||
)
|
||||
async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
first_on_pass: str,
|
||||
expected_calls: tuple[int, int],
|
||||
expected_header: list[str],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", first_on_pass),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == expected_calls
|
||||
assert result["metadata"]["applied_guardrails"] == expected_header
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert result["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail):
|
||||
def __init__(self, guardrail_name: str, error: Exception) -> None:
|
||||
super().__init__(guardrail_name, "pre_call")
|
||||
self.error = error
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
|
||||
) -> dict[str, object]:
|
||||
self.calls += 1
|
||||
raise self.error
|
||||
|
||||
|
||||
_PIPELINE_STEP_ERRORS: Final = (
|
||||
pytest.param(HTTPException(status_code=400, detail="flagged"), id="flagged"),
|
||||
pytest.param(RuntimeError("vendor unreachable"), id="errored"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS)
|
||||
async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
error: Exception,
|
||||
) -> None:
|
||||
first: Final = _RaisingPipelineStepGuardrail("first", error)
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert result["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS)
|
||||
async def test_deployment_hook_still_runs_a_pipeline_step_that_raised(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
error: Exception,
|
||||
) -> None:
|
||||
first: Final = _RaisingPipelineStepGuardrail("first", error)
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(type(error)):
|
||||
await first.async_pre_call_deployment_hook(
|
||||
kwargs={**result, "guardrails": ["first"]}, call_type=CallTypes.completion
|
||||
)
|
||||
await second.async_pre_call_deployment_hook(
|
||||
kwargs={**result, "guardrails": ["second"]}, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (2, 1)
|
||||
|
||||
|
||||
class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail):
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
|
||||
) -> dict[str, object]:
|
||||
self.calls += 1
|
||||
data["litellm_metadata"] = {}
|
||||
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
|
||||
return data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _OpensLitellmMetadataGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "next"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert "applied_guardrails" not in result["metadata"]
|
||||
assert result["litellm_metadata"]["applied_guardrails"] == ["second", "first"]
|
||||
assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("log_parked", "call_type", "bucket", "expected_after_response"),
|
||||
[
|
||||
pytest.param(True, "acompletion", "metadata", ["first"], id="success-log-waits-for-the-response-hooks"),
|
||||
pytest.param(True, "acompletion", "litellm_metadata", ["first"], id="success-log-waits-litellm-metadata"),
|
||||
pytest.param(False, "acompletion", "metadata", ["first", "second"], id="success-log-already-built"),
|
||||
pytest.param(
|
||||
True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out"
|
||||
),
|
||||
pytest.param(
|
||||
True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call"
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
log_parked: bool,
|
||||
call_type: str,
|
||||
bucket: str,
|
||||
expected_after_response: list[str],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
request: Final = _two_step_pipeline_data("post_call", "allow")
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": request["model"],
|
||||
"messages": request["messages"],
|
||||
bucket: request["metadata"],
|
||||
"litellm_logging_obj": _logging_obj(log_parked=log_parked, call_type=call_type),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
assert data[bucket]["applied_guardrails"] == ["first", "second"]
|
||||
logged_metadata: Final = data[bucket].copy()
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data[bucket]["applied_guardrails"] == expected_after_response
|
||||
assert logged_metadata["applied_guardrails"] == expected_after_response
|
||||
|
||||
|
||||
def _three_step_two_pipeline_data() -> dict[str, object]:
|
||||
runs_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="first", on_pass="next", on_fail="block"),
|
||||
PipelineStep(guardrail="second", on_pass="next", on_fail="block"),
|
||||
],
|
||||
)
|
||||
stops_before_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="third", on_pass="allow", on_fail="block"),
|
||||
PipelineStep(guardrail="second", on_pass="next", on_fail="block"),
|
||||
],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second", "third"],
|
||||
"_guardrail_pipelines": [("runs-second", runs_second), ("stops-before-second", stops_before_second)],
|
||||
"_pipeline_managed_guardrails": {"first", "second", "third"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
**_three_step_two_pipeline_data(),
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1]
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={
|
||||
**_three_step_two_pipeline_data(),
|
||||
"stream": True,
|
||||
"litellm_logging_obj": logging_obj,
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
delivered: Final[list[object]] = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
)
|
||||
]
|
||||
|
||||
assert len(delivered) == 2
|
||||
assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1]
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"]
|
||||
|
||||
|
||||
def _ends_before_a_later_pipeline_data(ending: str) -> dict[str, object]:
|
||||
ends: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass=ending)])
|
||||
never_reached: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="second")])
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [("ends", ends), ("never-reached", never_reached)],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("ending", "raised"), [("block", HTTPException), ("modify_response", ModifyResponseException)])
|
||||
async def test_post_call_pipeline_that_ends_the_request_drops_the_steps_of_a_later_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
ending: str,
|
||||
raised: type[Exception],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_ends_before_a_later_pipeline_data(ending), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(raised):
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first"]
|
||||
|
||||
|
||||
def _blocks_before_a_later_request_pipeline_data() -> dict[str, object]:
|
||||
blocks: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass="block")])
|
||||
request: Final = GuardrailPipeline(
|
||||
mode="pre_call",
|
||||
steps=[PipelineStep(guardrail="gate", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "gate", "second"],
|
||||
"_guardrail_pipelines": [("blocks", blocks), ("request", request)],
|
||||
"_pipeline_managed_guardrails": {"first", "gate", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_that_blocks_keeps_a_later_request_step_whose_guardrail_runs_outside_pipelines(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
gate: Final = _PipelineStepGuardrail("gate", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, gate, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_blocks_before_a_later_request_pipeline_data(), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert sorted(data["metadata"]["applied_guardrails"]) == ["first", "gate", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("ending", "stream_end"), [("block", pytest.raises(HTTPException)), ("modify_response", nullcontext())]
|
||||
)
|
||||
async def test_held_stream_pipeline_that_ends_the_stream_drops_the_steps_of_a_later_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
ending: str,
|
||||
stream_end: AbstractContextManager[object],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={**_ends_before_a_later_pipeline_data(ending), "stream": True, "litellm_logging_obj": logging_obj},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with stream_end:
|
||||
async for _item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
):
|
||||
pass
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("steps", "expected"),
|
||||
[(("first", "second"), ["second"]), (("first",), None)],
|
||||
ids=["skipped-then-ran", "only-step-skipped"],
|
||||
)
|
||||
async def test_post_call_pipeline_drops_a_step_its_guardrail_skips_on_a_whole_response(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
steps: tuple[str, ...],
|
||||
expected: list[str] | None,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in steps]
|
||||
guardrails[0].apply_stream_scope("streaming")
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
pipeline: Final = GuardrailPipeline(
|
||||
mode="post_call", steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="block") for name in steps]
|
||||
)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": list(steps),
|
||||
"_guardrail_pipelines": [("governance", pipeline)],
|
||||
"_pipeline_managed_guardrails": set(steps),
|
||||
},
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert [guardrail.calls for guardrail in guardrails] == [0, 1][: len(steps)]
|
||||
assert data["metadata"].get("applied_guardrails") == expected
|
||||
|
||||
|
||||
class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
) -> None:
|
||||
self.calls += 1
|
||||
data["litellm_metadata"] = {"user_api_key_hash": "hashed"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "next"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
stops_before_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")],
|
||||
)
|
||||
stops_before_first: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")],
|
||||
)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [
|
||||
("stops-before-second", stops_before_second),
|
||||
("stops-before-first", stops_before_first),
|
||||
],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
skip_guardrails=True,
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (0, 0)
|
||||
assert result["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_names_a_skipped_pre_call_step_it_ran(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
await second.async_pre_call_deployment_hook(
|
||||
kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
assert second.deployment_hook_ran(data)
|
||||
assert not first.deployment_hook_ran(data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS)
|
||||
async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
error: Exception,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _RaisingPipelineStepGuardrail("second", error)
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(type(error)):
|
||||
await second.async_pre_call_deployment_hook(
|
||||
kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_skipped_step_a_deployment_hook_ran(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
await second.async_post_call_success_deployment_hook(
|
||||
request_data={**data, "guardrails": ["second"]},
|
||||
response=litellm.ModelResponse(),
|
||||
call_type=CallTypes.acompletion,
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_leaves_a_request_without_policy_pipelines_untouched() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": {"guardrails": ["second"]},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert kwargs["metadata"] == {"guardrails": ["second"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_names_a_pipeline_step_it_ran_before_any_guardrail_was_named() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
metadata: Final[dict[str, object]] = {
|
||||
"guardrails": ["second"],
|
||||
"_pipeline_managed_guardrails": frozenset({"second"}),
|
||||
}
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert metadata["applied_guardrails"] == ["second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_names_a_pipeline_step_requested_only_in_the_other_metadata_bucket() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": {
|
||||
"guardrails": ["first"],
|
||||
"applied_guardrails": ["first"],
|
||||
"_pipeline_managed_guardrails": frozenset({"first", "second"}),
|
||||
},
|
||||
"litellm_metadata": {"guardrails": ["second"]},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert sorted(get_logging_caching_headers(kwargs)["x-litellm-applied-guardrails"].split(",")) == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("named_bucket", ["metadata", "litellm_metadata"])
|
||||
async def test_deployment_hook_names_a_pipeline_step_next_to_the_steps_already_named_when_both_buckets_exist(
|
||||
named_bucket: str,
|
||||
) -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
metadata: Final[dict[str, object]] = {
|
||||
"guardrails": ["first", "second"],
|
||||
"_pipeline_managed_guardrails": frozenset({"first", "second"}),
|
||||
}
|
||||
litellm_metadata: Final[dict[str, object]] = {}
|
||||
{"metadata": metadata, "litellm_metadata": litellm_metadata}[named_bucket]["applied_guardrails"] = ["first"]
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": metadata,
|
||||
"litellm_metadata": litellm_metadata,
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert sorted(get_logging_caching_headers(kwargs)["x-litellm-applied-guardrails"].split(",")) == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one_the_request_did_not_list() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": {"_pipeline_managed_guardrails": frozenset({"second"})},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.deployment_hook_ran(kwargs)
|
||||
assert not _PipelineStepGuardrail("first", "pre_call").deployment_hook_ran(kwargs)
|
||||
assert "applied_guardrails" not in kwargs["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("log_waits", "expected_after_stream"),
|
||||
[(True, ["first"]), (False, ["first", "second"])],
|
||||
ids=["stream-log-waits-for-the-pipeline", "stream-log-not-deferred"],
|
||||
)
|
||||
async def test_held_stream_pipeline_drops_the_steps_that_did_not_run_only_while_its_log_waits(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
log_waits: bool,
|
||||
expected_after_stream: list[str],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log if log_waits else None
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "stream": True, "litellm_logging_obj": logging_obj},
|
||||
call_type="completion",
|
||||
)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
logged_metadata: Final = data["metadata"].copy()
|
||||
|
||||
delivered: Final[list[object]] = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
)
|
||||
]
|
||||
|
||||
assert len(delivered) == 2
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == expected_after_stream
|
||||
assert logged_metadata["applied_guardrails"] == expected_after_stream
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_replacement_response_reaches_caller(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
|
|
@ -170,8 +170,9 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
|
|||
proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
|
||||
invoked = {}
|
||||
|
||||
def fake_process(d):
|
||||
def fake_process(d: object, *, skip_guardrails: bool) -> None:
|
||||
invoked["data"] = d
|
||||
invoked["skip_guardrails"] = skip_guardrails
|
||||
|
||||
proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
|
|
@ -180,7 +181,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
|
|||
call_type="completion",
|
||||
)
|
||||
assert out is data
|
||||
assert invoked["data"] is data
|
||||
assert invoked == {"data": data, "skip_guardrails": False}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue