mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(guardrails): keep naming pipeline steps that ran at the deployment level, in a later pipeline, or on another request
This commit is contained in:
yucheng
parent
c984395879
commit
dcf0f62400
5 changed files with 351 additions and 41 deletions
|
|
@ -3,7 +3,7 @@ import copy
|
|||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
|
@ -70,6 +70,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
|
|||
# proxy's metadata sanitizer.
|
||||
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
|
||||
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
|
||||
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
|
||||
|
||||
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
|
||||
|
||||
|
|
@ -219,6 +220,25 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
|
|||
)
|
||||
|
||||
|
||||
def _str_collection(value: object) -> frozenset[str]:
|
||||
if not isinstance(value, (list, tuple, set, frozenset)):
|
||||
return frozenset()
|
||||
items: Final = cast("Iterable[object]", value) # cast-ok: a metadata collection of guardrail names
|
||||
return frozenset(item for item in items if isinstance(item, str))
|
||||
|
||||
|
||||
def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> bool:
|
||||
if marker is None:
|
||||
return False
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list) and marker in executed:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
|
@ -825,11 +845,14 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
return False
|
||||
|
||||
def _pre_call_marker(self, *, raised: bool = False) -> str | None:
|
||||
def _marker(self, token: str) -> str | None:
|
||||
name: Final = self.guardrail_name
|
||||
if not name:
|
||||
return None
|
||||
return f"{_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
return f"{token}:{name}"
|
||||
|
||||
def _pre_call_marker(self, *, raised: bool = False) -> str | None:
|
||||
return self._marker(_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN)
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None:
|
||||
"""
|
||||
|
|
@ -864,16 +887,47 @@ class CustomGuardrail(CustomLogger):
|
|||
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"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list) and marker in executed:
|
||||
return True
|
||||
return False
|
||||
return _request_carries_marker(data, self._pre_call_marker(raised=raised))
|
||||
|
||||
def deployment_hook_ran(self, data: Mapping[str, object]) -> bool:
|
||||
"""Whether a deployment-level hook ran this guardrail as a policy pipeline step of this request."""
|
||||
return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN))
|
||||
|
||||
def _record_deployment_hook_run(
|
||||
self,
|
||||
data: dict[str, object], # mutable-ok: the request payload the run is recorded on
|
||||
) -> None:
|
||||
"""
|
||||
Record that a deployment-level hook ran this guardrail when it is a policy pipeline step of
|
||||
this request, and name it in the applied guardrails when the request's metadata lists it.
|
||||
"""
|
||||
name: Final = self.guardrail_name
|
||||
marker: Final = self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)
|
||||
buckets: Final = tuple(
|
||||
cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed
|
||||
for bucket in (data.get("metadata"), data.get("litellm_metadata"))
|
||||
if isinstance(bucket, dict)
|
||||
)
|
||||
policy_state: Final = next(
|
||||
(bucket for bucket in buckets if name in _str_collection(bucket.get("_pipeline_managed_guardrails"))),
|
||||
None,
|
||||
)
|
||||
if marker is None or policy_state is None:
|
||||
return
|
||||
recorded: Final = policy_state.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if not isinstance(recorded, list):
|
||||
policy_state[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
elif marker not in recorded:
|
||||
cast("list[object]", recorded).append(marker) # cast-ok: markers are only compared as strings
|
||||
requested: Final = next((bucket.get("guardrails") for bucket in buckets if "guardrails" in bucket), None)
|
||||
if name not in _str_collection(requested):
|
||||
return
|
||||
metadata_bucket: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str keys
|
||||
applied: Final = metadata_bucket.get("applied_guardrails")
|
||||
if not isinstance(applied, list):
|
||||
metadata_bucket["applied_guardrails"] = [name]
|
||||
elif name not in applied:
|
||||
cast("list[object]", applied).append(name) # cast-ok: entries are only compared with names
|
||||
|
||||
@classmethod
|
||||
def uses_apply_guardrail_interface(cls) -> bool:
|
||||
|
|
@ -922,12 +976,15 @@ class CustomGuardrail(CustomLogger):
|
|||
target: Final = self._deployment_hook_target()
|
||||
if target is not self:
|
||||
kwargs["guardrail_to_apply"] = self
|
||||
result: Final = await target.async_pre_call_hook(
|
||||
user_api_key_dict=_user_api_key_auth_from_request(kwargs),
|
||||
cache=dc,
|
||||
data=kwargs,
|
||||
call_type="completion" if call_type == CallTypes.completion else "acompletion",
|
||||
)
|
||||
try:
|
||||
result: Final = await target.async_pre_call_hook(
|
||||
user_api_key_dict=_user_api_key_auth_from_request(kwargs),
|
||||
cache=dc,
|
||||
data=kwargs,
|
||||
call_type="completion" if call_type == CallTypes.completion else "acompletion",
|
||||
)
|
||||
finally:
|
||||
self._record_deployment_hook_run(kwargs)
|
||||
|
||||
if result is not None and isinstance(result, dict):
|
||||
result_messages: Final = result.get("messages")
|
||||
|
|
@ -1011,6 +1068,9 @@ class CustomGuardrail(CustomLogger):
|
|||
finally:
|
||||
if target is not self:
|
||||
request_data.pop("guardrail_to_apply", None)
|
||||
self._record_deployment_hook_run(
|
||||
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
|
||||
)
|
||||
|
||||
if not self._is_valid_response_type(result):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -713,6 +713,7 @@ def _ran_or_runs_outside_pipelines(
|
|||
) -> bool:
|
||||
return callback is not None and (
|
||||
callback.pre_call_hook_ran_or_raised(data)
|
||||
or callback.deployment_hook_ran(data)
|
||||
or any(
|
||||
callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks
|
||||
for hook in GuardrailEventHooks
|
||||
|
|
@ -766,13 +767,25 @@ def _pipelines_settled_by(
|
|||
return pipelines[index:] if result.terminal_action != "allow" else pipelines[index : index + 1]
|
||||
|
||||
|
||||
def _steps_later_pipelines_still_run(
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
|
||||
) -> frozenset[str]:
|
||||
"""The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on."""
|
||||
if result.terminal_action != "allow":
|
||||
return frozenset()
|
||||
later: Final = tuple(
|
||||
(policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"
|
||||
)
|
||||
return _pipeline_step_guardrail_names(later)
|
||||
|
||||
|
||||
def _drop_response_steps_that_did_not_run(
|
||||
data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]],
|
||||
steps_that_ran: frozenset[str],
|
||||
steps_to_keep: frozenset[str],
|
||||
) -> None:
|
||||
"""
|
||||
Drop from the applied guardrails the steps of these response pipelines that no pipeline ran,
|
||||
Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``,
|
||||
unless their guardrail also runs outside pipelines. Every step is named up front, so a
|
||||
log built before the pipeline finished, such as a cache hit's, still lists the steps
|
||||
that ran. The list is trimmed in place, since a route's logging object may hold a
|
||||
|
|
@ -781,7 +794,7 @@ def _drop_response_steps_that_did_not_run(
|
|||
pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
post_call: Final = _pipeline_step_guardrail_names(pipelines)
|
||||
not_run: Final = frozenset(
|
||||
name for name in post_call - steps_that_ran if _runs_only_in_pipelines(name, data, pre_call, post_call)
|
||||
name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call)
|
||||
)
|
||||
if not not_run:
|
||||
return
|
||||
|
|
@ -2735,8 +2748,12 @@ class ProxyLogging:
|
|||
data.pop("prompt_version", None)
|
||||
data.pop("prompt_environment", None)
|
||||
|
||||
def _process_guardrail_metadata(self, data: dict) -> None:
|
||||
"""Process guardrails from metadata and add to applied_guardrails."""
|
||||
def _process_guardrail_metadata(self, data: dict, *, skip_guardrails: bool = False) -> None:
|
||||
"""
|
||||
Process guardrails from metadata and add to applied_guardrails, leaving out the pre_call
|
||||
policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on
|
||||
another request, so every step is added.
|
||||
"""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
|
@ -2760,8 +2777,10 @@ class ProxyLogging:
|
|||
if not isinstance(applied_guardrails, list):
|
||||
applied_guardrails = []
|
||||
|
||||
not_run: Final = _request_steps_that_have_not_run(
|
||||
cast(dict[str, object], data) # cast-ok: request payload
|
||||
not_run: Final = (
|
||||
frozenset[str]()
|
||||
if skip_guardrails
|
||||
else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload
|
||||
)
|
||||
for guardrail_name in guardrails_in_metadata:
|
||||
if (
|
||||
|
|
@ -2822,7 +2841,9 @@ class ProxyLogging:
|
|||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if log_waits:
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request, _pipelines_settled_by(pipelines, index, result), steps_that_ran
|
||||
request,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
|
||||
data = self._handle_pipeline_result(
|
||||
|
|
@ -3064,7 +3085,7 @@ class ProxyLogging:
|
|||
and (guardrails_only or not caps.has_pre_call_override)
|
||||
):
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
return data
|
||||
|
||||
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
|
||||
|
|
@ -3168,17 +3189,17 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
|
||||
return data
|
||||
except SensitiveDataRouteException as e:
|
||||
data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict)
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
return data
|
||||
except Exception:
|
||||
if data is not None:
|
||||
self._process_guardrail_metadata(data)
|
||||
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
|
||||
raise
|
||||
|
||||
async def _run_parallel_pre_call_guardrails(
|
||||
|
|
@ -4757,7 +4778,9 @@ class ProxyLogging:
|
|||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if _stream_log_waits_for_its_pipelines(request_data):
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request_data, _pipelines_settled_by(pipelines, index, result), steps_that_ran
|
||||
request_data,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
try:
|
||||
ProxyLogging._handle_pipeline_result(
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ _RUNS_EVERY_STEP: Final = "runs-every-step"
|
|||
_BLOCKS_AT_FIRST: Final = "blocks-at-first"
|
||||
_RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first"
|
||||
_RESPONSE_RUNS_EVERY_STEP: Final = "response-runs-every-step"
|
||||
_RESPONSE_STOPS_AFTER_SECOND: Final = "response-stops-after-second"
|
||||
_SHIELDS_THEN_CHECKS: Final = "shields-then-checks"
|
||||
_MOVES_PAST_FIRST: Final = "moves-past-first"
|
||||
_RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first"
|
||||
|
|
@ -418,6 +419,9 @@ def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path:
|
|||
_RESPONSE_RUNS_EVERY_STEP: _policy(
|
||||
"post_call", ((_RESPONSE_FIRST, "next", "block"), (_RESPONSE_SECOND, "next", "block"))
|
||||
),
|
||||
_RESPONSE_STOPS_AFTER_SECOND: _policy(
|
||||
"post_call", ((_RESPONSE_SECOND, "allow", "block"), (_RESPONSE_FIRST, "next", "block"))
|
||||
),
|
||||
_SHIELDS_THEN_CHECKS: _policy(
|
||||
"pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block"))
|
||||
),
|
||||
|
|
@ -695,6 +699,35 @@ def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran(
|
|||
] * len(spans), spans
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tag", "model_guardrail", "expected"),
|
||||
[
|
||||
pytest.param(_STOPS_AFTER_FIRST, _REQUEST_SECOND, [_REQUEST_FIRST, _REQUEST_SECOND], id="request-pipeline"),
|
||||
pytest.param(
|
||||
_RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_SECOND, [_RESPONSE_FIRST, _RESPONSE_SECOND], id="response-pipeline"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_model_level_guardrail_the_pipeline_skipped_is_named_everywhere_once_the_model_runs_it(
|
||||
rig: _Rig, tag: str, model_guardrail: str, expected: list[str]
|
||||
) -> None:
|
||||
since: Final = recorded_spans(rig.traces)[0]
|
||||
with rig.gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture", guardrails=[model_guardrail]
|
||||
)
|
||||
team: Final = scenario.team(metadata={"tags": [tag]})
|
||||
key: Final = scenario.key(team_id=team)
|
||||
rig.vendor.drain()
|
||||
response: Final = rig.gateway.request(
|
||||
"POST", "/v1/chat/completions", _body("/v1/chat/completions", model, False), key=key
|
||||
)
|
||||
header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response)
|
||||
assert response.status_code == 200, response.text
|
||||
assert sorted(_steps_that_ran(rig)) == expected
|
||||
assert (header, rows, spans) == (expected, [expected], [expected] * len(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)
|
||||
|
|
@ -1005,17 +1038,22 @@ def agent(rig: _Rig) -> Iterator[str]:
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tag", "expected"),
|
||||
("tags", "expected"),
|
||||
[
|
||||
pytest.param(_SCANS_AFTER_RELEASE, [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"),
|
||||
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"
|
||||
(_RESPONSE_MOVES_PAST_FIRST,), [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on"
|
||||
),
|
||||
pytest.param((_RESPONSE_STOPS_AFTER_FIRST,), [_RESPONSE_FIRST], id="first-step-allows"),
|
||||
pytest.param(
|
||||
(_RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_STOPS_AFTER_SECOND),
|
||||
[_RESPONSE_FIRST, _RESPONSE_SECOND],
|
||||
id="two-policies-each-run-the-step-the-other-skipped",
|
||||
),
|
||||
pytest.param(_RESPONSE_STOPS_AFTER_FIRST, [_RESPONSE_FIRST], id="first-step-allows"),
|
||||
],
|
||||
)
|
||||
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]
|
||||
rig: _Rig, agent: str, tags: tuple[str, ...], expected: list[str]
|
||||
) -> None:
|
||||
since: Final = recorded_spans(rig.traces)[0]
|
||||
rig.vendor.drain()
|
||||
|
|
@ -1026,7 +1064,7 @@ def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_an
|
|||
"parts": [{"kind": "text", "text": f"hello {uuid.uuid4().hex}"}],
|
||||
}
|
||||
with rig.gateway.scenario() as scenario:
|
||||
team: Final = scenario.team(metadata={"tags": [tag]})
|
||||
team: Final = scenario.team(metadata={"tags": list(tags)})
|
||||
response: Final = rig.gateway.request(
|
||||
"POST",
|
||||
f"/a2a/{agent}",
|
||||
|
|
@ -1035,7 +1073,7 @@ def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_an
|
|||
headers={"Accept": "application/json"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert list(_steps_that_ran(rig)) == expected
|
||||
assert sorted(_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,
|
||||
|
|
|
|||
|
|
@ -1672,6 +1672,194 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_
|
|||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
stops_before_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")],
|
||||
)
|
||||
stops_before_first: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")],
|
||||
)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [
|
||||
("stops-before-second", stops_before_second),
|
||||
("stops-before-first", stops_before_first),
|
||||
],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
|
||||
result: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
skip_guardrails=True,
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (0, 0)
|
||||
assert result["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_names_a_skipped_pre_call_step_it_ran(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
await second.async_pre_call_deployment_hook(
|
||||
kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
assert second.deployment_hook_ran(data)
|
||||
assert not first.deployment_hook_ran(data)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS)
|
||||
async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
error: Exception,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "pre_call")
|
||||
second: Final = _RaisingPipelineStepGuardrail("second", error)
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=_two_step_pipeline_data("pre_call", "allow"),
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(type(error)):
|
||||
await second.async_pre_call_deployment_hook(
|
||||
kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_skipped_step_a_deployment_hook_ran(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
await second.async_post_call_success_deployment_hook(
|
||||
request_data={**data, "guardrails": ["second"]},
|
||||
response=litellm.ModelResponse(),
|
||||
call_type=CallTypes.acompletion,
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_leaves_a_request_without_policy_pipelines_untouched() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": {"guardrails": ["second"]},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert kwargs["metadata"] == {"guardrails": ["second"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_names_a_pipeline_step_it_ran_before_any_guardrail_was_named() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
metadata: Final[dict[str, object]] = {
|
||||
"guardrails": ["second"],
|
||||
"_pipeline_managed_guardrails": frozenset({"second"}),
|
||||
}
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.calls == 1
|
||||
assert metadata["applied_guardrails"] == ["second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one_the_request_did_not_list() -> None:
|
||||
guardrail: Final = _PipelineStepGuardrail("second", "pre_call")
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrails": ["second"],
|
||||
"metadata": {"_pipeline_managed_guardrails": frozenset({"second"})},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion)
|
||||
|
||||
assert guardrail.deployment_hook_ran(kwargs)
|
||||
assert not _PipelineStepGuardrail("first", "pre_call").deployment_hook_ran(kwargs)
|
||||
assert "applied_guardrails" not in kwargs["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("log_waits", "expected_after_stream"),
|
||||
|
|
|
|||
|
|
@ -170,8 +170,9 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
|
|||
proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
|
||||
invoked = {}
|
||||
|
||||
def fake_process(d):
|
||||
def fake_process(d, *, skip_guardrails):
|
||||
invoked["data"] = d
|
||||
invoked["skip_guardrails"] = skip_guardrails
|
||||
|
||||
proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
|
|
@ -180,7 +181,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
|
|||
call_type="completion",
|
||||
)
|
||||
assert out is data
|
||||
assert invoked["data"] is data
|
||||
assert invoked == {"data": data, "skip_guardrails": False}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue