scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

Merge 5434ce4675 into 73cd296418

Browse source
This commit is contained in:
yucheng-berri 2026-10-28 00:11:49 +00:00 • committed by GitHub
commit fb67e5c11f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1442 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/integrations/custom_guardrail.py
View file

@ -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"):

23
litellm/proxy/a2a/version_convert.py
View file

@ -24,7 +24,8 @@ The two wire shapes:
``Task`` result is a bare object without ``kind``.
"""
from collections.abc import Callable
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel
@ -39,6 +40,26 @@ JsonDict = dict[str, object]
_V1_SEND_ENVELOPE_KEYS: Final = frozenset({"message", "task"})
_V1_STREAM_ENVELOPE_KEYS: Final = frozenset({"message", "task", "statusUpdate", "artifactUpdate"})
A2A_PASCAL_TO_WIRE: Final[Mapping[str, str]] = MappingProxyType(
{
"SendMessage": "message/send",
"SendStreamingMessage": "message/stream",
"GetTask": "tasks/get",
"ListTasks": "tasks/list",
"CancelTask": "tasks/cancel",
"SubscribeToTask": "tasks/resubscribe",
"CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set",
"GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get",
"ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list",
"DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete",
"GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard",
}
)
_STREAMING_WIRE_METHODS: Final = frozenset({"message/stream", "tasks/resubscribe"})
A2A_STREAMING_METHODS: Final = _STREAMING_WIRE_METHODS | frozenset(
pascal for pascal, wire in A2A_PASCAL_TO_WIRE.items() if wire in _STREAMING_WIRE_METHODS
)
def _dump_03(model: BaseModel) -> JsonDict:
"""Dump a compat (0.3) pydantic model to its camelCase wire dict."""

20
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -28,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.a2a.version_convert import (
A2A_PASCAL_TO_WIRE,
A2AVersion,
normalize_agent_card,
normalize_jsonrpc_response,
@ -65,21 +66,6 @@ _SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
}
)
_PASCAL_TO_WIRE: Final[Mapping[str, str]] = {
"SendMessage": "message/send",
"SendStreamingMessage": "message/stream",
"GetTask": "tasks/get",
"ListTasks": "tasks/list",
"CancelTask": "tasks/cancel",
"SubscribeToTask": "tasks/resubscribe",
"CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set",
"GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get",
"ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list",
"DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete",
"GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard",
}
_DECODED_JSON: Final = TypeAdapter(object)
@ -132,7 +118,7 @@ def _served_version(agent: "AgentResponse", request: Request, original_method: s
configured: Final = (agent.agent_card_params or {}).get("protocolVersion")
if configured in ("0.3", "1.0"):
return configured
if original_method in _PASCAL_TO_WIRE:
if original_method in A2A_PASCAL_TO_WIRE:
return "1.0"
return "1.0" if request.headers.get("a2a-version", "").startswith("1.") else "0.3"
@ -699,7 +685,7 @@ async def invoke_agent_a2a(
params = body.get("params", {})
if method:
method = _PASCAL_TO_WIRE.get(method, method)
method = A2A_PASCAL_TO_WIRE.get(method, method)
if isinstance(params, dict):
# extract any litellm params from the params - eg. 'guardrails'

74
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -329,6 +329,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,
@ -389,6 +431,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()
@ -407,6 +452,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,
@ -480,6 +526,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
@ -489,8 +536,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 +572,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(
@ -539,6 +585,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,
@ -582,16 +629,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 (
@ -609,9 +647,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,

102
litellm/proxy/utils.py
View file

@ -179,6 +179,7 @@ from litellm.proxy._types import (
SpecialModelNames,
UserAPIKeyAuth,
)
from litellm.proxy.a2a.version_convert import A2A_STREAMING_METHODS
from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
@ -706,6 +707,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
)
)
def _post_call_steps_named_before_headers(
data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str
) -> 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. The LLM pass-through route builds
its headers before any response step runs, so none of its steps.
"""
if call_type == CallTypes.allm_passthrough_route.value:
return frozenset()
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, call_type: str
) -> 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, and the LLM pass-through route runs
its response steps after its headers are built, 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, call_type)
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 +984,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(
@ -2698,7 +2759,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, call_type: str) -> None:
"""Process guardrails from metadata and add to applied_guardrails."""
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -2723,8 +2784,17 @@ 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,
call_type,
)
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(
@ -3011,7 +3081,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, call_type)
return data
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
@ -3115,17 +3185,17 @@ class ProxyLogging:
)
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, user_api_key_dict, call_type)
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, call_type)
return data
except Exception:
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, user_api_key_dict, call_type)
raise
async def _run_parallel_pre_call_guardrails(

900
tests/integration/observability/test_guardrail_pipeline_applied_header.py Normal file
View file

@ -0,0 +1,900 @@
from __future__ import annotations
import ast
import asyncio
import itertools
import json
import uuid
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anthropic
import httpx
import openai
import pytest
import yaml
from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value
from tests.integration._support.otlp_sink import Span, SpanSinks, recorded_spans
from tests.integration._support.process import owned_proxy
from tests.integration._support.upstream import aws_event_stream_frame
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"
_REQUEST_ERRORING: Final = "request-erroring-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"
_JUDGE: Final = "judge-step"
_JUDGE_MODEL: Final = "pipeline-judge"
_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"
_ERRORS_PAST_FIRST: Final = "errors-past-first"
_JUDGES_THEN_CHECKS: Final = "judges-then-checks"
_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages")
_BEDROCK_MODEL: Final = "pipeline-bedrock"
_BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
@dataclass(frozen=True, slots=True)
class _Rig:
gateway: Gateway
vendor: Wire
upstream: Wire
model: str
traces: 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 _event(event: dict[str, object]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def _response(identity: str) -> dict[str, object]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"id": f"msg_{identity}",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "guarded", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4},
}
def _responses_reply(request: Request) -> Reply:
identity: Final = f"resp_{uuid.uuid4().hex}"
response: Final = _response(identity)
if json.loads(request.body).get("stream") is not True:
return Reply(body=json.dumps(response).encode())
return Reply(
content_type="text/event-stream",
chunks=(
_event(
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
}
),
_event(
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": f"msg_{identity}",
"output_index": 0,
"content_index": 0,
"delta": "guarded",
}
),
_event({"type": "response.completed", "sequence_number": 2, "response": response}),
),
)
def _bedrock_event(event_type: str, payload: dict[str, object]) -> bytes:
return aws_event_stream_frame(
{":event-type": event_type, ":content-type": "application/json", ":message-type": "event"},
json.dumps(payload).encode(),
)
def _bedrock_reply(request: Request) -> Reply:
action: Final = request.target.rsplit("/", 1)[1]
if action == "converse-stream":
return Reply(
content_type="application/vnd.amazon.eventstream",
chunks=(
_bedrock_event("messageStart", {"role": "assistant"}),
_bedrock_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "guarded"}}),
_bedrock_event("contentBlockStop", {"contentBlockIndex": 0}),
_bedrock_event("messageStop", {"stopReason": "end_turn"}),
_bedrock_event(
"metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, "metrics": {}}
),
),
)
if action == "invoke":
return Reply(
body=json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": _BEDROCK_MODEL_ID,
"content": [{"type": "text", "text": "guarded"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 3, "output_tokens": 1},
}
).encode()
)
assert action == "converse", request.target
return Reply(
body=json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "guarded"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4},
"metrics": {"latencyMs": 1},
}
).encode()
)
def _completion(request: Request) -> Reply:
if request.target.startswith("/model/"):
return _bedrock_reply(request)
if request.target == "/v1/responses":
return _responses_reply(request)
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 == f"/{_JUDGE}/v1/decisions":
return Reply(
body=json.dumps(
{
"model": "gpt-6-luna",
"answers": [{"type": "predicate", "name": "unsafe_output", "probability": 0.01}],
"usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
}
).encode()
)
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_ERRORING}/"):
return Reply(status=503, body=json.dumps({"error": "guardrail service unavailable"}).encode())
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], ...], *, on_error: str | None = None) -> 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, "on_error": on_error}
for guardrail, on_pass, on_fail in steps
],
},
}
def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path:
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["litellm_settings"] = {**configuration.get("litellm_settings", {}), "callbacks": ["otel"]}
configuration["model_list"] = [
*configuration.get("model_list", []),
{
"model_name": _BEDROCK_MODEL,
"litellm_params": {
"model": f"bedrock/{_BEDROCK_MODEL_ID}",
"api_base": upstream.url,
"aws_access_key_id": "AKIASYNTHETICPROVIDER",
"aws_secret_access_key": "synthetic-provider-secret",
"aws_region_name": "us-east-1",
},
},
{
"model_name": _JUDGE_MODEL,
"litellm_params": {
"model": "openai/gpt-6-luna",
"api_base": f"{vendor.url}/{_JUDGE}/v1",
"api_key": "synthetic-judge-key",
},
},
]
configuration["guardrails"] = [
*(
_guardrail(name, "pre_call", vendor)
for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER, _REQUEST_ERRORING)
),
*(_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": _JUDGE,
"litellm_params": {
"guardrail": "decision_model",
"mode": "post_call",
"default_on": False,
"decision_model": _JUDGE_MODEL,
"checks": [
{
"name": "unsafe_output",
"instructions": "Does the response say anything unsafe?",
"action": "block",
"threshold": 0.5,
}
],
},
},
{
"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"))
),
_JUDGES_THEN_CHECKS: _policy("post_call", ((_JUDGE, "next", "block"), (_RESPONSE_SECOND, "next", "block"))),
_ERRORS_PAST_FIRST: _policy(
"pre_call", ((_REQUEST_ERRORING, "next", "block"), (_REQUEST_SECOND, "next", "block")), on_error="next"
),
}
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, audit_sinks: SpanSinks) -> 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,
{"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "OTEL_BSP_SCHEDULE_DELAY": "100"},
config=_config(directory, vendor, upstream),
workers=2,
) as owned,
owned.scenario() as scenario,
):
model: Final = scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture")
yield _Rig(owned, vendor, upstream, model, audit_sinks.operator)
def _body(path: str, model: str, stream: bool) -> dict[str, object]:
prompt: Final = f"hello {uuid.uuid4().hex}"
if path == "/v1/responses":
return {"model": model, "input": prompt, "stream": stream}
messages: Final = [{"role": "user", "content": prompt}]
if path == "/v1/messages":
return {"model": model, "max_tokens": 20, "messages": messages, "stream": stream}
return {"model": model, "messages": messages, "stream": stream}
def _ask_as_new_team(
rig: _Rig,
tag: str,
*,
stream: bool = False,
path: str = "/v1/chat/completions",
extra: Mapping[str, object] = MappingProxyType({}),
) -> tuple[str, 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 team, rig.gateway.request("POST", path, {**_body(path, rig.model, stream), **extra}, key=key)
def _ask_as_team_tagged(
rig: _Rig,
tag: str,
*,
stream: bool = False,
path: str = "/v1/chat/completions",
extra: Mapping[str, object] = MappingProxyType({}),
) -> httpx.Response:
return _ask_as_new_team(rig, tag, stream=stream, path=path, extra=extra)[1]
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
)
@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"])
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint(
rig: _Rig, path: str, stream: bool
) -> None:
response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=stream, path=path)
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)
@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"])
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint(
rig: _Rig, path: str, stream: bool
) -> None:
response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=stream, path=path)
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_request_pipeline_names_a_step_whose_service_errored_before_moving_on(rig: _Rig) -> None:
response: Final = _ask_as_team_tagged(rig, _ERRORS_PAST_FIRST)
assert response.status_code == 200, response.text
assert _steps_that_ran(rig) == (_REQUEST_ERRORING, _REQUEST_SECOND)
assert len(rig.upstream.drain()) == 1
assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_ERRORING, _REQUEST_SECOND}, dict(
response.headers
)
@pytest.mark.parametrize(
"forged",
[
pytest.param(
{
"metadata": {
"applied_guardrails": [_REQUEST_SECOND],
"_pre_call_executed_guardrails": [f"deadbeef:{_REQUEST_SECOND}", _REQUEST_SECOND],
}
},
id="names-in-its-metadata",
),
pytest.param({"guardrails": [_REQUEST_SECOND]}, id="asks-for-the-skipped-step"),
],
)
def test_a_request_cannot_get_a_step_that_never_ran_named(rig: _Rig, forged: Mapping[str, object]) -> None:
response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, extra=forged)
assert response.status_code == 200, response.text
assert _steps_that_ran(rig) == (_REQUEST_FIRST,)
assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers)
def _spans_naming_guardrails(rig: _Rig, since: int, team: str) -> tuple[Span, ...]:
spans: Final = recorded_spans(rig.traces, since)[1]
return tuple(
span
for span in spans
if span["attributes"].get("metadata.user_api_key_team_id") == team
and "metadata.applied_guardrails" in span["attributes"]
)
@pytest.mark.parametrize(
("tag", "path", "stream", "expected"),
[
pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_REQUEST_FIRST], id="request-pipeline"),
pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_REQUEST_FIRST], id="request-pipeline-stream"),
pytest.param(_STOPS_AFTER_FIRST, "/v1/responses", True, [_REQUEST_FIRST], id="request-pipeline-responses"),
pytest.param(
_RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_RESPONSE_FIRST], id="response-pipeline"
),
pytest.param(
_RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_RESPONSE_FIRST], id="response-pipeline-stream"
),
pytest.param(
_RESPONSE_STOPS_AFTER_FIRST, "/v1/messages", True, [_RESPONSE_FIRST], id="response-pipeline-messages"
),
],
)
def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran(
rig: _Rig, tag: str, path: str, stream: bool, expected: list[str]
) -> None:
since: Final = recorded_spans(rig.traces)[0]
team, response = _ask_as_new_team(rig, tag, stream=stream, path=path)
assert response.status_code == 200, response.text
rows: Final = eventually(
lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
lambda found: len(found) == 1,
seconds=70,
)
spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), bool, seconds=30)
assert object_value(rows[0]["metadata"]).get("applied_guardrails") == expected, rows
assert [ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"])) for span in spans] == [
expected
] * len(spans), spans
@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"])
def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig: _Rig, stream: bool) -> None:
team, response = _ask_as_new_team(rig, _JUDGES_THEN_CHECKS, stream=stream)
assert response.status_code == 200, response.text
assert _steps_that_ran(rig) == (_JUDGE, _RESPONSE_SECOND)
assert sorted(response.headers.get(_HEADER, "").split(",")) == [_JUDGE, _RESPONSE_SECOND], dict(response.headers)
rows: Final = eventually(
lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
lambda found: len(found) == 1,
seconds=70,
)
assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND]
def _bedrock_body(action: str) -> dict[str, object]:
prompt: Final = f"hello {uuid.uuid4().hex}"
if action == "invoke":
return {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 20,
"messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
}
return {"messages": [{"role": "user", "content": [{"text": prompt}]}]}
@pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream"])
def test_a_pass_through_route_names_its_response_pipeline_up_front_because_its_headers_are_built_first(
rig: _Rig, action: str
) -> None:
rig.vendor.drain()
rig.upstream.drain()
with rig.gateway.scenario() as scenario:
team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]})
response: Final = rig.gateway.request(
"POST", f"/bedrock/model/{_BEDROCK_MODEL}/{action}", _bedrock_body(action), key=scenario.key(team_id=team)
)
assert response.status_code == 200, response.text
assert _steps_that_ran(rig) == (_RESPONSE_FIRST,)
assert [request.target.rsplit("/", 1)[1] for request in rig.upstream.drain()] == [action]
assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _RESPONSE_SECOND], dict(
response.headers
)
rows: Final = eventually(
lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
lambda found: len(found) == 1,
seconds=70,
)
assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [
_RESPONSE_FIRST,
_RESPONSE_SECOND,
], rows
def _agent_card(name: str, url: str) -> dict[str, object]:
return {
"protocolVersion": "0.3",
"name": name,
"description": "Scripted agent that answers every message with one text part",
"version": "1.0.0",
"url": url,
"capabilities": {"streaming": True},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
"skills": [],
}
@pytest.fixture(scope="module")
def agent(rig: _Rig) -> Iterator[str]:
name: Final = f"pipeline-agent-{uuid.uuid4().hex}"
def answer(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=json.dumps(_agent_card(name, f"{peer.url}/")).encode())
message: Final = {
"kind": "message",
"role": "agent",
"messageId": uuid.uuid4().hex,
"parts": [{"kind": "text", "text": "guarded"}],
}
return Reply(
body=json.dumps({"jsonrpc": "2.0", "id": json.loads(request.body)["id"], "result": message}).encode()
)
with wire_server(answer) as peer:
identity: Final = str(
rig.gateway.post(
"/v1/agents", {"agent_name": name, "agent_card_params": _agent_card(name, f"{peer.url}/")}
)["agent_id"]
)
yield identity
rig.gateway.request("DELETE", f"/v1/agents/{identity}")
@pytest.mark.parametrize(
("tag", "expected"),
[
pytest.param(_SCANS_AFTER_RELEASE, [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"),
pytest.param(
_RESPONSE_MOVES_PAST_FIRST, [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on"
),
],
)
def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_and_the_trace(
rig: _Rig, agent: str, tag: str, expected: list[str]
) -> None:
since: Final = recorded_spans(rig.traces)[0]
rig.vendor.drain()
message: Final = {
"kind": "message",
"role": "user",
"messageId": uuid.uuid4().hex,
"parts": [{"kind": "text", "text": f"hello {uuid.uuid4().hex}"}],
}
with rig.gateway.scenario() as scenario:
team: Final = scenario.team(metadata={"tags": [tag]})
response: Final = rig.gateway.request(
"POST",
f"/a2a/{agent}",
{"jsonrpc": "2.0", "id": uuid.uuid4().hex, "method": "message/send", "params": {"message": message}},
key=scenario.key(team_id=team),
headers={"Accept": "application/json"},
)
assert response.status_code == 200, response.text
assert list(_steps_that_ran(rig)) == expected
rows: Final = eventually(
lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
lambda found: len(found) == 1,
seconds=70,
)
spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), bool, seconds=30)
assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == expected, rows
assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [
expected
] * len(spans), spans
def _openai_chat(base: str, key: str, model: str) -> str | None:
with (
httpx.Client(timeout=30, trust_env=False) as transport,
openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client,
):
response: Final = client.chat.completions.with_raw_response.create(
model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}]
)
return response.headers.get(_HEADER)
def _openai_responses(base: str, key: str, model: str) -> str | None:
with (
httpx.Client(timeout=30, trust_env=False) as transport,
openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client,
):
return client.responses.with_raw_response.create(model=model, input=f"hello {uuid.uuid4().hex}").headers.get(
_HEADER
)
def _async_openai_streamed_chat(base: str, key: str, model: str) -> str | None:
async def run() -> str | None:
async with (
httpx.AsyncClient(timeout=30, trust_env=False) as transport,
openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client,
client.chat.completions.with_streaming_response.create(
model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], stream=True
) as response,
):
lines: Final = [line async for line in response.iter_lines()]
assert "data: [DONE]" in lines, lines
return response.headers.get(_HEADER)
return asyncio.run(run())
def _async_openai_streamed_responses(base: str, key: str, model: str) -> str | None:
async def run() -> str | None:
async with (
httpx.AsyncClient(timeout=30, trust_env=False) as transport,
openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client,
client.responses.with_streaming_response.create(
model=model, input=f"hello {uuid.uuid4().hex}", stream=True
) as response,
):
lines: Final = [line async for line in response.iter_lines()]
assert any("response.completed" in line for line in lines), lines
return response.headers.get(_HEADER)
return asyncio.run(run())
def _anthropic_streamed_messages(base: str, key: str, model: str) -> str | None:
with (
httpx.Client(timeout=30, trust_env=False) as transport,
anthropic.Anthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client,
client.messages.with_streaming_response.create(
model=model,
max_tokens=20,
messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}],
stream=True,
) as response,
):
lines: Final = tuple(response.iter_lines())
assert any("message_stop" in line for line in lines), lines
return response.headers.get(_HEADER)
def _async_anthropic_messages(base: str, key: str, model: str) -> str | None:
async def run() -> str | None:
async with (
httpx.AsyncClient(timeout=30, trust_env=False) as transport,
anthropic.AsyncAnthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client,
):
response: Final = await client.messages.with_raw_response.create(
model=model, max_tokens=20, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}]
)
return response.headers.get(_HEADER)
return asyncio.run(run())
@pytest.mark.parametrize(
"client",
[
pytest.param(_openai_chat, id="openai-chat"),
pytest.param(_async_openai_streamed_chat, id="async-openai-chat-stream"),
pytest.param(_openai_responses, id="openai-responses"),
pytest.param(_async_openai_streamed_responses, id="async-openai-responses-stream"),
pytest.param(_anthropic_streamed_messages, id="anthropic-messages-stream"),
pytest.param(_async_anthropic_messages, id="async-anthropic-messages"),
],
)
def test_every_sdk_client_sees_only_the_step_that_ran(rig: _Rig, client: Callable[[str, str, str], str | None]) -> None:
rig.vendor.drain()
with rig.gateway.scenario() as scenario:
key: Final = scenario.key(team_id=scenario.team(metadata={"tags": [_STOPS_AFTER_FIRST]}))
header: Final = client(str(rig.gateway.client.base_url).rstrip("/"), key, rig.model)
assert _steps_that_ran(rig) == (_REQUEST_FIRST,)
assert header == _REQUEST_FIRST
async def _send_concurrently(
rig: _Rig, keys: Mapping[str, str], cases: tuple[tuple[str, str, bool], ...]
) -> tuple[httpx.Response, ...]:
async with httpx.AsyncClient(base_url=rig.gateway.client.base_url, timeout=60, trust_env=False) as client:
return tuple(
await asyncio.gather(
*(
client.post(
path, json=_body(path, rig.model, stream), headers={"Authorization": f"Bearer {keys[tag]}"}
)
for tag, path, stream in cases
)
)
)
def test_a_concurrent_burst_names_only_each_requests_own_steps(rig: _Rig) -> None:
expected: Final = MappingProxyType(
{
_STOPS_AFTER_FIRST: _REQUEST_FIRST,
_RUNS_EVERY_STEP: f"{_REQUEST_FIRST},{_REQUEST_SECOND}",
_RESPONSE_STOPS_AFTER_FIRST: _RESPONSE_FIRST,
}
)
cases: Final = tuple(itertools.product(expected, _ENDPOINTS, (False, True))) * 2
with rig.gateway.scenario() as scenario:
keys: Final = MappingProxyType(
{tag: scenario.key(team_id=scenario.team(metadata={"tags": [tag]})) for tag in expected}
)
responses: Final = asyncio.run(_send_concurrently(rig, keys, cases))
rig.vendor.drain()
rig.upstream.drain()
assert [
(tag, path, stream, response.status_code, response.headers.get(_HEADER))
for (tag, path, stream), response in zip(cases, responses, strict=True)
] == [(tag, path, stream, 200, expected[tag]) for tag, path, stream in cases]

60
tests/unit/proxy/policy_engine/test_pipeline_executor.py
View file

@ -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()])

304
tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -30,12 +30,17 @@ 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.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,
@ -275,7 +280,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(), "completion")
snapshot = {
"call_count": len(calls),
"first_name": calls[0]["name"],
@ -300,7 +305,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(), "completion")
assert calls == ["g2"]
@ -312,12 +317,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(), "completion")
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(), "completion") # type: ignore[arg-type]
# ---------------------------------------------------------------------------
@ -1160,6 +1165,295 @@ 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"},
},
}
@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
async def test_post_call_pipeline_steps_are_named_only_once_they_run(
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"),
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"]
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"),
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(
("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: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
asks_for: dict[str, object],
request_route: str,
streams_live: bool,
expected_header: list[str] | None,
) -> None:
guardrails: Final = [_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: Final = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next")
data.update(asks_for)
result: Final = 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
@pytest.mark.parametrize("stream", [False, True], ids=["whole-response", "stream"])
async def test_pre_call_hook_names_up_front_the_post_call_steps_the_llm_pass_through_route_runs_after_its_headers(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
stream: bool,
) -> None:
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")]
monkeypatch.setattr(litellm, "callbacks", guardrails)
data: Final = _two_step_pipeline_data("post_call", "allow")
data["stream"] = stream
result: Final = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/bedrock/model/m/converse"),
data=data,
call_type="allm_passthrough_route",
)
assert result["metadata"].get("applied_guardrails") == ["first", "second"]
@pytest.mark.asyncio
async def test_post_call_pipeline_replacement_response_reaches_caller(
proxy_logging, make_user_api_key_auth, monkeypatch

2
tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -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: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str) -> None:
invoked["data"] = d
proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]