diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index 5bc1b9df23a..bf1b7bba876 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -31,7 +31,7 @@ from RestrictedPython.Eval import default_guarded_getitem from RestrictedPython.Guards import ( full_write_guard, guarded_iter_unpack_sequence, - guarded_unpack_sequence, + guarded_unpack_sequence, # pyright: ignore[reportUnknownVariableType] # RestrictedPython is untyped safer_getattr, ) from RestrictedPython.transformer import copy_locations diff --git a/tests/unit/proxy/guardrails/test_custom_code_security.py b/tests/unit/proxy/guardrails/test_custom_code_security.py index 7ad2f96363f..dc959137b43 100644 --- a/tests/unit/proxy/guardrails/test_custom_code_security.py +++ b/tests/unit/proxy/guardrails/test_custom_code_security.py @@ -2,6 +2,7 @@ import asyncio import http.server import threading import time +from collections.abc import AsyncIterator from http.server import ThreadingHTTPServer import pytest @@ -953,7 +954,7 @@ def test_async_for_gets_the_same_guards_as_for(): assert async_unpack & {"_agetiter_", "_aiter_unpack_sequence_"} -async def _arange(*values): +async def _arange(*values: object) -> AsyncIterator[object]: for v in values: yield v @@ -974,7 +975,7 @@ async def _arange(*values): ), ], ) -async def test_async_for_still_executes(source, arg, expected): +async def test_async_for_still_executes(source: str, arg: tuple[object, ...], expected: int): """Guarding `async for` must not break it, in either target shape.""" from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals,