mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(guardrails): drop the response pipeline steps that never ran from x-litellm-applied-guardrails (#49264)
* refactor(guardrails): record which policy pipeline steps and deployment-level hooks ran * fix(guardrails): name only the request pipeline steps that ran in x-litellm-applied-guardrails * fix(guardrails): drop the response pipeline steps that never ran from x-litellm-applied-guardrails * test(guardrails): integration cells for request pipeline steps in x-litellm-applied-guardrails * fix(guardrails): keep the agent cells on the shared proxy and shorten a long suppression line * fix(guardrails): keep request steps when a response pipeline blocks before a later request pipeline
This commit is contained in:
yucheng-berri
GitHub
parent
ef9385c712
commit
969ded2b54
3 changed files with 627 additions and 2 deletions
|
|
@ -746,6 +746,88 @@ def _request_steps_that_have_not_run(data: Mapping[str, object]) -> frozenset[st
|
|||
return frozenset(name for name in pre_call if _runs_only_in_pipelines(name, data, pre_call, frozenset()))
|
||||
|
||||
|
||||
def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]:
|
||||
return frozenset(step.guardrail_name for step in result.step_results if step.outcome != "skip")
|
||||
|
||||
|
||||
def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool:
|
||||
"""
|
||||
The success log is parked until the response hooks finish, and the headers are built
|
||||
after them. The LLM pass-through route builds its headers first, and the generateContent
|
||||
adapter logs its nested completion before the response hooks run.
|
||||
"""
|
||||
logging_obj: Final = data.get("litellm_logging_obj")
|
||||
return (
|
||||
call_type not in (CallTypes.allm_passthrough_route.value, CallTypes.agenerate_content.value)
|
||||
and isinstance(logging_obj, Logging)
|
||||
and logging_obj.enqueue_deferred_logging is not None
|
||||
)
|
||||
|
||||
|
||||
def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool:
|
||||
return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None
|
||||
|
||||
|
||||
def _pipelines_settled_by(
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
|
||||
) -> Sequence[tuple[str, "GuardrailPipeline"]]:
|
||||
"""This pipeline, plus every later response pipeline when its verdict ends the request before they run."""
|
||||
if result.terminal_action == "allow":
|
||||
return pipelines[index : index + 1]
|
||||
return (
|
||||
pipelines[index],
|
||||
*((policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"),
|
||||
)
|
||||
|
||||
|
||||
def _steps_later_pipelines_still_run(
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
|
||||
) -> frozenset[str]:
|
||||
"""The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on."""
|
||||
if result.terminal_action != "allow":
|
||||
return frozenset()
|
||||
later: Final = tuple(
|
||||
(policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call"
|
||||
)
|
||||
return _pipeline_step_guardrail_names(later)
|
||||
|
||||
|
||||
def _drop_response_steps_that_did_not_run(
|
||||
data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place
|
||||
pipelines: Sequence[tuple[str, "GuardrailPipeline"]],
|
||||
steps_to_keep: frozenset[str],
|
||||
) -> None:
|
||||
"""
|
||||
Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``,
|
||||
unless their guardrail also runs outside pipelines. Every step is named up front, so a
|
||||
log built before the pipeline finished, such as a cache hit's, still lists the steps
|
||||
that ran. The list is trimmed in place, since a route's logging object may hold a
|
||||
shallow copy of the metadata.
|
||||
"""
|
||||
pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
||||
post_call: Final = _pipeline_step_guardrail_names(pipelines)
|
||||
not_run: Final = frozenset(
|
||||
name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call)
|
||||
)
|
||||
if not not_run:
|
||||
return
|
||||
for bucket_key in ("metadata", "litellm_metadata"):
|
||||
_drop_applied_guardrails_in_place(data.get(bucket_key), not_run)
|
||||
|
||||
|
||||
def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None:
|
||||
if not isinstance(bucket, dict):
|
||||
return
|
||||
metadata: Final = cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed
|
||||
applied: Final = metadata.get("applied_guardrails")
|
||||
if not isinstance(applied, list):
|
||||
return
|
||||
entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names
|
||||
entries[:] = [name for name in entries if name not in names]
|
||||
if not entries:
|
||||
metadata.pop("applied_guardrails")
|
||||
|
||||
|
||||
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
|
||||
resolved: Final = tuple(
|
||||
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
||||
|
|
@ -2814,7 +2896,10 @@ class ProxyLogging:
|
|||
return data, None
|
||||
|
||||
current_response = response # rebind-ok: chains each pipeline's replacement response into the next
|
||||
for policy_name, pipeline in pipelines:
|
||||
steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran
|
||||
request: Final = cast(dict[str, object], data) # cast-ok: request payload, updated in place by every pipeline
|
||||
log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type)
|
||||
for index, (policy_name, pipeline) in enumerate(pipelines):
|
||||
if pipeline.mode != event_hook:
|
||||
continue
|
||||
|
||||
|
|
@ -2830,6 +2915,14 @@ class ProxyLogging:
|
|||
raw_request_snapshot=raw_request_snapshot,
|
||||
)
|
||||
|
||||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if log_waits:
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
|
||||
data = self._handle_pipeline_result(
|
||||
result=result,
|
||||
data=data,
|
||||
|
|
@ -4781,7 +4874,8 @@ class ProxyLogging:
|
|||
|
||||
call_type, endpoint_translation = translation
|
||||
|
||||
for policy_name, pipeline in pipelines:
|
||||
steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran
|
||||
for index, (policy_name, pipeline) in enumerate(pipelines):
|
||||
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode="post_call",
|
||||
|
|
@ -4792,6 +4886,13 @@ class ProxyLogging:
|
|||
streaming_chunks=buffered,
|
||||
endpoint_translation=endpoint_translation,
|
||||
)
|
||||
steps_that_ran = steps_that_ran | _steps_that_ran(result)
|
||||
if _stream_log_waits_for_its_pipelines(request_data):
|
||||
_drop_response_steps_that_did_not_run(
|
||||
request_data,
|
||||
_pipelines_settled_by(pipelines, index, result),
|
||||
steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result),
|
||||
)
|
||||
try:
|
||||
ProxyLogging._handle_pipeline_result(
|
||||
result, data=request_data, policy_name=policy_name, original_response=buffered
|
||||
|
|
|
|||
|
|
@ -546,6 +546,21 @@ def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(r
|
|||
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
|
||||
|
|
@ -596,6 +611,17 @@ def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step_o
|
|||
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
|
||||
|
|
@ -644,6 +670,15 @@ def _spans_naming_guardrails(rig: _Rig, since: int, team: str) -> tuple[Span, ..
|
|||
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(
|
||||
|
|
@ -1009,6 +1044,7 @@ def agent(rig: _Rig) -> Iterator[str]:
|
|||
pytest.param(
|
||||
(_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],
|
||||
|
|
@ -1174,6 +1210,7 @@ def test_a_concurrent_burst_names_only_each_requests_own_steps(rig: _Rig) -> Non
|
|||
{
|
||||
_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
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, Iterator
|
||||
from contextlib import AbstractContextManager, nullcontext
|
||||
from copy import deepcopy
|
||||
import logging
|
||||
from typing import Any, Callable, Dict, Final, List
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
|
|
@ -1206,6 +1208,21 @@ def _two_step_pipeline_data(
|
|||
}
|
||||
|
||||
|
||||
def _logging_obj(*, log_parked: bool, call_type: str = "acompletion") -> LiteLLMLoggingObj:
|
||||
logging_obj: Final = LiteLLMLoggingObj(
|
||||
model="m",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type=call_type,
|
||||
start_time=None,
|
||||
litellm_call_id="call-1",
|
||||
function_id="fn-1",
|
||||
)
|
||||
logging_obj.defer_async_logging = True
|
||||
logging_obj.enqueue_deferred_logging = (lambda: None) if log_parked else None
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("first_on_pass", "expected_calls", "expected_header"),
|
||||
|
|
@ -1343,6 +1360,403 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat
|
|||
assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("log_parked", "call_type", "bucket", "expected_after_response"),
|
||||
[
|
||||
pytest.param(True, "acompletion", "metadata", ["first"], id="success-log-waits-for-the-response-hooks"),
|
||||
pytest.param(True, "acompletion", "litellm_metadata", ["first"], id="success-log-waits-litellm-metadata"),
|
||||
pytest.param(False, "acompletion", "metadata", ["first", "second"], id="success-log-already-built"),
|
||||
pytest.param(
|
||||
True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out"
|
||||
),
|
||||
pytest.param(
|
||||
True, "agenerate_content", "metadata", ["first", "second"], id="generate-content-adapter-logged-its-call"
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
log_parked: bool,
|
||||
call_type: str,
|
||||
bucket: str,
|
||||
expected_after_response: list[str],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
request: Final = _two_step_pipeline_data("post_call", "allow")
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": request["model"],
|
||||
"messages": request["messages"],
|
||||
bucket: request["metadata"],
|
||||
"litellm_logging_obj": _logging_obj(log_parked=log_parked, call_type=call_type),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
assert data[bucket]["applied_guardrails"] == ["first", "second"]
|
||||
logged_metadata: Final = data[bucket].copy()
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data[bucket]["applied_guardrails"] == expected_after_response
|
||||
assert logged_metadata["applied_guardrails"] == expected_after_response
|
||||
|
||||
|
||||
def _three_step_two_pipeline_data() -> dict[str, object]:
|
||||
runs_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="first", on_pass="next", on_fail="block"),
|
||||
PipelineStep(guardrail="second", on_pass="next", on_fail="block"),
|
||||
],
|
||||
)
|
||||
stops_before_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[
|
||||
PipelineStep(guardrail="third", on_pass="allow", on_fail="block"),
|
||||
PipelineStep(guardrail="second", on_pass="next", on_fail="block"),
|
||||
],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second", "third"],
|
||||
"_guardrail_pipelines": [("runs-second", runs_second), ("stops-before-second", stops_before_second)],
|
||||
"_pipeline_managed_guardrails": {"first", "second", "third"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
**_three_step_two_pipeline_data(),
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1]
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")]
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={
|
||||
**_three_step_two_pipeline_data(),
|
||||
"stream": True,
|
||||
"litellm_logging_obj": logging_obj,
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
delivered: Final[list[object]] = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
)
|
||||
]
|
||||
|
||||
assert len(delivered) == 2
|
||||
assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1]
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"]
|
||||
|
||||
|
||||
def _ends_before_a_later_pipeline_data(ending: str) -> dict[str, object]:
|
||||
ends: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass=ending)])
|
||||
never_reached: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="second")])
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [("ends", ends), ("never-reached", never_reached)],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("ending", "raised"), [("block", HTTPException), ("modify_response", ModifyResponseException)])
|
||||
async def test_post_call_pipeline_that_ends_the_request_drops_the_steps_of_a_later_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
ending: str,
|
||||
raised: type[Exception],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_ends_before_a_later_pipeline_data(ending), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(raised):
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first"]
|
||||
|
||||
|
||||
def _blocks_before_a_later_request_pipeline_data() -> dict[str, object]:
|
||||
blocks: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass="block")])
|
||||
request: Final = GuardrailPipeline(
|
||||
mode="pre_call",
|
||||
steps=[PipelineStep(guardrail="gate", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")],
|
||||
)
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "gate", "second"],
|
||||
"_guardrail_pipelines": [("blocks", blocks), ("request", request)],
|
||||
"_pipeline_managed_guardrails": {"first", "gate", "second"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_that_blocks_keeps_a_later_request_step_whose_guardrail_runs_outside_pipelines(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
gate: Final = _PipelineStepGuardrail("gate", "pre_call")
|
||||
second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, gate, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_blocks_before_a_later_request_pipeline_data(), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert sorted(data["metadata"]["applied_guardrails"]) == ["first", "gate", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("ending", "stream_end"), [("block", pytest.raises(HTTPException)), ("modify_response", nullcontext())]
|
||||
)
|
||||
async def test_held_stream_pipeline_that_ends_the_stream_drops_the_steps_of_a_later_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
ending: str,
|
||||
stream_end: AbstractContextManager[object],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={**_ends_before_a_later_pipeline_data(ending), "stream": True, "litellm_logging_obj": logging_obj},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
with stream_end:
|
||||
async for _item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
):
|
||||
pass
|
||||
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"])
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("steps", "expected"),
|
||||
[(("first", "second"), ["second"]), (("first",), None)],
|
||||
ids=["skipped-then-ran", "only-step-skipped"],
|
||||
)
|
||||
async def test_post_call_pipeline_drops_a_step_its_guardrail_skips_on_a_whole_response(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
steps: tuple[str, ...],
|
||||
expected: list[str] | None,
|
||||
) -> None:
|
||||
guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in steps]
|
||||
guardrails[0].apply_stream_scope("streaming")
|
||||
monkeypatch.setattr(litellm, "callbacks", guardrails)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
pipeline: Final = GuardrailPipeline(
|
||||
mode="post_call", steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="block") for name in steps]
|
||||
)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": list(steps),
|
||||
"_guardrail_pipelines": [("governance", pipeline)],
|
||||
"_pipeline_managed_guardrails": set(steps),
|
||||
},
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert [guardrail.calls for guardrail in guardrails] == [0, 1][: len(steps)]
|
||||
assert data["metadata"].get("applied_guardrails") == expected
|
||||
|
||||
|
||||
class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
) -> None:
|
||||
self.calls += 1
|
||||
data["litellm_metadata"] = {"user_api_key_hash": "hashed"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={**_two_step_pipeline_data("post_call", "next"), "litellm_logging_obj": _logging_obj(log_parked=True)},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
stops_before_second: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")],
|
||||
)
|
||||
stops_before_first: Final = GuardrailPipeline(
|
||||
mode="post_call",
|
||||
steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")],
|
||||
)
|
||||
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data={
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"guardrails": ["first", "second"],
|
||||
"_guardrail_pipelines": [
|
||||
("stops-before-second", stops_before_second),
|
||||
("stops-before-first", stops_before_first),
|
||||
],
|
||||
"_pipeline_managed_guardrails": {"first", "second"},
|
||||
},
|
||||
"litellm_logging_obj": _logging_obj(log_parked=True),
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert (first.calls, second.calls) == (1, 1)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request(
|
||||
proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -1409,6 +1823,33 @@ async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised(
|
|||
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")
|
||||
|
|
@ -1509,6 +1950,51 @@ async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one
|
|||
assert "applied_guardrails" not in kwargs["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("log_waits", "expected_after_stream"),
|
||||
[(True, ["first"]), (False, ["first", "second"])],
|
||||
ids=["stream-log-waits-for-the-pipeline", "stream-log-not-deferred"],
|
||||
)
|
||||
async def test_held_stream_pipeline_drops_the_steps_that_did_not_run_only_while_its_log_waits(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
log_waits: bool,
|
||||
expected_after_stream: list[str],
|
||||
) -> None:
|
||||
first: Final = _PipelineStepGuardrail("first", "post_call")
|
||||
second: Final = _PipelineStepGuardrail("second", "post_call")
|
||||
monkeypatch.setattr(litellm, "callbacks", [first, second])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
logging_obj: Final = _logging_obj(log_parked=False)
|
||||
|
||||
async def stream_log(*args: object) -> None:
|
||||
return None
|
||||
|
||||
logging_obj._on_deferred_stream_complete = stream_log if log_waits else None
|
||||
auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
|
||||
data: Final = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
data={**_two_step_pipeline_data("post_call", "allow"), "stream": True, "litellm_logging_obj": logging_obj},
|
||||
call_type="completion",
|
||||
)
|
||||
assert data["metadata"]["applied_guardrails"] == ["first", "second"]
|
||||
logged_metadata: Final = data["metadata"].copy()
|
||||
|
||||
delivered: Final[list[object]] = [
|
||||
item
|
||||
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
|
||||
)
|
||||
]
|
||||
|
||||
assert len(delivered) == 2
|
||||
assert (first.calls, second.calls) == (1, 0)
|
||||
assert data["metadata"]["applied_guardrails"] == expected_after_stream
|
||||
assert logged_metadata["applied_guardrails"] == expected_after_stream
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_pipeline_replacement_response_reaches_caller(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
@ -3769,3 +4255,4 @@ async def test_live_pipeline_disconnect_scan_that_cannot_run_records_every_step_
|
|||
assert seen.get("count") is None
|
||||
assert _post_call_statuses(data) == ["guardrail_failed_to_respond"]
|
||||
assert any("cut short" in message and "RuntimeError" in message for message in _warnings(caplog))
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue