mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(guardrails): run post_call policy pipelines on background Responses retrieval
A POST /v1/responses with background: true returns a queued response, so the
post_call pipelines attached at submit time had nothing to inspect. They now
defer on queued and in_progress responses and run on GET /v1/responses/{id}
instead: the retrieval resolves the response id back to its deployment,
re-attaches the policies that governed the original model, and reports them
in the x-litellm-applied-* headers of the retrieval response.
This commit is contained in:
parent
08b60c409a
commit
3caa3b60d5
6 changed files with 485 additions and 43 deletions
|
|
@ -184,6 +184,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
refresh_proxy_server_request_body_snapshot,
|
||||
reject_url_valued_destination,
|
||||
)
|
||||
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
|
|
@ -1883,7 +1884,6 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Calculate request queue time after add_litellm_data_to_request
|
||||
# which sets arrival_time in proxy_server_request. Ends at start_time
|
||||
# (not a freshly captured time.time() here) so this window is exactly
|
||||
|
|
@ -2031,6 +2031,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
call_type=route_type,
|
||||
)
|
||||
if route_type == "aget_responses":
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
# Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may
|
||||
# have mutated `self.data` in place, and the audit-trail snapshot taken in
|
||||
|
|
|
|||
116
litellm/proxy/policy_engine/response_retrieval.py
Normal file
116
litellm/proxy/policy_engine/response_retrieval.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
add_policy_sources_to_metadata,
|
||||
add_policy_to_applied_policies_header,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.proxy.policy_engine import PolicyMatchContext
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router import Router
|
||||
|
||||
PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...]
|
||||
|
||||
_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines)
|
||||
|
||||
|
||||
def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | None:
|
||||
if llm_router is None or not isinstance(response_id, str):
|
||||
return None
|
||||
model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
|
||||
if model_id is None:
|
||||
return None
|
||||
deployment: Final = llm_router.get_deployment(model_id)
|
||||
return deployment.model_name if deployment is not None else None
|
||||
|
||||
|
||||
def _retrieval_context(
|
||||
data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str
|
||||
) -> PolicyMatchContext:
|
||||
team_alias: Final = user_api_key_dict.team_alias
|
||||
key_alias: Final = user_api_key_dict.key_alias
|
||||
return PolicyMatchContext(
|
||||
team_alias=team_alias if isinstance(team_alias, str) else None,
|
||||
key_alias=key_alias if isinstance(key_alias, str) else None,
|
||||
model=model_group,
|
||||
tags=get_tags_from_request_body(data) or None,
|
||||
)
|
||||
|
||||
|
||||
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
|
||||
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
|
||||
if not matches:
|
||||
return (), MappingProxyType({})
|
||||
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
|
||||
policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list
|
||||
context=context,
|
||||
)
|
||||
post_call_pipelines: Final = tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context(
|
||||
context=context, policy_names=applied_policy_names
|
||||
)
|
||||
if pipeline.mode == "post_call"
|
||||
)
|
||||
return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches})
|
||||
|
||||
|
||||
def attach_post_call_pipelines_to_retrieval(
|
||||
data: dict, # mutable-ok: the proxy's request-state dict, written in place like every other policy engine hook
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
llm_router: "Router | None",
|
||||
) -> None:
|
||||
if not get_policy_registry().is_initialized():
|
||||
return
|
||||
model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router)
|
||||
if model_group is None:
|
||||
return
|
||||
context: Final = _retrieval_context(data, user_api_key_dict, model_group)
|
||||
post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context)
|
||||
_, bucket = get_or_create_metadata_bucket(data)
|
||||
already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ())
|
||||
attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached)
|
||||
added: Final = tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in post_call_pipelines
|
||||
if policy_name not in attached_policy_names
|
||||
)
|
||||
if not added:
|
||||
return
|
||||
pipelines: Final = (*already_attached, *added)
|
||||
bucket["_guardrail_pipelines"] = pipelines
|
||||
bucket["_pipeline_managed_guardrails"] = frozenset(
|
||||
step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps
|
||||
)
|
||||
for policy_name, _pipeline in added:
|
||||
add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name)
|
||||
for _policy_name, pipeline in added:
|
||||
for step in pipeline.steps:
|
||||
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
|
||||
add_policy_sources_to_metadata(
|
||||
request_data=data,
|
||||
policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict
|
||||
policy_name: policy_sources[policy_name] for policy_name, _pipeline in added
|
||||
},
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s",
|
||||
data.get("response_id"),
|
||||
model_group,
|
||||
", ".join(policy_name for policy_name, _pipeline in added),
|
||||
)
|
||||
|
|
@ -171,6 +171,7 @@ from litellm.repositories.verification_token_repository import (
|
|||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.mcp import (
|
||||
MCPDuringCallResponseObject,
|
||||
MCPPreCallRequestObject,
|
||||
|
|
@ -525,15 +526,21 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr
|
|||
)
|
||||
|
||||
|
||||
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
|
||||
if data.get("background") is not True:
|
||||
return
|
||||
_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress"))
|
||||
|
||||
|
||||
def _is_pending_background_response(response: LLMResponseTypes) -> bool:
|
||||
return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES
|
||||
|
||||
|
||||
def _log_deferred_post_call_pipelines(data: Mapping[str, object], response: ResponsesAPIResponse) -> None:
|
||||
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
|
||||
if not policy_names:
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"Policies with post_call guardrail pipelines do not run on background responses yet; "
|
||||
"the response is released ungoverned by them: %s",
|
||||
verbose_proxy_logger.debug(
|
||||
"Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s",
|
||||
response.id,
|
||||
response.status,
|
||||
", ".join(policy_names),
|
||||
)
|
||||
|
||||
|
|
@ -1956,8 +1963,6 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
try:
|
||||
_warn_background_skips_post_call_pipelines(data)
|
||||
|
||||
# Execute guardrail pipelines before the normal callback loop
|
||||
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
|
||||
data=data,
|
||||
|
|
@ -2927,6 +2932,24 @@ class ProxyLogging:
|
|||
daemon=True,
|
||||
).start()
|
||||
|
||||
async def _run_post_call_pipelines(
|
||||
self,
|
||||
data: dict, # mutable-ok: same request-payload shape as post_call_success_hook's data
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: LLMResponseTypes,
|
||||
) -> LLMResponseTypes | None:
|
||||
if _is_pending_background_response(response):
|
||||
_log_deferred_post_call_pipelines(data, response)
|
||||
return None
|
||||
_, pipeline_response = await self._maybe_execute_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
|
||||
event_hook="post_call",
|
||||
response=response,
|
||||
)
|
||||
return pipeline_response
|
||||
|
||||
async def post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -2946,11 +2969,9 @@ class ProxyLogging:
|
|||
from litellm.proxy.proxy_server import llm_router
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
_, pipeline_response = await self._maybe_execute_pipelines(
|
||||
pipeline_response: Final = await self._run_post_call_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
|
||||
event_hook="post_call",
|
||||
response=response,
|
||||
)
|
||||
if pipeline_response is not None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,157 @@
|
|||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
|
||||
GOVERNED_MODEL_ID = "deployment-governed"
|
||||
UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini"
|
||||
UNGOVERNED_MODEL_ID = "deployment-ungoverned"
|
||||
|
||||
|
||||
class FakeRouter:
|
||||
def __init__(self, deployments: dict[str, Deployment]):
|
||||
self._deployments = deployments
|
||||
|
||||
def get_deployment(self, model_id: str) -> Deployment | None:
|
||||
return self._deployments.get(model_id)
|
||||
|
||||
|
||||
def _deployment(model_group: str, model_id: str) -> Deployment:
|
||||
return Deployment(
|
||||
model_name=model_group,
|
||||
litellm_params=LiteLLM_Params(model=f"openai/{model_group}"),
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
|
||||
|
||||
def _router() -> FakeRouter:
|
||||
return FakeRouter(
|
||||
{
|
||||
GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID),
|
||||
UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _encoded_response_id(model_id: str) -> str:
|
||||
return ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream"
|
||||
)
|
||||
|
||||
|
||||
def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]:
|
||||
return {
|
||||
"guardrails": {"add": [guardrail]},
|
||||
"pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def policy_engine():
|
||||
policy_registry = get_policy_registry()
|
||||
attachment_registry = get_attachment_registry()
|
||||
policy_registry.load_policies(
|
||||
{
|
||||
"response-governance": _pipeline_policy("output-word-filter"),
|
||||
"input-governance": _pipeline_policy("input-word-filter", mode="pre_call"),
|
||||
"team-governance": _pipeline_policy("team-word-filter"),
|
||||
}
|
||||
)
|
||||
attachment_registry.load_attachments(
|
||||
[
|
||||
{"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]},
|
||||
{"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]},
|
||||
{"policy": "team-governance", "teams": ["governed-team"]},
|
||||
]
|
||||
)
|
||||
yield
|
||||
policy_registry.clear()
|
||||
attachment_registry.clear()
|
||||
|
||||
|
||||
def _retrieval_data(model_id: str) -> dict:
|
||||
return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}}
|
||||
|
||||
|
||||
def _attached_pipelines(data: dict) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(
|
||||
(policy_name, ",".join(step.guardrail for step in pipeline.steps))
|
||||
for policy_name, pipeline in data["litellm_metadata"]["_guardrail_pipelines"]
|
||||
)
|
||||
|
||||
|
||||
def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine):
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"})
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"]
|
||||
assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"}
|
||||
assert "model" not in data
|
||||
assert "guardrails" not in data["litellm_metadata"]
|
||||
|
||||
|
||||
def test_key_and_team_context_also_governs_retrieval(policy_engine):
|
||||
data = _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router()
|
||||
)
|
||||
|
||||
assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),)
|
||||
|
||||
|
||||
def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine):
|
||||
data = _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
|
||||
def test_already_attached_policy_is_not_attached_twice(policy_engine):
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
router = _router()
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"response_id",
|
||||
["resp_plain_upstream_id", _encoded_response_id("deployment-missing-from-router"), None],
|
||||
)
|
||||
def test_unresolvable_response_id_attaches_nothing(policy_engine, response_id):
|
||||
data = {"response_id": response_id, "litellm_metadata": {}}
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == {"response_id": response_id, "litellm_metadata": {}}
|
||||
|
||||
|
||||
def test_without_a_router_attaches_nothing(policy_engine):
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None)
|
||||
|
||||
assert data == _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
|
||||
def test_without_policy_engine_attaches_nothing():
|
||||
get_policy_registry().clear()
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
|
@ -8259,3 +8259,112 @@ class TestPassthroughHeadersAcceptImmutableMappings:
|
|||
assert merged["content-type"] == "text/event-stream"
|
||||
# the excluded hop-by-hop header is still dropped
|
||||
assert "transfer-encoding" not in merged
|
||||
|
||||
|
||||
class TestBackgroundResponseRetrievalGovernance:
|
||||
"""LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines."""
|
||||
|
||||
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
|
||||
GOVERNED_MODEL_ID = "deployment-governed"
|
||||
|
||||
@pytest.fixture
|
||||
def policy_engine(self):
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
get_policy_registry().load_policies(
|
||||
{
|
||||
"response-governance": {
|
||||
"guardrails": {"add": ["output-word-filter"]},
|
||||
"pipeline": {
|
||||
"mode": "post_call",
|
||||
"steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}],
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
get_attachment_registry().load_attachments(
|
||||
[{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}]
|
||||
)
|
||||
yield
|
||||
get_policy_registry().clear()
|
||||
get_attachment_registry().clear()
|
||||
|
||||
def _router(self) -> MagicMock:
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
router = MagicMock()
|
||||
router.get_deployment.side_effect = lambda model_id: (
|
||||
Deployment(
|
||||
model_name=self.GOVERNED_MODEL_GROUP,
|
||||
litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"),
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
if model_id == self.GOVERNED_MODEL_ID
|
||||
else None
|
||||
)
|
||||
return router
|
||||
|
||||
async def _pre_call(self, route_type: str, monkeypatch) -> dict:
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
client_facing_response_id = "resp_opaque-client-facing-id"
|
||||
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream"
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"response_id": client_facing_response_id, "litellm_metadata": {}}
|
||||
)
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
async def passthrough_add_litellm_data_to_request(data, **kwargs):
|
||||
return data
|
||||
|
||||
async def decrypting_pre_call_hook(user_api_key_dict, data, call_type):
|
||||
if data.get("response_id") == client_facing_response_id:
|
||||
data["response_id"] = encoded_response_id
|
||||
return data
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing,
|
||||
"add_litellm_data_to_request",
|
||||
passthrough_add_litellm_data_to_request,
|
||||
)
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook)
|
||||
proxy_config = MagicMock(spec=ProxyConfig)
|
||||
proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None)
|
||||
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
|
||||
request=mock_request,
|
||||
general_settings={},
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
proxy_config=proxy_config,
|
||||
route_type=route_type,
|
||||
llm_router=self._router(),
|
||||
)
|
||||
return returned_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline(
|
||||
self, policy_engine, monkeypatch
|
||||
):
|
||||
data = await self._pre_call("aget_responses", monkeypatch)
|
||||
|
||||
assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f")
|
||||
pipelines = data["litellm_metadata"]["_guardrail_pipelines"]
|
||||
assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [
|
||||
("response-governance", ["output-word-filter"])
|
||||
]
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
assert data["model"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id(
|
||||
self, policy_engine, monkeypatch
|
||||
):
|
||||
data = await self._pre_call("aresponses", monkeypatch)
|
||||
|
||||
assert "_guardrail_pipelines" not in data["litellm_metadata"]
|
||||
assert "applied_policies" not in data["litellm_metadata"]
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_g
|
|||
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import (
|
||||
GuardrailPipeline,
|
||||
PipelineStep,
|
||||
|
|
@ -1419,8 +1420,69 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver
|
|||
assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def _background_response(status: str, text: str = "") -> ResponsesAPIResponse:
|
||||
output = (
|
||||
[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else []
|
||||
)
|
||||
return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status)
|
||||
|
||||
|
||||
def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]:
|
||||
class OutputBlockingGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["response"] = response
|
||||
raise HTTPException(status_code=400, detail={"error": "output blocked"})
|
||||
|
||||
return [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
|
||||
@pytest.mark.parametrize("pending_status", ["queued", "in_progress"])
|
||||
async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog, pending_status
|
||||
):
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(background=True)
|
||||
response = _background_response(pending_status)
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.post_call_success_hook(
|
||||
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert out is response
|
||||
assert "response" not in seen
|
||||
assert not _warnings(caplog)
|
||||
assert any(
|
||||
"response-governance" in record.getMessage() and pending_status in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("final_status", ["completed", "incomplete"])
|
||||
async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, final_status
|
||||
):
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data()
|
||||
response = _background_response(final_status, text="kumquat")
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert info.value.detail["error"] == "output blocked"
|
||||
assert seen["response"] is response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
|
@ -1436,36 +1498,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
|
|||
|
||||
assert out is not None
|
||||
assert out.get("background") is True
|
||||
assert any("response-governance" in message and "background" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"background": True,
|
||||
"metadata": {
|
||||
"_guardrail_pipelines": [("request-governance", pre_call)],
|
||||
"_pipeline_managed_guardrails": {"gr-post"},
|
||||
},
|
||||
}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
guardrails_only=True,
|
||||
)
|
||||
|
||||
assert out is not None
|
||||
assert not any("background" in message for message in _warnings(caplog))
|
||||
assert not _warnings(caplog)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue