From 5bd75f813e72b5a634b3e169f974aae8242567b8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:44:03 -0700 Subject: [PATCH] test(guardrails): type the pipeline step test helpers --- .../policy_engine/test_pipeline_executor.py | 13 ++++-- .../proxy_logging/test_guardrail_pipeline.py | 40 ++++++++++--------- 2 files changed, 32 insertions(+), 21 deletions(-) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index 068102aac08..bd8b9eb24d7 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -15,6 +15,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( @@ -30,7 +31,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs try: from fastapi.exceptions import HTTPException @@ -1450,10 +1451,16 @@ async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(m class _LitellmMetadataOpeningGuardrail(CustomGuardrail): """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does.""" - def __init__(self): + def __init__(self) -> None: super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: request_data["litellm_metadata"] = {"user_api_key_hash": "hashed"} return inputs diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index c8d19c465f8..29438c63a56 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -39,7 +39,8 @@ from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_inf from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks -from litellm.types.utils import CallTypes +from litellm.caching.dual_cache import DualCache +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -1169,20 +1170,26 @@ class _PipelineStepGuardrail(CustomGuardrail): super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False) self.calls = 0 - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 return data - async def async_post_call_success_hook(self, data, user_api_key_dict, response): + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: self.calls += 1 -def _two_step_pipeline_data(mode: str, first_on_pass: str, first_on_fail: str = "block") -> dict[str, Any]: +def _two_step_pipeline_data( + mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block" +) -> dict[str, object]: pipeline = GuardrailPipeline( mode=mode, steps=[ PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), - PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail=second_on_fail), ], ) return { @@ -1241,7 +1248,9 @@ class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): super().__init__(guardrail_name, "pre_call") self.error = error - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 raise self.error @@ -1297,7 +1306,9 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 data["litellm_metadata"] = {} add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -1349,7 +1360,9 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run( class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): - async def async_post_call_success_hook(self, data, user_api_key_dict, response): + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: self.calls += 1 data["litellm_metadata"] = {"user_api_key_hash": "hashed"} @@ -1398,16 +1411,7 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live monkeypatch.setattr(litellm, "callbacks", guardrails) - data = _two_step_pipeline_data("post_call", "next") - data["metadata"]["_guardrail_pipelines"] = [ - ( - "governance", - GuardrailPipeline( - mode="post_call", - steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="next") for name in ("first", "second")], - ), - ) - ] + data = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") data.update(asks_for) result = await proxy_logging.pre_call_hook(