test(guardrails): type the pipeline step test helpers

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
 

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(