diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index bd8b9eb24d7..c56b6364bb8 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1440,7 +1440,7 @@ class _SelfRecordingGuardrail(_TextReturningGuardrail): @pytest.mark.asyncio -async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch): +async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) @@ -1466,7 +1466,9 @@ class _LitellmMetadataOpeningGuardrail(CustomGuardrail): @pytest.mark.asyncio -async def test_streaming_step_is_named_in_the_metadata_the_request_started_with(monkeypatch): +async def test_streaming_step_is_named_in_the_metadata_the_request_started_with( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_LitellmMetadataOpeningGuardrail()]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) @@ -1476,7 +1478,9 @@ async def test_streaming_step_is_named_in_the_metadata_the_request_started_with( @pytest.mark.asyncio -async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(monkeypatch): +async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) 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 29438c63a56..2426ee8b7f0 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1209,8 +1209,11 @@ def _two_step_pipeline_data( [("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])], ) async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( - proxy_logging, make_user_api_key_auth, monkeypatch, first_on_pass, expected_calls, expected_header -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + first_on_pass: str, + expected_calls: tuple[int, int], + expected_header: list[str], +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1227,8 +1230,8 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( @pytest.mark.asyncio async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1264,8 +1267,9 @@ _PIPELINE_STEP_ERRORS: Final = ( @pytest.mark.asyncio @pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on( - proxy_logging, make_user_api_key_auth, monkeypatch, error -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: first = _RaisingPipelineStepGuardrail("first", error) second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1283,8 +1287,9 @@ async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_ @pytest.mark.asyncio @pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( - proxy_logging, make_user_api_key_auth, monkeypatch, error -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: first = _RaisingPipelineStepGuardrail("first", error) second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1317,8 +1322,8 @@ class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): @pytest.mark.asyncio async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _OpensLitellmMetadataGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1337,8 +1342,8 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat @pytest.mark.asyncio async def test_post_call_pipeline_steps_are_named_only_once_they_run( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "post_call") second = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1369,8 +1374,8 @@ class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): @pytest.mark.asyncio async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") second = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1405,8 +1410,12 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ ], ) async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers( - proxy_logging, make_user_api_key_auth, monkeypatch, asks_for, request_route, streams_live, expected_header -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + asks_for: dict[str, object], + request_route: str, + streams_live: bool, + expected_header: list[str] | None, +) -> None: guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 615e39c495d..48dac33d010 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ 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, user_api_key_dict): + def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]