mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(guardrails): name only the pipeline steps that ran in x-litellm-applied-guardrails
The pre_call header sweep claimed every policy pipeline step up front, so steps a pipeline never reached showed up in x-litellm-applied-guardrails and in spend logs. The sweep now leaves pipeline steps to the executor, which names each step when its hook runs, whatever it decides. Steps a stream runs after its headers go out (a pipeline the stream cannot hold, or one whose steps all stream live) are still named up front, as before. Resolves LIT-10096
This commit is contained in:
yucheng-berri
parent
1c022d1d06
commit
4fcddb0bbd
7 changed files with 693 additions and 44 deletions
|
|
@ -69,6 +69,7 @@ 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)
|
||||
|
||||
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
|
||||
|
||||
|
|
@ -824,13 +825,13 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
return False
|
||||
|
||||
def _pre_call_marker(self) -> str | None:
|
||||
def _pre_call_marker(self, *, raised: bool = False) -> str | None:
|
||||
name: Final = self.guardrail_name
|
||||
if not name:
|
||||
return None
|
||||
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
return f"{_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
|
||||
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 +840,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,8 +859,12 @@ 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()
|
||||
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:
|
||||
marker: Final = self._pre_call_marker(raised=raised)
|
||||
if marker is None:
|
||||
return False
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
|
|
|
|||
|
|
@ -329,6 +329,32 @@ 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 _release_original_chunks(
|
||||
guardrail_name: str,
|
||||
reason: str,
|
||||
|
|
@ -489,8 +515,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 = (
|
||||
|
|
@ -525,12 +551,12 @@ class PipelineExecutor:
|
|||
except UndeliverableStreamRewrite as undeliverable:
|
||||
_release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals)
|
||||
return
|
||||
finally:
|
||||
if not callback.records_own_guardrail_information:
|
||||
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=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(
|
||||
|
|
@ -582,16 +608,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 (
|
||||
|
|
@ -612,6 +629,7 @@ class PipelineExecutor:
|
|||
)
|
||||
response = None
|
||||
elif mode == "post_call":
|
||||
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
|
||||
response = await target.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -706,6 +706,62 @@ 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 any(
|
||||
callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks
|
||||
for hook in GuardrailEventHooks
|
||||
if hook.value not in pipeline_modes
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# The A2A JSON-RPC methods answered with a stream, as both protocol versions spell them
|
||||
_A2A_STREAMING_METHODS: Final = frozenset(
|
||||
("message/stream", "tasks/resubscribe", "SendStreamingMessage", "SubscribeToTask")
|
||||
)
|
||||
|
||||
|
||||
def _post_call_steps_named_before_headers(
|
||||
data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> frozenset[str]:
|
||||
"""
|
||||
Post_call pipeline steps the executor names, as it runs them, before the response
|
||||
headers go out: every step of a whole response, and of a streamed one the steps of
|
||||
the pipelines that hold the stream until they finish.
|
||||
"""
|
||||
if data.get("stream") is not True and data.get("method") not in _A2A_STREAMING_METHODS:
|
||||
return pipeline_managed_guardrail_names(data, "post_call")
|
||||
gated: Final = _stream_gated_pipelines(data, user_api_key_dict)
|
||||
return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated)
|
||||
|
||||
|
||||
def _pipeline_steps_that_have_not_run(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]:
|
||||
"""
|
||||
Guardrails only policy pipelines run whose pre_call hook has not run for this request:
|
||||
pre_call steps a pipeline skipped and the post_call steps the executor names when it
|
||||
runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one
|
||||
whose steps all stream live after the headers go out, so those steps are named up front
|
||||
like any other guardrail.
|
||||
"""
|
||||
pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict)
|
||||
return frozenset(
|
||||
name
|
||||
for name in pre_call | post_call
|
||||
if 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 _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
|
||||
resolved: Final = tuple(
|
||||
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
|
|
@ -927,19 +983,23 @@ def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseT
|
|||
return None if resolved is None else resolved[1]
|
||||
|
||||
|
||||
def _stream_gated_pipelines(
|
||||
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
|
||||
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
|
||||
if translation is None:
|
||||
return ()
|
||||
return tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in _post_call_pipelines(request_data)
|
||||
if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
||||
)
|
||||
|
||||
|
||||
def stream_gated_guardrail_names(
|
||||
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> frozenset[str]:
|
||||
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 not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
||||
)
|
||||
)
|
||||
return _pipeline_step_guardrail_names(_stream_gated_pipelines(request_data, user_api_key_dict))
|
||||
|
||||
|
||||
def _streamable_post_call_pipelines(
|
||||
|
|
@ -2639,7 +2699,7 @@ class ProxyLogging:
|
|||
data.pop("prompt_version", None)
|
||||
data.pop("prompt_environment", None)
|
||||
|
||||
def _process_guardrail_metadata(self, data: dict) -> None:
|
||||
def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
"""Process guardrails from metadata and add to applied_guardrails."""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
|
|
@ -2664,8 +2724,16 @@ class ProxyLogging:
|
|||
if not isinstance(applied_guardrails, list):
|
||||
applied_guardrails = []
|
||||
|
||||
not_run: Final = _pipeline_steps_that_have_not_run(
|
||||
cast(dict[str, object], data), # cast-ok: request payload
|
||||
user_api_key_dict,
|
||||
)
|
||||
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(
|
||||
|
|
@ -2952,7 +3020,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, user_api_key_dict)
|
||||
return data
|
||||
|
||||
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
|
||||
|
|
@ -3056,17 +3124,17 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, user_api_key_dict)
|
||||
|
||||
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, user_api_key_dict)
|
||||
return data
|
||||
except Exception:
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, user_api_key_dict)
|
||||
raise
|
||||
|
||||
async def _run_parallel_pre_call_guardrails(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,314 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from tests.integration._support.client import Gateway, gateway_from_environment
|
||||
from tests.integration._support.process import owned_proxy
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
_HEADER: Final = "x-litellm-applied-guardrails"
|
||||
_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api"
|
||||
_PROMPT_SHIELD_PATH: Final = "/contentsafety/text:shieldPrompt"
|
||||
_PROMPT_SHIELD: Final = "prompt-shield-step"
|
||||
_REQUEST_FIRST: Final = "request-first-step"
|
||||
_REQUEST_SECOND: Final = "request-second-step"
|
||||
_REQUEST_BLOCKER: Final = "request-blocking-step"
|
||||
_RESPONSE_FIRST: Final = "response-first-step"
|
||||
_RESPONSE_SECOND: Final = "response-second-step"
|
||||
_RESPONSE_BLOCKER: Final = "response-blocking-step"
|
||||
_LIVE_FIRST: Final = "live-first-step"
|
||||
_LIVE_SECOND: Final = "live-second-step"
|
||||
_UNGATED_STEP: Final = "ungated-step"
|
||||
_STOPS_AFTER_FIRST: Final = "stops-after-first"
|
||||
_RUNS_EVERY_STEP: Final = "runs-every-step"
|
||||
_BLOCKS_AT_FIRST: Final = "blocks-at-first"
|
||||
_RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first"
|
||||
_SHIELDS_THEN_CHECKS: Final = "shields-then-checks"
|
||||
_MOVES_PAST_FIRST: Final = "moves-past-first"
|
||||
_RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first"
|
||||
_SCANS_AFTER_RELEASE: Final = "scans-after-release"
|
||||
_CANNOT_HOLD_STREAM: Final = "cannot-hold-stream"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
gateway: Gateway
|
||||
vendor: Wire
|
||||
upstream: Wire
|
||||
model: str
|
||||
|
||||
|
||||
def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes:
|
||||
chunk: Final = {
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
return f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
|
||||
|
||||
def _completion(request: Request) -> Reply:
|
||||
assert request.target == "/v1/chat/completions", request.target
|
||||
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
if json.loads(request.body).get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_completion_chunk(identity, {"role": "assistant", "content": "guarded"}, None),
|
||||
_completion_chunk(identity, {}, "stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _vendor(request: Request) -> Reply:
|
||||
if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"):
|
||||
return Reply(
|
||||
body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode()
|
||||
)
|
||||
if request.target.startswith(f"/{_UNGATED_STEP}/"):
|
||||
return Reply(body=json.dumps({"finalAction": "NONE", "appliedPolicies": []}).encode())
|
||||
assert request.target.endswith(_VENDOR_PATH), request.target
|
||||
if request.target.startswith((f"/{_REQUEST_BLOCKER}/", f"/{_RESPONSE_BLOCKER}/")):
|
||||
return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "first step refused"}).encode())
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
|
||||
def _guardrail(name: str, mode: str, vendor: Wire, **params: object) -> dict[str, object]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": f"{vendor.url}/{name}",
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
**params,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _policy(mode: str, steps: tuple[tuple[str, str, str], ...]) -> dict[str, object]:
|
||||
return {
|
||||
"guardrails": {"add": [guardrail for guardrail, _, _ in steps]},
|
||||
"pipeline": {
|
||||
"mode": mode,
|
||||
"steps": [
|
||||
{"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail}
|
||||
for guardrail, on_pass, on_fail in steps
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _config(directory: Path, vendor: Wire) -> Path:
|
||||
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
configuration["guardrails"] = [
|
||||
*(_guardrail(name, "pre_call", vendor) for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER)),
|
||||
*(_guardrail(name, "post_call", vendor) for name in (_RESPONSE_FIRST, _RESPONSE_SECOND, _RESPONSE_BLOCKER)),
|
||||
*(
|
||||
_guardrail(name, "post_call", vendor, streaming_buffer_until_moderated=False)
|
||||
for name in (_LIVE_FIRST, _LIVE_SECOND)
|
||||
),
|
||||
{
|
||||
"guardrail_name": _UNGATED_STEP,
|
||||
"litellm_params": {
|
||||
"guardrail": "dynamoai",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"api_base": f"{vendor.url}/{_UNGATED_STEP}",
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": _PROMPT_SHIELD,
|
||||
"litellm_params": {
|
||||
"guardrail": "azure/prompt_shield",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": f"{vendor.url}/{_PROMPT_SHIELD}",
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
},
|
||||
]
|
||||
configuration["policies"] = {
|
||||
_STOPS_AFTER_FIRST: _policy(
|
||||
"pre_call", ((_REQUEST_FIRST, "allow", "block"), (_REQUEST_SECOND, "next", "block"))
|
||||
),
|
||||
_RUNS_EVERY_STEP: _policy("pre_call", ((_REQUEST_FIRST, "next", "block"), (_REQUEST_SECOND, "next", "block"))),
|
||||
_BLOCKS_AT_FIRST: _policy(
|
||||
"pre_call", ((_REQUEST_BLOCKER, "next", "block"), (_REQUEST_SECOND, "next", "block"))
|
||||
),
|
||||
_RESPONSE_STOPS_AFTER_FIRST: _policy(
|
||||
"post_call", ((_RESPONSE_FIRST, "allow", "block"), (_RESPONSE_SECOND, "next", "block"))
|
||||
),
|
||||
_SHIELDS_THEN_CHECKS: _policy(
|
||||
"pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block"))
|
||||
),
|
||||
_MOVES_PAST_FIRST: _policy(
|
||||
"pre_call", ((_REQUEST_BLOCKER, "next", "next"), (_REQUEST_SECOND, "next", "block"))
|
||||
),
|
||||
_RESPONSE_MOVES_PAST_FIRST: _policy(
|
||||
"post_call", ((_RESPONSE_BLOCKER, "next", "next"), (_RESPONSE_SECOND, "next", "block"))
|
||||
),
|
||||
_SCANS_AFTER_RELEASE: _policy("post_call", ((_LIVE_FIRST, "next", "next"), (_LIVE_SECOND, "next", "next"))),
|
||||
_CANNOT_HOLD_STREAM: _policy(
|
||||
"post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block"))
|
||||
),
|
||||
}
|
||||
configuration["policy_attachments"] = [{"policy": policy, "tags": [policy]} for policy in configuration["policies"]]
|
||||
path: Final = directory / "guardrail-pipeline-applied-header.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("guardrail-pipeline-applied-header")
|
||||
with (
|
||||
wire_server(_vendor) as vendor,
|
||||
wire_server(_completion) as upstream,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(shared, directory, {}, config=_config(directory, vendor)) as owned,
|
||||
owned.scenario() as scenario,
|
||||
):
|
||||
yield _Rig(owned, vendor, upstream, scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture"))
|
||||
|
||||
|
||||
def _ask_as_team_tagged(rig: _Rig, tag: str, *, stream: bool = False) -> httpx.Response:
|
||||
rig.vendor.drain()
|
||||
rig.upstream.drain()
|
||||
with rig.gateway.scenario() as scenario:
|
||||
team: Final = scenario.team(metadata={"tags": [tag]})
|
||||
key: Final = scenario.key(team_id=team)
|
||||
return rig.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": rig.model,
|
||||
"messages": [{"role": "user", "content": f"hello {uuid.uuid4().hex}"}],
|
||||
"stream": stream,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _steps_that_ran(rig: _Rig) -> tuple[str, ...]:
|
||||
return tuple(request.target.split("/")[1] for request in rig.vendor.drain())
|
||||
|
||||
|
||||
def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_REQUEST_FIRST,)
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_streamed_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=True)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.rstrip().endswith("data: [DONE]"), response.text
|
||||
assert _steps_that_ran(rig) == (_REQUEST_FIRST,)
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_request_pipeline_that_runs_every_step_names_every_step_in_order(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _RUNS_EVERY_STEP)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_REQUEST_FIRST, _REQUEST_SECOND)
|
||||
assert response.headers.get(_HEADER) == f"{_REQUEST_FIRST},{_REQUEST_SECOND}", dict(response.headers)
|
||||
|
||||
|
||||
def test_a_request_pipeline_names_a_prompt_shield_step_that_ran_before_a_generic_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _SHIELDS_THEN_CHECKS)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_PROMPT_SHIELD, _REQUEST_FIRST)
|
||||
assert set(response.headers.get(_HEADER, "").split(",")) == {_PROMPT_SHIELD, _REQUEST_FIRST}, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _BLOCKS_AT_FIRST)
|
||||
assert response.status_code == 400, response.text
|
||||
assert _steps_that_ran(rig) == (_REQUEST_BLOCKER,)
|
||||
assert rig.upstream.drain() == ()
|
||||
assert response.headers.get(_HEADER) == _REQUEST_BLOCKER, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_RESPONSE_FIRST,)
|
||||
assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_streamed_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=True)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.rstrip().endswith("data: [DONE]"), response.text
|
||||
assert _steps_that_ran(rig) == (_RESPONSE_FIRST,)
|
||||
assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_request_pipeline_names_a_step_that_flagged_the_request_before_moving_on(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _MOVES_PAST_FIRST)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_REQUEST_BLOCKER, _REQUEST_SECOND)
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_BLOCKER, _REQUEST_SECOND}, dict(
|
||||
response.headers
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"])
|
||||
def test_a_response_pipeline_names_a_step_that_flagged_the_response_before_moving_on(rig: _Rig, stream: bool) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _RESPONSE_MOVES_PAST_FIRST, stream=stream)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _steps_that_ran(rig) == (_RESPONSE_BLOCKER, _RESPONSE_SECOND)
|
||||
assert set(response.headers.get(_HEADER, "").split(",")) == {_RESPONSE_BLOCKER, _RESPONSE_SECOND}, dict(
|
||||
response.headers
|
||||
)
|
||||
|
||||
|
||||
def test_a_stream_released_before_its_pipeline_scans_it_names_the_steps_in_its_headers(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _SCANS_AFTER_RELEASE, stream=True)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.rstrip().endswith("data: [DONE]"), response.text
|
||||
assert _steps_that_ran(rig) == (_LIVE_FIRST, _LIVE_SECOND)
|
||||
assert set(response.headers.get(_HEADER, "").split(",")) == {_LIVE_FIRST, _LIVE_SECOND}, dict(response.headers)
|
||||
|
||||
|
||||
def test_a_stream_its_pipeline_cannot_hold_names_the_steps_that_run_on_their_own(rig: _Rig) -> None:
|
||||
response: Final = _ask_as_team_tagged(rig, _CANNOT_HOLD_STREAM, stream=True)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.rstrip().endswith("data: [DONE]"), response.text
|
||||
assert _RESPONSE_SECOND in _steps_that_ran(rig)
|
||||
assert set(response.headers.get(_HEADER, "").split(",")) == {_UNGATED_STEP, _RESPONSE_SECOND}, dict(
|
||||
response.headers
|
||||
)
|
||||
|
|
@ -1270,7 +1270,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 +1433,16 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he
|
|||
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(monkeypatch):
|
||||
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()])
|
||||
|
|
|
|||
|
|
@ -30,12 +30,16 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
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.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.types.utils import CallTypes
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import (
|
||||
GuardrailPipeline,
|
||||
|
|
@ -275,7 +279,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat
|
|||
|
||||
monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add)
|
||||
data = {"metadata": {"guardrails": ["g1", "g2"]}}
|
||||
proxy_logging._process_guardrail_metadata(data)
|
||||
proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth())
|
||||
snapshot = {
|
||||
"call_count": len(calls),
|
||||
"first_name": calls[0]["name"],
|
||||
|
|
@ -300,7 +304,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp
|
|||
|
||||
monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add)
|
||||
data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}}
|
||||
proxy_logging._process_guardrail_metadata(data)
|
||||
proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth())
|
||||
assert calls == ["g2"]
|
||||
|
||||
|
||||
|
|
@ -312,12 +316,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat
|
|||
"add_guardrail_to_applied_guardrails_header",
|
||||
MagicMock(side_effect=AssertionError("should not be called")),
|
||||
)
|
||||
proxy_logging._process_guardrail_metadata({})
|
||||
proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth())
|
||||
|
||||
|
||||
def test_process_guardrail_metadata_invalid_data_raises(proxy_logging):
|
||||
with pytest.raises(AttributeError):
|
||||
proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type]
|
||||
proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth()) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1160,6 +1164,233 @@ 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, cache, data, call_type):
|
||||
self.calls += 1
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
self.calls += 1
|
||||
|
||||
|
||||
def _two_step_pipeline_data(mode: str, first_on_pass: str, first_on_fail: str = "block") -> dict[str, Any]:
|
||||
pipeline = 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="block"),
|
||||
],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [("governance", pipeline)],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@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, make_user_api_key_auth, monkeypatch, first_on_pass, expected_calls, expected_header
|
||||
):
|
||||
first = _PipelineStepGuardrail("first", "pre_call")
|
||||
second = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result = 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, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
first = _PipelineStepGuardrail("first", "pre_call")
|
||||
second = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result = 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, cache, data, call_type):
|
||||
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, make_user_api_key_auth, monkeypatch, error
|
||||
):
|
||||
first = _RaisingPipelineStepGuardrail("first", error)
|
||||
second = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result = 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, make_user_api_key_auth, monkeypatch, error
|
||||
):
|
||||
first = _RaisingPipelineStepGuardrail("first", error)
|
||||
second = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
result = 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, cache, data, call_type):
|
||||
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, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
first = _PipelineStepGuardrail("first", "pre_call")
|
||||
second = _OpensLitellmMetadataGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result = 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
|
||||
async def test_post_call_pipeline_steps_are_named_only_once_they_run(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
first = _PipelineStepGuardrail("first", "post_call")
|
||||
second = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("post_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
assert "applied_guardrails" not in data["metadata"]
|
||||
|
||||
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"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("asks_for", "request_route", "streams_live", "expected_header"),
|
||||
[
|
||||
pytest.param({"stream": False}, "/custom/stream", False, None, id="whole-response"),
|
||||
pytest.param({"stream": True}, "/v1/chat/completions", False, None, id="stream-held-for-the-pipeline"),
|
||||
pytest.param({"stream": True}, "/v1/chat/completions", True, ["first", "second"], id="stream-released-live"),
|
||||
pytest.param(
|
||||
{"stream": True}, "/custom/stream", False, ["first", "second"], id="stream-runs-steps-on-their-own"
|
||||
),
|
||||
pytest.param({"method": "message/send"}, "/a2a/agent", False, None, id="a2a-whole-response"),
|
||||
pytest.param({"method": "message/stream"}, "/a2a/agent", False, ["first", "second"], id="a2a-stream"),
|
||||
pytest.param({"method": "SubscribeToTask"}, "/a2a/agent", False, ["first", "second"], id="a2a-v1-resubscribe"),
|
||||
],
|
||||
)
|
||||
async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, asks_for, request_route, streams_live, expected_header
|
||||
):
|
||||
guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")]
|
||||
for guardrail in guardrails:
|
||||
guardrail.streaming_buffer_until_moderated = not streams_live
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
data = _two_step_pipeline_data("post_call", "next")
|
||||
data["metadata"]["_guardrail_pipelines"] = [
|
||||
(
|
||||
"governance",
|
||||
GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="next") for name in ("first", "second")],
|
||||
),
|
||||
)
|
||||
]
|
||||
data.update(asks_for)
|
||||
|
||||
result = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(request_route=request_route),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert result["metadata"].get("applied_guardrails") == expected_header
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_replacement_response_reaches_caller(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
|
|
@ -170,7 +170,7 @@ 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, user_api_key_dict):
|
||||
invoked["data"] = d
|
||||
|
||||
proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue