mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
fix(guardrails): block private destinations in custom code http_request and bound guardrail execution time (#43280)
* fix(guardrails): block private destinations in custom code http_request and bound guardrail execution time * fix(guardrails): keep startup fail-closed on a custom code compile error and report a load timeout on the test endpoint A compile failure is no longer a ValueError, so a config-file custom code guardrail that does not compile stops the proxy at startup as it did before, while POST /guardrails catches it by name and still rolls back. The admin test endpoint reports a module-level timeout as an execution timeout instead of a compile error, a caller-supplied Host header is stripped from http_* requests while validation is on, and GET keeps the shared client's connect timeout. * test(guardrails): cover the http_request methods, header passthrough and cancellation paths --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
5ac640e49d
commit
dfb5d905ea
10 changed files with 1334 additions and 332 deletions
|
|
@ -2,7 +2,7 @@
|
|||
CRUD ENDPOINTS FOR GUARDRAILS
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
|
|
@ -22,6 +22,12 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
|||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.path_utils import safe_join
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import (
|
||||
ExecutionTimeoutError,
|
||||
await_with_timeout,
|
||||
call_off_loop_with_timeout,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeCompilationError
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
|
||||
build_sandbox_globals,
|
||||
compile_sandboxed,
|
||||
|
|
@ -401,7 +407,7 @@ async def create_guardrail(
|
|||
verbose_proxy_logger.info(
|
||||
"Immediate sync: Successfully initialized guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
|
||||
)
|
||||
except (ValueError, TypeError) as init_error:
|
||||
except (ValueError, TypeError, CustomCodeCompilationError) as init_error:
|
||||
# Configuration error — roll back the DB write so the guardrail isn't orphaned
|
||||
if prisma_client is not None:
|
||||
try:
|
||||
|
|
@ -421,6 +427,8 @@ async def create_guardrail(
|
|||
)
|
||||
|
||||
return result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error adding guardrail to db: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
|
@ -2124,15 +2132,20 @@ async def test_custom_code_guardrail(
|
|||
try:
|
||||
exec_globals: Final = build_sandbox_globals()
|
||||
|
||||
try:
|
||||
def load_module() -> None:
|
||||
compiled: Final[CodeType] = compile_sandboxed(request.custom_code)
|
||||
exec(compiled, exec_globals) # noqa: S102
|
||||
|
||||
try:
|
||||
await call_off_loop_with_timeout(load_module, EXECUTION_TIMEOUT_SECONDS, label="test:load")
|
||||
except SyntaxError as e:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
error=f"Syntax error in custom code: {e}",
|
||||
error_type="compilation",
|
||||
)
|
||||
except ExecutionTimeoutError:
|
||||
return _execution_timeout_response(EXECUTION_TIMEOUT_SECONDS)
|
||||
except Exception as e:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
|
|
@ -2178,16 +2191,9 @@ async def test_custom_code_guardrail(
|
|||
return apply_fn(test_inputs, safe_request_data, request.input_type)
|
||||
|
||||
try:
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future: Final = executor.submit(execute_guardrail)
|
||||
try:
|
||||
result: Final = future.result(timeout=EXECUTION_TIMEOUT_SECONDS)
|
||||
except concurrent.futures.TimeoutError:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
error=f"Execution timeout: code took longer than {EXECUTION_TIMEOUT_SECONDS} seconds",
|
||||
error_type="execution",
|
||||
)
|
||||
result: Final = await _run_test_guardrail(execute_guardrail, EXECUTION_TIMEOUT_SECONDS)
|
||||
except ExecutionTimeoutError:
|
||||
return _execution_timeout_response(EXECUTION_TIMEOUT_SECONDS)
|
||||
except Exception as e:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
|
|
@ -2219,6 +2225,23 @@ async def test_custom_code_guardrail(
|
|||
)
|
||||
|
||||
|
||||
def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
error=f"Execution timeout: code took longer than {timeout:g} seconds",
|
||||
error_type="execution",
|
||||
)
|
||||
|
||||
|
||||
async def _run_test_guardrail(execute_guardrail: Callable[[], object], timeout: float) -> object:
|
||||
deadline: Final = asyncio.get_running_loop().time() + timeout
|
||||
raw_result: Final = await call_off_loop_with_timeout(execute_guardrail, timeout, label="test")
|
||||
if not inspect.iscoroutine(raw_result):
|
||||
return raw_result
|
||||
remaining: Final = max(deadline - asyncio.get_running_loop().time(), 0.0)
|
||||
return await await_with_timeout(raw_result, remaining, label="test")
|
||||
|
||||
|
||||
def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: str) -> Literal["request", "response"]:
|
||||
"""Return the effective input_type, auto-upgrading to 'response' for post_call guardrails."""
|
||||
if input_type == "request":
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
custom_code_guardrail: Final = CustomCodeGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
custom_code=custom_code,
|
||||
execution_timeout=litellm_params.timeout,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,231 @@
|
|||
"""Wall-clock bounds for sandboxed guardrail code.
|
||||
|
||||
Sync guardrail code runs on a dedicated daemon thread so a runaway loop never stalls the event loop; async
|
||||
guardrail code is awaited as its own task. Either way the code's deadline is published through a context
|
||||
variable, and the sandbox compiler routes every ``while`` test, ``for`` iteration and comprehension through
|
||||
:func:`budget_ok`, which raises ``ExecutionInterrupted`` once that deadline has passed, whatever the code
|
||||
catches around the loop body. As a backstop, a worker thread still running at the deadline has
|
||||
``ExecutionInterrupted`` injected with ``PyThreadState_SetAsyncExc`` and a task still running is cancelled
|
||||
repeatedly. A long-running C call (a catastrophic regex, for one) only sees any of this once it returns, so
|
||||
the caller still gets its timeout on schedule while the worker keeps burning CPU until that call ends.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import contextvars
|
||||
import ctypes
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Generic, TypeVar
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
T: Final = TypeVar("T")
|
||||
|
||||
_INTERRUPT_GRACE_SECONDS: Final = 1.0
|
||||
_INTERRUPT_POLL_SECONDS: Final = 0.05
|
||||
_deadline: Final[contextvars.ContextVar[float | None]] = contextvars.ContextVar("guardrail_code_deadline", default=None)
|
||||
|
||||
|
||||
class ExecutionInterrupted(BaseException):
|
||||
"""Raised inside guardrail code once its budget is spent; a BaseException so sandboxed
|
||||
``except Exception`` clauses cannot swallow it."""
|
||||
|
||||
|
||||
class ExecutionTimeoutError(Exception):
|
||||
"""The guardrail code did not finish within its wall-clock budget."""
|
||||
|
||||
def __init__(self, timeout: float) -> None:
|
||||
super().__init__(f"exceeded the {timeout:g}s execution timeout")
|
||||
self.timeout: Final = timeout
|
||||
|
||||
|
||||
class SandboxExit(Exception):
|
||||
"""Sandboxed code raised something outside the ``Exception`` tree (``SystemExit``, ``KeyboardInterrupt``,
|
||||
a bare ``BaseException``). It is delivered as an ordinary exception so it can neither stop the event loop
|
||||
nor pass for a timeout."""
|
||||
|
||||
def __init__(self, cause: BaseException) -> None:
|
||||
super().__init__(f"{type(cause).__name__}: {cause}")
|
||||
|
||||
|
||||
def _past_deadline() -> bool:
|
||||
deadline: Final = _deadline.get()
|
||||
return deadline is not None and time.monotonic() > deadline
|
||||
|
||||
|
||||
def budget_ok() -> bool:
|
||||
"""Bound to ``_budget_ok_`` in the sandbox, where every ``while`` test starts with a call to it."""
|
||||
if _past_deadline():
|
||||
raise ExecutionInterrupted
|
||||
return True
|
||||
|
||||
|
||||
def budgeted_iter(iterable: Iterable[T]) -> Iterator[T]:
|
||||
"""Bound to ``_getiter_`` in the sandbox, so every ``for`` loop and comprehension checks the budget per item."""
|
||||
for item in iterable:
|
||||
budget_ok()
|
||||
yield item
|
||||
|
||||
|
||||
class _InterruptGate:
|
||||
"""Aims the interrupt at the worker thread only while it is inside the sandboxed call, so a thread id the
|
||||
OS recycles after the worker exits is never hit."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._thread_id: int | None = None
|
||||
|
||||
def open(self) -> None:
|
||||
self._thread_id = threading.get_ident()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self._thread_id = None
|
||||
|
||||
def is_open(self) -> bool:
|
||||
return self._thread_id is not None
|
||||
|
||||
def interrupt(self) -> bool:
|
||||
with self._lock:
|
||||
if self._thread_id is None:
|
||||
return False
|
||||
ctypes.pythonapi.PyThreadState_SetAsyncExc(
|
||||
ctypes.c_ulong(self._thread_id), ctypes.py_object(ExecutionInterrupted)
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Worker(Generic[T]):
|
||||
thread: threading.Thread
|
||||
outcome: concurrent.futures.Future[T]
|
||||
gate: _InterruptGate
|
||||
|
||||
|
||||
def _run(fn: Callable[[], T], timeout: float, gate: _InterruptGate) -> tuple[T | None, Exception | None]:
|
||||
_deadline.set(time.monotonic() + timeout)
|
||||
gate.open()
|
||||
try:
|
||||
result: Final = fn()
|
||||
except Exception as e: # noqa: BLE001 # every failure is handed to the waiting caller through the future
|
||||
return None, e
|
||||
except ExecutionInterrupted:
|
||||
return None, ExecutionTimeoutError(timeout)
|
||||
except BaseException as e: # noqa: BLE001 # a SystemExit must reach the caller as a failure, not end the worker silently
|
||||
return None, SandboxExit(e)
|
||||
finally:
|
||||
gate.close()
|
||||
if _past_deadline():
|
||||
return None, ExecutionTimeoutError(timeout)
|
||||
return result, None
|
||||
|
||||
|
||||
def _deliver(fn: Callable[[], T], timeout: float, outcome: concurrent.futures.Future[T], gate: _InterruptGate) -> None:
|
||||
try:
|
||||
_settle(outcome, *_run(fn, timeout, gate))
|
||||
except ExecutionInterrupted:
|
||||
_settle(outcome, exception=ExecutionTimeoutError(timeout))
|
||||
|
||||
|
||||
def _settle(outcome: concurrent.futures.Future[T], result: T | None = None, exception: Exception | None = None) -> None:
|
||||
try:
|
||||
if exception is not None:
|
||||
outcome.set_exception(exception)
|
||||
else:
|
||||
outcome.set_result(result) # pyright: ignore[reportArgumentType] # result is T whenever exception is None
|
||||
except concurrent.futures.InvalidStateError:
|
||||
return
|
||||
|
||||
|
||||
def _start_worker(fn: Callable[[], T], timeout: float, label: str) -> _Worker[T]:
|
||||
outcome: Final[concurrent.futures.Future[T]] = concurrent.futures.Future()
|
||||
gate: Final = _InterruptGate()
|
||||
thread: Final = threading.Thread(
|
||||
target=_deliver, args=(fn, timeout, outcome, gate), name=f"guardrail-code:{label}", daemon=True
|
||||
)
|
||||
thread.start()
|
||||
return _Worker(thread, outcome, gate)
|
||||
|
||||
|
||||
def _interrupt(worker: _Worker[T]) -> None:
|
||||
deadline: Final = time.monotonic() + _INTERRUPT_GRACE_SECONDS
|
||||
while worker.gate.interrupt() and time.monotonic() < deadline:
|
||||
worker.thread.join(_INTERRUPT_POLL_SECONDS)
|
||||
if worker.gate.is_open():
|
||||
verbose_proxy_logger.error(
|
||||
"%s is still running after its timeout; it is stuck in a call Python cannot interrupt", worker.thread.name
|
||||
)
|
||||
|
||||
|
||||
def call_with_timeout(fn: Callable[[], T], timeout: float, label: str) -> T:
|
||||
"""Run ``fn`` on a worker thread and wait for it, from sync code."""
|
||||
worker: Final = _start_worker(fn, timeout, label)
|
||||
try:
|
||||
return worker.outcome.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
worker.outcome.cancel()
|
||||
_interrupt(worker)
|
||||
raise ExecutionTimeoutError(timeout) from None
|
||||
|
||||
|
||||
async def call_off_loop_with_timeout(fn: Callable[[], T], timeout: float, label: str) -> T:
|
||||
"""Run ``fn`` on a worker thread and await it without blocking the event loop."""
|
||||
worker: Final = _start_worker(fn, timeout, label)
|
||||
try:
|
||||
return await asyncio.wait_for(asyncio.wrap_future(worker.outcome), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
await asyncio.to_thread(_interrupt, worker)
|
||||
raise ExecutionTimeoutError(timeout) from None
|
||||
except asyncio.CancelledError:
|
||||
threading.Thread(target=_interrupt, args=(worker,), name=f"guardrail-interrupt:{label}", daemon=True).start()
|
||||
raise
|
||||
|
||||
|
||||
def _discard_outcome(task: asyncio.Future[T]) -> None:
|
||||
if not task.cancelled():
|
||||
task.exception()
|
||||
|
||||
|
||||
async def _cancel(task: asyncio.Task[T], label: str) -> None:
|
||||
deadline: Final = time.monotonic() + _INTERRUPT_GRACE_SECONDS
|
||||
while not task.done() and time.monotonic() < deadline:
|
||||
task.cancel()
|
||||
await asyncio.wait((task,), timeout=_INTERRUPT_POLL_SECONDS)
|
||||
if not task.done():
|
||||
verbose_proxy_logger.error(
|
||||
"guardrail-code:%s is still running after its timeout; it keeps swallowing cancellation", label
|
||||
)
|
||||
|
||||
|
||||
async def _contain(pending: Awaitable[T], timeout: float) -> T:
|
||||
try:
|
||||
result: Final = await pending
|
||||
except (Exception, asyncio.CancelledError):
|
||||
raise
|
||||
except ExecutionInterrupted:
|
||||
raise ExecutionTimeoutError(timeout) from None
|
||||
except BaseException as e: # noqa: BLE001 # a SystemExit escaping a task stops the whole event loop
|
||||
raise SandboxExit(e) from e
|
||||
if _past_deadline():
|
||||
raise ExecutionTimeoutError(timeout)
|
||||
return result
|
||||
|
||||
|
||||
async def await_with_timeout(pending: Awaitable[object], timeout: float, label: str) -> object:
|
||||
"""Await ``pending`` on the event loop and give it up at the deadline, even if it swallows cancellation."""
|
||||
context: Final = contextvars.copy_context()
|
||||
context.run(_deadline.set, time.monotonic() + timeout)
|
||||
task: Final = context.run(asyncio.ensure_future, _contain(pending, timeout))
|
||||
task.add_done_callback(_discard_outcome)
|
||||
try:
|
||||
await asyncio.wait((task,), timeout=timeout)
|
||||
except asyncio.CancelledError:
|
||||
task.cancel()
|
||||
raise
|
||||
if task.done():
|
||||
return task.result()
|
||||
await _cancel(task, label)
|
||||
raise ExecutionTimeoutError(timeout)
|
||||
|
|
@ -35,12 +35,16 @@ Example: block when response rejects the user (input_type response only):
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import inspect
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, cast
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import Field
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -53,11 +57,19 @@ from litellm.types.guardrails import GuardrailEventHooks
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
from .bounded_execution import (
|
||||
ExecutionTimeoutError,
|
||||
await_with_timeout,
|
||||
call_off_loop_with_timeout,
|
||||
call_with_timeout,
|
||||
)
|
||||
from .sandbox import build_sandbox_globals, compile_sandboxed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
DEFAULT_EXECUTION_TIMEOUT_SECONDS: Final = 30.0
|
||||
|
||||
|
||||
def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]:
|
||||
bucket: Final = request_data.get(key)
|
||||
|
|
@ -73,7 +85,8 @@ class CustomCodeGuardrailError(Exception):
|
|||
|
||||
|
||||
class CustomCodeCompilationError(CustomCodeGuardrailError):
|
||||
"""Raised when custom code fails to compile."""
|
||||
"""Raised when custom code fails to compile. Deliberately not a ValueError: a config-file guardrail whose
|
||||
code does not compile must stop startup instead of being skipped, so the guardrail endpoints catch it by name."""
|
||||
|
||||
|
||||
class CustomCodeExecutionError(CustomCodeGuardrailError):
|
||||
|
|
@ -90,6 +103,15 @@ class CustomCodeGuardrailConfigModel(GuardrailConfigModel):
|
|||
custom_code: str
|
||||
"""The Python-like code containing the apply_guardrail function."""
|
||||
|
||||
timeout: float | None = Field(
|
||||
default=DEFAULT_EXECUTION_TIMEOUT_SECONDS,
|
||||
gt=0.0,
|
||||
description=(
|
||||
"Wall-clock limit in seconds for one run of apply_guardrail, module-level code included. "
|
||||
"A run that exceeds it fails the request instead of stalling the proxy."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CustomCodeGuardrail(CustomGuardrail):
|
||||
"""
|
||||
|
|
@ -97,7 +119,8 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
|
||||
The code runs in a sandboxed environment that provides:
|
||||
- Access to LiteLLM primitives (regex_match, json_parse, etc.)
|
||||
- No file I/O or network access
|
||||
- No file I/O; network access only through `http_get`/`http_post`/`http_request`, which refuse
|
||||
private, link-local and loopback destinations unless the host is allowlisted
|
||||
- No imports allowed
|
||||
|
||||
Users write an `apply_guardrail(inputs, request_data, input_type)` function
|
||||
|
|
@ -119,6 +142,7 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
self,
|
||||
custom_code: str,
|
||||
guardrail_name: str | None = "custom_code",
|
||||
execution_timeout: float | None = None,
|
||||
**kwargs: Unpack[_CustomGuardrailOptions],
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -127,9 +151,15 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
Args:
|
||||
custom_code: The source code containing apply_guardrail function
|
||||
guardrail_name: Name of this guardrail instance
|
||||
execution_timeout: Wall-clock budget in seconds for one run of the code
|
||||
**kwargs: Additional arguments passed to CustomGuardrail
|
||||
"""
|
||||
if execution_timeout is not None and not execution_timeout > 0:
|
||||
raise ValueError(f"execution_timeout must be positive, got {execution_timeout}")
|
||||
self.custom_code: str = custom_code
|
||||
self.execution_timeout: float = (
|
||||
DEFAULT_EXECUTION_TIMEOUT_SECONDS if execution_timeout is None else execution_timeout
|
||||
)
|
||||
self._compiled_function: Callable[..., object] | None = None
|
||||
self._compile_lock = threading.Lock()
|
||||
self._compile_error: str | None = None
|
||||
|
|
@ -163,7 +193,11 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
"""Internal compilation method without lock. Expected to run inside _compile_lock."""
|
||||
exec_globals: Final = build_sandbox_globals()
|
||||
compiled: Final = compile_sandboxed(self.custom_code)
|
||||
exec(compiled, exec_globals) # noqa: S102
|
||||
|
||||
def load_module() -> None:
|
||||
exec(compiled, exec_globals) # noqa: S102
|
||||
|
||||
call_with_timeout(load_module, self.execution_timeout, label=f"{self.guardrail_name}:load")
|
||||
|
||||
if "apply_guardrail" not in exec_globals:
|
||||
raise CustomCodeCompilationError(
|
||||
|
|
@ -241,18 +275,10 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
|
||||
start_time: Final = time.time()
|
||||
try:
|
||||
# Prepare inputs dict for the function
|
||||
|
||||
# Prepare request_data with safe subset of information
|
||||
safe_request_data: Final = self._prepare_safe_request_data(request_data)
|
||||
|
||||
# Execute the custom function - handle both sync and async functions
|
||||
raw_result: Final = self._compiled_function(inputs, safe_request_data, input_type)
|
||||
|
||||
# If the function is async (returns a coroutine), await it
|
||||
resolved_result: Final[object] = await raw_result if asyncio.iscoroutine(raw_result) else raw_result
|
||||
|
||||
# Process the result
|
||||
resolved_result: Final = await self._call_compiled(
|
||||
self._compiled_function, inputs, safe_request_data, input_type
|
||||
)
|
||||
return self._process_result(
|
||||
result=resolved_result,
|
||||
inputs=inputs,
|
||||
|
|
@ -267,6 +293,19 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
except ModifyResponseException:
|
||||
# Pre-call block uses passthrough; must not wrap as execution error (500)
|
||||
raise
|
||||
except ExecutionTimeoutError:
|
||||
verbose_proxy_logger.error(
|
||||
"Custom code guardrail '%s' exceeded its %gs execution timeout",
|
||||
self.guardrail_name,
|
||||
self.execution_timeout,
|
||||
)
|
||||
raise CustomCodeExecutionError(
|
||||
f"Custom code guardrail '{self.guardrail_name}' exceeded its "
|
||||
f"{self.execution_timeout:g}s execution timeout",
|
||||
details=MappingProxyType(
|
||||
{"guardrail_name": self.guardrail_name, "input_type": input_type, "timeout": self.execution_timeout}
|
||||
),
|
||||
) from None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Custom code guardrail '%s' execution error: %s", self.guardrail_name, e)
|
||||
raise CustomCodeExecutionError(
|
||||
|
|
@ -277,6 +316,31 @@ class CustomCodeGuardrail(CustomGuardrail):
|
|||
},
|
||||
) from e
|
||||
|
||||
async def _call_compiled(
|
||||
self,
|
||||
compiled_function: Callable[..., object],
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
safe_request_data: Mapping[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
) -> object:
|
||||
"""Run the user's function under the execution budget.
|
||||
|
||||
A coroutine function is awaited on the event loop, so the budget bounds it at its
|
||||
await points and it is given up at the deadline even if it swallows cancellation. A
|
||||
plain function runs on a worker thread, which keeps a busy loop from stalling every
|
||||
other request and lets the runner interrupt it at the deadline.
|
||||
"""
|
||||
label: Final = str(self.guardrail_name)
|
||||
if inspect.iscoroutinefunction(compiled_function):
|
||||
pending: Final = compiled_function(inputs, safe_request_data, input_type)
|
||||
return await await_with_timeout(pending, self.execution_timeout, label)
|
||||
call: Final = functools.partial(compiled_function, inputs, safe_request_data, input_type)
|
||||
deadline: Final = time.monotonic() + self.execution_timeout
|
||||
raw_result: Final = await call_off_loop_with_timeout(call, self.execution_timeout, label)
|
||||
if not asyncio.iscoroutine(raw_result):
|
||||
return raw_result
|
||||
return await await_with_timeout(raw_result, max(deadline - time.monotonic(), 0.0), label)
|
||||
|
||||
def _prepare_safe_request_data(self, request_data: Mapping[str, object]) -> dict[str, object]:
|
||||
"""
|
||||
Prepare a safe subset of request_data for code execution.
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ These functions are injected into the custom code execution environment
|
|||
and provide safe, sandboxed functionality for common guardrail operations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
|
@ -15,7 +16,9 @@ import httpx
|
|||
from pydantic import JsonValue
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, validate_url
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
|
|
@ -393,6 +396,8 @@ _HTTP_DEFAULT_TIMEOUT: Final = 30.0
|
|||
# Maximum allowed timeout (in seconds)
|
||||
_HTTP_MAX_TIMEOUT: Final = 60.0
|
||||
|
||||
_HTTP_ALLOWED_METHODS: Final = ("GET", "POST", "PUT", "DELETE", "PATCH")
|
||||
|
||||
|
||||
class HttpResponseResult(TypedDict):
|
||||
"""Outcome of an HTTP primitive call, as handed back to custom code."""
|
||||
|
|
@ -463,6 +468,11 @@ async def http_request(
|
|||
Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling
|
||||
and better performance.
|
||||
|
||||
Destinations go through LiteLLM's SSRF validation: private, link-local,
|
||||
loopback and cloud-metadata addresses are refused (every redirect hop
|
||||
included) unless the host is listed in ``litellm_settings.user_url_allowed_hosts``
|
||||
or ``litellm_settings.user_url_validation`` is turned off.
|
||||
|
||||
Args:
|
||||
url: The URL to request
|
||||
method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET.
|
||||
|
|
@ -492,35 +502,35 @@ async def http_request(
|
|||
body={"text": "content to check"}
|
||||
)
|
||||
"""
|
||||
# Validate URL
|
||||
if not is_valid_url(url):
|
||||
return _http_error_response(f"Invalid URL: {url}")
|
||||
|
||||
# Validate and normalize method
|
||||
method = method.upper()
|
||||
allowed_methods: Final = {"GET", "POST", "PUT", "DELETE", "PATCH"}
|
||||
if method not in allowed_methods:
|
||||
return _http_error_response(f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}")
|
||||
normalized_method: Final = method.upper()
|
||||
if normalized_method not in _HTTP_ALLOWED_METHODS:
|
||||
return _http_error_response(
|
||||
f"Invalid HTTP method: {normalized_method}. Allowed: {', '.join(_HTTP_ALLOWED_METHODS)}"
|
||||
)
|
||||
|
||||
# Apply timeout limits
|
||||
if timeout is None:
|
||||
timeout = _HTTP_DEFAULT_TIMEOUT
|
||||
else:
|
||||
timeout = min(max(0.1, timeout), _HTTP_MAX_TIMEOUT)
|
||||
effective_timeout: Final = _HTTP_DEFAULT_TIMEOUT if timeout is None else min(max(0.1, timeout), _HTTP_MAX_TIMEOUT)
|
||||
|
||||
# Get the global cached async HTTP client
|
||||
client: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
||||
params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)},
|
||||
params={
|
||||
"timeout": httpx.Timeout(timeout=effective_timeout, connect=5.0),
|
||||
"follow_redirects": not litellm.user_url_validation,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await _execute_http_request(client, method, url, headers, body, timeout)
|
||||
response: Final = await _execute_http_request(client, normalized_method, url, headers, body, effective_timeout)
|
||||
return _http_success_response(response)
|
||||
|
||||
except SSRFError as e:
|
||||
verbose_proxy_logger.warning("Custom code http_request blocked: %s", e)
|
||||
return _http_error_response(f"Blocked URL: {e}")
|
||||
except httpx.TimeoutException as e:
|
||||
verbose_proxy_logger.warning("Custom code http_request timeout: %s", e)
|
||||
return _http_error_response(f"Request timeout after {timeout}s")
|
||||
return _http_error_response(f"Request timeout after {effective_timeout}s")
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Return the response even for non-2xx status codes
|
||||
return _http_success_response(e.response)
|
||||
|
|
@ -542,21 +552,47 @@ async def _execute_http_request(
|
|||
) -> httpx.Response:
|
||||
"""Execute the HTTP request using the appropriate client method."""
|
||||
json_body, data_body = _prepare_http_body(body)
|
||||
outbound_headers: Final = _caller_headers(headers)
|
||||
|
||||
if method == "GET":
|
||||
return await client.get(url=url, headers=headers)
|
||||
elif method == "POST":
|
||||
return await client.post(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
|
||||
return await async_safe_get(client, url, headers=outbound_headers)
|
||||
|
||||
destination_url, destination_headers = await _validated_destination(url, outbound_headers)
|
||||
if method == "POST":
|
||||
return await client.post(
|
||||
url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout
|
||||
)
|
||||
elif method == "PUT":
|
||||
return await client.put(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
|
||||
return await client.put(
|
||||
url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout
|
||||
)
|
||||
elif method == "DELETE":
|
||||
return await client.delete(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
|
||||
return await client.delete(
|
||||
url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout
|
||||
)
|
||||
elif method == "PATCH":
|
||||
return await client.patch(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
|
||||
return await client.patch(
|
||||
url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported HTTP method: {method}")
|
||||
|
||||
|
||||
def _caller_headers(headers: dict[str, str] | None) -> dict[str, str]:
|
||||
if headers is None:
|
||||
return {}
|
||||
if not litellm.user_url_validation:
|
||||
return headers
|
||||
return {name: value for name, value in headers.items() if name.lower() != "host"}
|
||||
|
||||
|
||||
async def _validated_destination(url: str, headers: dict[str, str]) -> tuple[str, dict[str, str]]:
|
||||
if not litellm.user_url_validation:
|
||||
return url, headers
|
||||
destination_url, host_header = await asyncio.to_thread(validate_url, url)
|
||||
return destination_url, {**headers, "Host": host_header}
|
||||
|
||||
|
||||
async def http_get(
|
||||
url: str,
|
||||
headers: dict[str, str] | None = None,
|
||||
|
|
|
|||
|
|
@ -27,13 +27,15 @@ from RestrictedPython import (
|
|||
safe_builtins,
|
||||
utility_builtins,
|
||||
)
|
||||
from RestrictedPython.Eval import default_guarded_getitem, default_guarded_getiter
|
||||
from RestrictedPython.Eval import default_guarded_getitem
|
||||
from RestrictedPython.Guards import (
|
||||
full_write_guard,
|
||||
guarded_iter_unpack_sequence,
|
||||
safer_getattr,
|
||||
)
|
||||
from RestrictedPython.transformer import copy_locations
|
||||
|
||||
from .bounded_execution import budget_ok, budgeted_iter
|
||||
from .primitives import get_custom_code_primitives
|
||||
|
||||
|
||||
|
|
@ -46,11 +48,31 @@ class AsyncAwareTransformer(RestrictingNodeTransformer):
|
|||
check, print-scope wrapping, and any future additions to that method are
|
||||
inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to
|
||||
``node_contents_visit`` so their children still get transformed.
|
||||
|
||||
``visit_While`` rewrites ``while test:`` to ``while _budget_ok_() and test:``
|
||||
so a loop that never yields is still stopped at the execution deadline;
|
||||
``for`` loops and comprehensions get the same check through ``_getiter_``.
|
||||
"""
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST:
|
||||
return self.visit_FunctionDef(node)
|
||||
|
||||
def visit_While(self, node: ast.While) -> ast.AST:
|
||||
visited: Final = self.node_contents_visit(node)
|
||||
budget_check: Final = ast.Call(
|
||||
func=ast.Name(id="_budget_ok_", ctx=ast.Load()),
|
||||
args=[], # mutable-ok: ast accepts list fields only
|
||||
keywords=[], # mutable-ok: ast accepts list fields only
|
||||
)
|
||||
test: Final = ast.BoolOp(
|
||||
op=ast.And(),
|
||||
values=[budget_check, visited.test], # mutable-ok: ast accepts list fields only
|
||||
)
|
||||
copy_locations(test, visited.test)
|
||||
bounded: Final = ast.While(test=test, body=visited.body, orelse=visited.orelse)
|
||||
copy_locations(bounded, visited)
|
||||
return bounded
|
||||
|
||||
def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST:
|
||||
return self.node_contents_visit(node)
|
||||
|
||||
|
|
@ -113,10 +135,11 @@ def build_sandbox_globals() -> dict[str, object]:
|
|||
"__builtins__": _build_sandbox_builtins(),
|
||||
"_getattr_": safer_getattr,
|
||||
"_getitem_": default_guarded_getitem,
|
||||
"_getiter_": default_guarded_getiter,
|
||||
"_getiter_": budgeted_iter,
|
||||
"_iter_unpack_sequence_": guarded_iter_unpack_sequence,
|
||||
"_write_": full_write_guard,
|
||||
"_inplacevar_": _inplacevar_,
|
||||
"_budget_ok_": budget_ok,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,174 @@
|
|||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import (
|
||||
ExecutionTimeoutError,
|
||||
SandboxExit,
|
||||
await_with_timeout,
|
||||
call_off_loop_with_timeout,
|
||||
call_with_timeout,
|
||||
)
|
||||
|
||||
|
||||
def _worker_threads() -> list[str]:
|
||||
return [t.name for t in threading.enumerate() if t.name.startswith("guardrail-code:")]
|
||||
|
||||
|
||||
def _spin_forever() -> None:
|
||||
n = 0
|
||||
while True:
|
||||
n += 1
|
||||
|
||||
|
||||
def _spin_swallowing_exceptions() -> None:
|
||||
while True:
|
||||
try:
|
||||
_spin_forever()
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
|
||||
async def _swallow_cancellations_for(seconds: float) -> str:
|
||||
deadline = time.monotonic() + seconds
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
await asyncio.sleep(deadline - time.monotonic())
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
return "survived"
|
||||
|
||||
|
||||
def _exit_now() -> None:
|
||||
raise SystemExit("bye")
|
||||
|
||||
|
||||
def test_call_with_timeout_returns_the_result_and_reraises_failures():
|
||||
assert call_with_timeout(lambda: 42, 1.0, label="ok") == 42
|
||||
with pytest.raises(ZeroDivisionError):
|
||||
call_with_timeout(lambda: 1 // 0, 1.0, label="boom")
|
||||
|
||||
|
||||
def test_call_with_timeout_delivers_a_system_exit_at_once():
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(SandboxExit, match="SystemExit: bye"):
|
||||
call_with_timeout(_exit_now, 5.0, label="exit")
|
||||
|
||||
assert time.monotonic() - started < 1.0
|
||||
|
||||
|
||||
async def _exit_later() -> None:
|
||||
await asyncio.sleep(0)
|
||||
raise SystemExit("bye")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_with_timeout_contains_a_system_exit_instead_of_stopping_the_loop():
|
||||
with pytest.raises(SandboxExit, match="SystemExit: bye"):
|
||||
await await_with_timeout(_exit_later(), 1.0, label="exit")
|
||||
|
||||
assert await asyncio.sleep(0, result="loop still running") == "loop still running"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fn", [_spin_forever, _spin_swallowing_exceptions])
|
||||
def test_call_with_timeout_interrupts_a_busy_loop_and_reclaims_the_thread(fn):
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(ExecutionTimeoutError, match=r"exceeded the 0\.2s execution timeout") as exc_info:
|
||||
call_with_timeout(fn, 0.2, label="spin")
|
||||
|
||||
assert exc_info.value.timeout == 0.2
|
||||
assert time.monotonic() - started < 1.5
|
||||
time.sleep(0.2)
|
||||
assert _worker_threads() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_off_loop_with_timeout_keeps_the_loop_running_and_stops_the_worker():
|
||||
ticks = 0
|
||||
|
||||
async def tick_forever() -> None:
|
||||
nonlocal ticks
|
||||
while True:
|
||||
await asyncio.sleep(0.02)
|
||||
ticks += 1
|
||||
|
||||
ticker = asyncio.create_task(tick_forever())
|
||||
try:
|
||||
assert await call_off_loop_with_timeout(lambda: "done", 1.0, label="ok") == "done"
|
||||
with pytest.raises(ExecutionTimeoutError):
|
||||
await call_off_loop_with_timeout(_spin_forever, 0.3, label="spin")
|
||||
finally:
|
||||
ticker.cancel()
|
||||
|
||||
assert ticks >= 5
|
||||
await asyncio.sleep(0.2)
|
||||
assert _worker_threads() == []
|
||||
|
||||
|
||||
def _stragglers() -> list[asyncio.Task[object]]:
|
||||
return [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_with_timeout_keeps_cancelling_a_coroutine_that_swallows_cancellation():
|
||||
assert await await_with_timeout(_swallow_cancellations_for(0.0), 1.0, label="ok") == "survived"
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(ExecutionTimeoutError, match=r"exceeded the 0\.1s execution timeout"):
|
||||
await await_with_timeout(_swallow_cancellations_for(0.4), 0.1, label="stubborn")
|
||||
|
||||
assert time.monotonic() - started < 1.0
|
||||
assert _stragglers() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_with_timeout_abandons_a_coroutine_that_never_stops_swallowing_cancellation():
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(ExecutionTimeoutError):
|
||||
await await_with_timeout(_swallow_cancellations_for(2.0), 0.1, label="stubborn")
|
||||
|
||||
elapsed = time.monotonic() - started
|
||||
assert 1.0 <= elapsed < 1.8
|
||||
stragglers = _stragglers()
|
||||
assert len(stragglers) == 1
|
||||
with pytest.raises(ExecutionTimeoutError):
|
||||
await asyncio.gather(*stragglers)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_off_loop_with_timeout_stops_the_worker_when_the_caller_is_cancelled():
|
||||
waiting = asyncio.create_task(call_off_loop_with_timeout(_spin_forever, 30.0, label="spin"))
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
waiting.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await waiting
|
||||
|
||||
await asyncio.sleep(0.3)
|
||||
assert _worker_threads() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_with_timeout_cancels_the_code_when_the_caller_is_cancelled():
|
||||
interrupted = asyncio.Event()
|
||||
|
||||
async def sleep_until_cancelled() -> None:
|
||||
try:
|
||||
await asyncio.sleep(30)
|
||||
except asyncio.CancelledError:
|
||||
interrupted.set()
|
||||
raise
|
||||
|
||||
waiting = asyncio.create_task(await_with_timeout(sleep_until_cancelled(), 30.0, label="sleep"))
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
waiting.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await waiting
|
||||
|
||||
await asyncio.wait_for(interrupted.wait(), timeout=1.0)
|
||||
|
|
@ -1,11 +1,22 @@
|
|||
import asyncio
|
||||
import http.server
|
||||
import threading
|
||||
import time
|
||||
from http.server import ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
|
||||
DEFAULT_EXECUTION_TIMEOUT_SECONDS,
|
||||
CustomCodeCompilationError,
|
||||
CustomCodeExecutionError,
|
||||
CustomCodeGuardrail,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
# str.mro() + generator gi_code + code.replace(co_names=...) + __setattr__
|
||||
# to swap a function's bytecode and read http_get's real builtins dict.
|
||||
|
|
@ -77,18 +88,14 @@ def test_nfkc_homoglyph_rejected_at_compile():
|
|||
[
|
||||
# Literal dunder attribute access.
|
||||
"def apply_guardrail(i, r, t):\n return str.__class__\n",
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
" return ().__class__.__bases__[0].__subclasses__()\n",
|
||||
"def apply_guardrail(i, r, t):\n return ().__class__.__bases__[0].__subclasses__()\n",
|
||||
# gi_code — on the transformer's restricted-names list.
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
" def g():\n yield 1\n"
|
||||
" return g().gi_code\n",
|
||||
"def apply_guardrail(i, r, t):\n def g():\n yield 1\n return g().gi_code\n",
|
||||
# Import forms.
|
||||
"import os\ndef apply_guardrail(i, r, t):\n return allow()\n",
|
||||
"from subprocess import call\n"
|
||||
"def apply_guardrail(i, r, t):\n return allow()\n",
|
||||
"from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()\n",
|
||||
# __import__ is rejected as an underscore-prefixed name.
|
||||
"def apply_guardrail(i, r, t):\n" ' return __import__("os")\n',
|
||||
'def apply_guardrail(i, r, t):\n return __import__("os")\n',
|
||||
],
|
||||
)
|
||||
def test_compile_time_rejections(snippet: str):
|
||||
|
|
@ -100,8 +107,7 @@ def test_compile_time_rejections(snippet: str):
|
|||
"snippet",
|
||||
[
|
||||
# getattr is not in the sandbox builtins — NameError at call time.
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
' return getattr(str, "_"+"_class_"+"_")\n',
|
||||
'def apply_guardrail(i, r, t):\n return getattr(str, "_"+"_class_"+"_")\n',
|
||||
# setattr is guarded_setattr + full_write_guard — setting any attribute
|
||||
# on a user-defined object raises TypeError, whether the name is a
|
||||
# dunder or not.
|
||||
|
|
@ -139,10 +145,7 @@ def test_documented_ssn_example_compiles_and_runs():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_guardrail_compiles_and_runs():
|
||||
code = (
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" return allow()\n"
|
||||
)
|
||||
code = "async def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
|
||||
guardrail = _compile(code)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
|
@ -156,10 +159,7 @@ async def test_async_guardrail_compiles_and_runs():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_code_pre_call_block_uses_passthrough():
|
||||
code = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' return block("blocked by test")\n'
|
||||
)
|
||||
code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n'
|
||||
guardrail = _compile(code)
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
|
|
@ -176,10 +176,7 @@ async def test_custom_code_pre_call_block_uses_passthrough():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_code_post_call_block_raises_http_400():
|
||||
code = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' return block("blocked by test")\n'
|
||||
)
|
||||
code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n'
|
||||
guardrail = _compile(code)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -333,10 +330,7 @@ async def test_custom_code_allow_still_records_success_not_flagged():
|
|||
|
||||
|
||||
def test_typical_sync_guardrail_still_works():
|
||||
code = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" return allow()\n"
|
||||
)
|
||||
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
|
||||
guardrail = _compile(code)
|
||||
assert guardrail._compiled_function is not None
|
||||
|
||||
|
|
@ -363,3 +357,492 @@ def test_augmented_assignment_works():
|
|||
def test_missing_apply_guardrail_raises():
|
||||
with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"):
|
||||
_compile("x = 1\n")
|
||||
|
||||
|
||||
class _QuietServer(ThreadingHTTPServer):
|
||||
def handle_error(self, request: object, client_address: object) -> None:
|
||||
return
|
||||
|
||||
|
||||
def _guardrail_worker_threads() -> list[str]:
|
||||
return [t.name for t in threading.enumerate() if t.name.startswith("guardrail-code:")]
|
||||
|
||||
|
||||
class _LocalServer:
|
||||
"""Loopback HTTP server that records every request it receives."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.hits: list[tuple[str, str]] = []
|
||||
self.received_headers: list[list[tuple[str, str]]] = []
|
||||
server = self
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
server.hits.append(("GET", self.path))
|
||||
server.received_headers.append(list(self.headers.items()))
|
||||
if self.path.startswith("/redirect-to/"):
|
||||
self._redirect()
|
||||
return
|
||||
if self.path == "/slow":
|
||||
time.sleep(2)
|
||||
self._reply(b"marker")
|
||||
|
||||
def do_POST(self) -> None:
|
||||
server.hits.append(("POST", self.path))
|
||||
server.received_headers.append(list(self.headers.items()))
|
||||
if self.path.startswith("/redirect-to/"):
|
||||
self._redirect()
|
||||
return
|
||||
self._reply(b"posted")
|
||||
|
||||
def do_PUT(self) -> None:
|
||||
self._record_and_reply(b"put")
|
||||
|
||||
def do_DELETE(self) -> None:
|
||||
self._record_and_reply(b"deleted")
|
||||
|
||||
def do_PATCH(self) -> None:
|
||||
self._record_and_reply(b"patched")
|
||||
|
||||
def _record_and_reply(self, body: bytes) -> None:
|
||||
server.hits.append((self.command, self.path))
|
||||
server.received_headers.append(list(self.headers.items()))
|
||||
self._reply(body)
|
||||
|
||||
def _redirect(self) -> None:
|
||||
target_port = self.path.rsplit("/", 1)[1]
|
||||
self.send_response(302)
|
||||
self.send_header("Location", f"http://127.0.0.1:{target_port}/marker")
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
|
||||
def _reply(self, body: bytes) -> None:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args: object) -> None:
|
||||
return
|
||||
|
||||
self.httpd = _QuietServer(("127.0.0.1", 0), Handler)
|
||||
self.port = self.httpd.server_address[1]
|
||||
threading.Thread(target=self.httpd.serve_forever, daemon=True).start()
|
||||
|
||||
def close(self) -> None:
|
||||
self.httpd.shutdown()
|
||||
self.httpd.server_close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_server():
|
||||
server = _LocalServer()
|
||||
yield server
|
||||
server.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def second_server():
|
||||
server = _LocalServer()
|
||||
yield server
|
||||
server.close()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_http_client_and_url_policy(monkeypatch):
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
monkeypatch.setattr(litellm, "user_url_validation", True)
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [])
|
||||
|
||||
|
||||
def _reporting_guardrail(call: str) -> CustomCodeGuardrail:
|
||||
code = (
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
f" r = await {call}\n"
|
||||
' return block("status=" + str(r["status_code"]) + " body=" + str(r["body"])'
|
||||
' + " error=" + str(r["error"]))\n'
|
||||
)
|
||||
return _compile(code)
|
||||
|
||||
|
||||
async def _block_reason(guardrail: CustomCodeGuardrail) -> str:
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request")
|
||||
return exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_refuses_loopback_by_default(local_server):
|
||||
guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")')
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=0" in reason
|
||||
assert "error=Blocked URL" in reason
|
||||
assert "user_url_allowed_hosts" in reason
|
||||
assert local_server.hits == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_post_refuses_loopback_by_default(local_server):
|
||||
guardrail = _reporting_guardrail(f'http_post("http://127.0.0.1:{local_server.port}/hook", body={{"a": 1}})')
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "error=Blocked URL" in reason
|
||||
assert local_server.hits == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_reaches_an_allowlisted_host(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")')
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=200 body=marker error=None" in reason
|
||||
assert local_server.hits == [("GET", "/marker")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_refuses_a_redirect_into_a_blocked_host(local_server, second_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(
|
||||
f'http_get("http://127.0.0.1:{local_server.port}/redirect-to/{second_server.port}")'
|
||||
)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "error=Blocked URL" in reason
|
||||
assert local_server.hits == [("GET", f"/redirect-to/{second_server.port}")]
|
||||
assert second_server.hits == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_post_does_not_follow_redirects(local_server, second_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(
|
||||
f'http_post("http://127.0.0.1:{local_server.port}/redirect-to/{second_server.port}", body={{"a": 1}})'
|
||||
)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=302" in reason
|
||||
assert second_server.hits == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call", ["http_post", "http_get"])
|
||||
async def test_caller_host_header_never_reaches_the_validated_destination(local_server, monkeypatch, call):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(
|
||||
f'{call}("http://127.0.0.1:{local_server.port}/marker", headers={{"host": "spoofed", "X-Extra": "kept"}})'
|
||||
)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=200" in reason
|
||||
(received,) = local_server.received_headers
|
||||
assert [value for name, value in received if name.lower() == "host"] == [f"127.0.0.1:{local_server.port}"]
|
||||
assert ("x-extra", "kept") in received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_headers_pass_through_untouched_when_url_validation_is_disabled(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
guardrail = _reporting_guardrail(
|
||||
f'http_post("http://127.0.0.1:{local_server.port}/marker", headers={{"host": "spoofed", "X-Extra": "kept"}})'
|
||||
)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=200" in reason
|
||||
(received,) = local_server.received_headers
|
||||
assert [value for name, value in received if name.lower() == "host"] == ["spoofed"]
|
||||
assert ("x-extra", "kept") in received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("method", "body"), [("PUT", "put"), ("DELETE", "deleted"), ("PATCH", "patched")])
|
||||
async def test_http_request_other_methods_reach_an_allowlisted_host(local_server, monkeypatch, method, body):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(
|
||||
f'http_request("http://127.0.0.1:{local_server.port}/marker", method="{method}")'
|
||||
)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert f"status=200 body={body}" in reason
|
||||
assert local_server.hits == [(method, "/marker")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_request_refuses_a_method_outside_the_allowlist(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(f'http_request("http://127.0.0.1:{local_server.port}/marker", method="TRACE")')
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "error=Invalid HTTP method: TRACE" in reason
|
||||
assert local_server.hits == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_gives_up_at_its_own_timeout(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/slow", timeout=0.5)')
|
||||
|
||||
started = time.monotonic()
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert time.monotonic() - started < 1.5
|
||||
assert "error=Request timeout after 0.5s" in reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_guardrail_returning_a_coroutine_has_it_awaited():
|
||||
code = (
|
||||
"async def decide():\n"
|
||||
' return block("decided late")\n'
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" return decide()\n"
|
||||
)
|
||||
guardrail = _compile(code)
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "decided late" in reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_is_unvalidated_when_url_validation_is_disabled(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")')
|
||||
|
||||
reason = await _block_reason(guardrail)
|
||||
|
||||
assert "status=200 body=marker" in reason
|
||||
assert local_server.hits == [("GET", "/marker")]
|
||||
|
||||
|
||||
BUSY_LOOP_GUARDRAIL = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n n = 0\n while True:\n n += 1\n"
|
||||
)
|
||||
|
||||
SWALLOWING_BUSY_LOOP_GUARDRAIL = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" n = 0\n"
|
||||
" while True:\n"
|
||||
" try:\n"
|
||||
" n += 1\n"
|
||||
" except Exception:\n"
|
||||
" n = 0\n"
|
||||
)
|
||||
|
||||
|
||||
async def _expect_execution_timeout(guardrail: CustomCodeGuardrail) -> float:
|
||||
started = time.monotonic()
|
||||
with pytest.raises(CustomCodeExecutionError, match=r"exceeded its 0\.3s execution timeout"):
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request")
|
||||
return time.monotonic() - started
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("code", [BUSY_LOOP_GUARDRAIL, SWALLOWING_BUSY_LOOP_GUARDRAIL])
|
||||
async def test_sync_busy_loop_is_stopped_at_the_execution_timeout(code):
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3)
|
||||
|
||||
elapsed = await _expect_execution_timeout(guardrail)
|
||||
|
||||
assert elapsed < 2.0
|
||||
await asyncio.sleep(0.2)
|
||||
assert _guardrail_worker_threads() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_busy_loop_does_not_stall_the_event_loop():
|
||||
guardrail = CustomCodeGuardrail(custom_code=BUSY_LOOP_GUARDRAIL, guardrail_name="busy", execution_timeout=0.3)
|
||||
ticks = 0
|
||||
|
||||
async def tick_forever() -> None:
|
||||
nonlocal ticks
|
||||
while True:
|
||||
await asyncio.sleep(0.02)
|
||||
ticks += 1
|
||||
|
||||
ticker = asyncio.create_task(tick_forever())
|
||||
try:
|
||||
await _expect_execution_timeout(guardrail)
|
||||
finally:
|
||||
ticker.cancel()
|
||||
|
||||
assert ticks >= 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_guardrail_is_stopped_at_the_execution_timeout(local_server, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
code = (
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
f' await http_get("http://127.0.0.1:{local_server.port}/slow")\n'
|
||||
" return allow()\n"
|
||||
)
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3)
|
||||
|
||||
elapsed = await _expect_execution_timeout(guardrail)
|
||||
|
||||
assert elapsed < 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_guardrail_that_swallows_cancellation_is_stopped_at_the_execution_timeout(
|
||||
local_server, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"])
|
||||
code = (
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" attempts = 0\n"
|
||||
" while attempts < 3:\n"
|
||||
" try:\n"
|
||||
f' await http_get("http://127.0.0.1:{local_server.port}/slow")\n'
|
||||
" except BaseException:\n"
|
||||
" attempts += 1\n"
|
||||
" return allow()\n"
|
||||
)
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="stubborn", execution_timeout=0.3)
|
||||
|
||||
elapsed = await _expect_execution_timeout(guardrail)
|
||||
|
||||
assert elapsed < 1.5
|
||||
|
||||
|
||||
LOOP_SHAPES_GUARDRAIL = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" n = 0\n"
|
||||
" while n < 3:\n"
|
||||
" n += 1\n"
|
||||
" else:\n"
|
||||
" n += 10\n"
|
||||
" pairs = [(k, v) for k, v in request_data['metadata'].items()]\n"
|
||||
" for k, v in pairs:\n"
|
||||
" n += v\n"
|
||||
" for i, (k, v) in zip(range(len(pairs)), pairs):\n"
|
||||
" n += i\n"
|
||||
" keys = sorted(k for k, v in pairs)\n"
|
||||
" return block(reason=str(n) + ' ' + ' '.join(keys))\n"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_checks_keep_every_loop_shape_working():
|
||||
guardrail = CustomCodeGuardrail(custom_code=LOOP_SHAPES_GUARDRAIL, guardrail_name="loops")
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["x"]}, request_data={"model": "m", "metadata": {"b": 2, "a": 5}}, input_type="request"
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "21 a b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(10)
|
||||
@pytest.mark.parametrize(
|
||||
"code",
|
||||
[
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n while True:\n pass\n",
|
||||
(
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" for a in range(500):\n"
|
||||
" for b in range(500):\n"
|
||||
" for c in range(500):\n"
|
||||
" pass\n"
|
||||
" return allow()\n"
|
||||
),
|
||||
(
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" try:\n"
|
||||
" while True:\n"
|
||||
" pass\n"
|
||||
" except BaseException:\n"
|
||||
" pass\n"
|
||||
" return allow()\n"
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_async_loop_that_never_yields_is_stopped_at_the_execution_timeout(code):
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="spin", execution_timeout=0.3)
|
||||
|
||||
elapsed = await _expect_execution_timeout(guardrail)
|
||||
|
||||
assert elapsed < 1.5
|
||||
assert await asyncio.sleep(0, result="loop still running") == "loop still running"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"code",
|
||||
[
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n",
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n",
|
||||
],
|
||||
)
|
||||
async def test_system_exit_from_guardrail_code_is_an_execution_error_not_a_timeout(code):
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="exit", execution_timeout=5.0)
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(CustomCodeExecutionError, match="execution failed: SystemExit: bye"):
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request")
|
||||
|
||||
assert time.monotonic() - started < 1.0
|
||||
|
||||
|
||||
def test_module_level_busy_loop_fails_compilation_at_the_execution_timeout():
|
||||
code = "n = 0\nwhile True:\n n += 1\n" + BUSY_LOOP_GUARDRAIL
|
||||
started = time.monotonic()
|
||||
|
||||
with pytest.raises(CustomCodeCompilationError, match=r"exceeded the 0\.3s execution timeout"):
|
||||
CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3)
|
||||
|
||||
assert time.monotonic() - started < 2.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("execution_timeout", [0, -1.0])
|
||||
def test_execution_timeout_must_be_positive(execution_timeout):
|
||||
with pytest.raises(ValueError, match="execution_timeout must be positive"):
|
||||
CustomCodeGuardrail(
|
||||
custom_code="def apply_guardrail(i, r, t):\n return allow()\n", execution_timeout=execution_timeout
|
||||
)
|
||||
|
||||
|
||||
def _initialize_from_config(guardrail_name: str, litellm_params: dict[str, object]) -> CustomCodeGuardrail:
|
||||
InMemoryGuardrailHandler().initialize_guardrail(
|
||||
guardrail={
|
||||
"guardrail_name": guardrail_name,
|
||||
"litellm_params": {
|
||||
"guardrail": SupportedGuardrailIntegrations.CUSTOM_CODE.value,
|
||||
"mode": "pre_call",
|
||||
"custom_code": "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n",
|
||||
**litellm_params,
|
||||
},
|
||||
}
|
||||
)
|
||||
initialized = [
|
||||
callback
|
||||
for callback in litellm.callbacks
|
||||
if isinstance(callback, CustomCodeGuardrail) and callback.guardrail_name == guardrail_name
|
||||
]
|
||||
assert initialized, f"{guardrail_name} was not registered as a callback"
|
||||
return initialized[-1]
|
||||
|
||||
|
||||
def test_config_timeout_reaches_the_guardrail():
|
||||
assert _initialize_from_config("custom-code-timeout", {"timeout": 0.2}).execution_timeout == 0.2
|
||||
|
||||
|
||||
def test_config_without_timeout_uses_the_default():
|
||||
assert (
|
||||
_initialize_from_config("custom-code-default-timeout", {}).execution_timeout
|
||||
== DEFAULT_EXECUTION_TIMEOUT_SECONDS
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
|
@ -13,6 +14,7 @@ from litellm.proxy.guardrails.guardrail_endpoints import (
|
|||
CreateGuardrailRequest,
|
||||
PatchGuardrailRequest,
|
||||
RegisterGuardrailRequest,
|
||||
TestCustomCodeGuardrailRequest,
|
||||
UpdateGuardrailRequest,
|
||||
apply_guardrail,
|
||||
approve_guardrail_submission,
|
||||
|
|
@ -28,6 +30,9 @@ from litellm.proxy.guardrails.guardrail_endpoints import (
|
|||
reject_guardrail_submission,
|
||||
update_guardrail,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_endpoints import (
|
||||
test_custom_code_guardrail as run_custom_code_test_endpoint,
|
||||
)
|
||||
|
||||
MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
from litellm.proxy.guardrails.guardrail_registry import (
|
||||
|
|
@ -87,12 +92,8 @@ def mock_prisma_client(mocker):
|
|||
# Create async mocks for the database methods
|
||||
mock_client.db = mocker.Mock()
|
||||
mock_client.db.litellm_guardrailstable = mocker.Mock()
|
||||
mock_client.db.litellm_guardrailstable.find_many = AsyncMock(
|
||||
return_value=[MOCK_DB_GUARDRAIL]
|
||||
)
|
||||
mock_client.db.litellm_guardrailstable.find_unique = AsyncMock(
|
||||
return_value=MOCK_DB_GUARDRAIL
|
||||
)
|
||||
mock_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[MOCK_DB_GUARDRAIL])
|
||||
mock_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
|
||||
return mock_client
|
||||
|
||||
|
||||
|
|
@ -118,17 +119,13 @@ def mock_guardrail_registry(mocker):
|
|||
return_value={**MOCK_DB_GUARDRAIL, "guardrail_id": "new-test-guardrail-id"}
|
||||
)
|
||||
mock_registry.delete_guardrail_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
|
||||
mock_registry.get_guardrail_by_id_from_db = AsyncMock(
|
||||
return_value=MOCK_DB_GUARDRAIL
|
||||
)
|
||||
mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
|
||||
mock_registry.update_guardrail_in_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
|
||||
return mock_registry
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_guardrails_v2_with_db_and_config(
|
||||
mocker, mock_prisma_client, mock_in_memory_handler
|
||||
):
|
||||
async def test_list_guardrails_v2_with_db_and_config(mocker, mock_prisma_client, mock_in_memory_handler):
|
||||
"""Test listing guardrails from both DB and config"""
|
||||
# Mock the prisma client
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -144,17 +141,13 @@ async def test_list_guardrails_v2_with_db_and_config(
|
|||
assert len(response.guardrails) == 2
|
||||
|
||||
# Check DB guardrail
|
||||
db_guardrail = next(
|
||||
g for g in response.guardrails if g.guardrail_id == "test-db-guardrail"
|
||||
)
|
||||
db_guardrail = next(g for g in response.guardrails if g.guardrail_id == "test-db-guardrail")
|
||||
assert db_guardrail.guardrail_name == "Test DB Guardrail"
|
||||
assert db_guardrail.guardrail_definition_location == "db"
|
||||
assert isinstance(db_guardrail.litellm_params, BaseLitellmParams)
|
||||
|
||||
# Check config guardrail
|
||||
config_guardrail = next(
|
||||
g for g in response.guardrails if g.guardrail_id == "test-config-guardrail"
|
||||
)
|
||||
config_guardrail = next(g for g in response.guardrails if g.guardrail_id == "test-config-guardrail")
|
||||
assert config_guardrail.guardrail_name == "Test Config Guardrail"
|
||||
assert config_guardrail.guardrail_definition_location == "config"
|
||||
assert isinstance(config_guardrail.litellm_params, BaseLitellmParams)
|
||||
|
|
@ -196,9 +189,7 @@ async def test_list_guardrails_v2_skips_stale_db_backed_in_memory_entries(mocker
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_404s_stale_db_backed_entry(
|
||||
mocker, mock_prisma_client, mock_in_memory_handler
|
||||
):
|
||||
async def test_get_guardrail_info_404s_stale_db_backed_entry(mocker, mock_prisma_client, mock_in_memory_handler):
|
||||
"""
|
||||
Stale DB-backed entry (in-memory but not in DB) must 404 instead of being
|
||||
returned as if it were a config-loaded guardrail.
|
||||
|
|
@ -208,9 +199,7 @@ async def test_get_guardrail_info_404s_stale_db_backed_entry(
|
|||
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
|
||||
mock_in_memory_handler,
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
|
||||
# In-memory still has it, but it's tagged as 'db' (stale, awaiting reconcile)
|
||||
mock_in_memory_handler.get_source.return_value = "db"
|
||||
|
||||
|
|
@ -241,9 +230,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker):
|
|||
mock_prisma_client = mocker.Mock()
|
||||
mock_prisma_client.db = mocker.Mock()
|
||||
mock_prisma_client.db.litellm_guardrailstable = mocker.Mock()
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
|
||||
return_value=[db_guardrail_with_secrets]
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[db_guardrail_with_secrets])
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
|
||||
|
|
@ -263,11 +250,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker):
|
|||
if isinstance(litellm_params, dict):
|
||||
params = litellm_params
|
||||
else:
|
||||
params = (
|
||||
litellm_params.model_dump()
|
||||
if hasattr(litellm_params, "model_dump")
|
||||
else dict(litellm_params)
|
||||
)
|
||||
params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params)
|
||||
|
||||
# Sensitive keys (containing "key", "secret", "token", etc.) should be masked
|
||||
assert params["api_key"] != "sk-1234567890abcdef"
|
||||
|
|
@ -299,9 +282,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock
|
|||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = [
|
||||
config_guardrail_with_secrets
|
||||
]
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = [config_guardrail_with_secrets]
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -318,11 +299,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock
|
|||
if isinstance(litellm_params, dict):
|
||||
params = litellm_params
|
||||
else:
|
||||
params = (
|
||||
litellm_params.model_dump()
|
||||
if hasattr(litellm_params, "model_dump")
|
||||
else dict(litellm_params)
|
||||
)
|
||||
params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params)
|
||||
|
||||
# Sensitive keys should be masked
|
||||
assert params["api_key"] != "my-secret-bedrock-key"
|
||||
|
|
@ -355,9 +332,7 @@ async def test_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are
|
|||
mock_prisma_client = mocker.Mock()
|
||||
mock_prisma_client.db = mocker.Mock()
|
||||
mock_prisma_client.db.litellm_guardrailstable = mocker.Mock()
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
|
||||
return_value=[other_team_guardrail]
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[other_team_guardrail])
|
||||
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
mock_in_memory_handler.list_in_memory_guardrails.return_value = []
|
||||
|
|
@ -372,9 +347,7 @@ async def test_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are
|
|||
AsyncMock(return_value=[]),
|
||||
)
|
||||
|
||||
viewer_auth = UserAPIKeyAuth(
|
||||
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
)
|
||||
viewer_auth = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
||||
response = await list_guardrails_v2(user_api_key_dict=viewer_auth)
|
||||
|
||||
assert [g.guardrail_id for g in response.guardrails] == ["other-team-guardrail"]
|
||||
|
|
@ -421,16 +394,10 @@ async def test_list_guardrails_v2_masks_sensitive_data_for_admin_viewer(mocker):
|
|||
AsyncMock(return_value=[]),
|
||||
)
|
||||
|
||||
viewer_auth = UserAPIKeyAuth(
|
||||
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
)
|
||||
viewer_auth = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
||||
response = await list_guardrails_v2(user_api_key_dict=viewer_auth)
|
||||
|
||||
guardrail = next(
|
||||
g
|
||||
for g in response.guardrails
|
||||
if g.guardrail_id == "other-team-secret-guardrail"
|
||||
)
|
||||
guardrail = next(g for g in response.guardrails if g.guardrail_id == "other-team-secret-guardrail")
|
||||
params = guardrail.litellm_params.model_dump()
|
||||
assert params["api_key"] != "sk-viewer-must-not-see-this"
|
||||
assert "****" in str(params["api_key"])
|
||||
|
|
@ -451,9 +418,7 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_from_config(
|
||||
mocker, mock_prisma_client, mock_in_memory_handler
|
||||
):
|
||||
async def test_get_guardrail_info_from_config(mocker, mock_prisma_client, mock_in_memory_handler):
|
||||
"""Test getting guardrail info from config when not found in DB"""
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -462,9 +427,7 @@ async def test_get_guardrail_info_from_config(
|
|||
)
|
||||
|
||||
# Mock DB to return None
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
response = await get_guardrail_info("test-config-guardrail")
|
||||
|
||||
|
|
@ -475,9 +438,7 @@ async def test_get_guardrail_info_from_config(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_not_found(
|
||||
mocker, mock_prisma_client, mock_in_memory_handler
|
||||
):
|
||||
async def test_get_guardrail_info_not_found(mocker, mock_prisma_client, mock_in_memory_handler):
|
||||
"""Test getting guardrail info when not found in either DB or config"""
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -486,9 +447,7 @@ async def test_get_guardrail_info_not_found(
|
|||
)
|
||||
|
||||
# Mock both DB and in-memory handler to return None
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None)
|
||||
mock_in_memory_handler.get_guardrail_by_id.return_value = None
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
@ -499,9 +458,7 @@ async def test_get_guardrail_info_not_found(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_guardrails_v2_without_prisma_returns_config_guardrails(
|
||||
mocker, mock_in_memory_handler
|
||||
):
|
||||
async def test_list_guardrails_v2_without_prisma_returns_config_guardrails(mocker, mock_in_memory_handler):
|
||||
"""
|
||||
A proxy without a DB must still list config-defined guardrails instead of
|
||||
raising 500 'Prisma client not initialized'.
|
||||
|
|
@ -535,18 +492,14 @@ async def test_list_guardrails_v2_without_prisma_non_admin_sees_unrestricted_con
|
|||
mock_in_memory_handler,
|
||||
)
|
||||
|
||||
non_admin_auth = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal-user-1"
|
||||
)
|
||||
non_admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal-user-1")
|
||||
response = await list_guardrails_v2(user_api_key_dict=non_admin_auth)
|
||||
|
||||
assert [g.guardrail_id for g in response.guardrails] == ["test-config-guardrail"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_without_prisma_returns_config_guardrail(
|
||||
mocker, mock_in_memory_handler
|
||||
):
|
||||
async def test_get_guardrail_info_without_prisma_returns_config_guardrail(mocker, mock_in_memory_handler):
|
||||
"""
|
||||
The info endpoint must serve config-defined guardrails from the in-memory
|
||||
registry when no DB is attached instead of raising 500.
|
||||
|
|
@ -565,9 +518,7 @@ async def test_get_guardrail_info_without_prisma_returns_config_guardrail(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_guardrail_info_without_prisma_404s_unknown_id(
|
||||
mocker, mock_in_memory_handler
|
||||
):
|
||||
async def test_get_guardrail_info_without_prisma_404s_unknown_id(mocker, mock_in_memory_handler):
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", None)
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
|
||||
|
|
@ -630,10 +581,7 @@ def test_get_provider_specific_params():
|
|||
assert "optional_params" in fields
|
||||
|
||||
# Check the structure of a simple field
|
||||
assert (
|
||||
fields["api_key"]["description"]
|
||||
== "API key for the Azure Content Safety Prompt Shield guardrail"
|
||||
)
|
||||
assert fields["api_key"]["description"] == "API key for the Azure Content Safety Prompt Shield guardrail"
|
||||
assert fields["api_key"]["required"] == False
|
||||
assert fields["api_key"]["type"] == "string" # Should be string, not None
|
||||
|
||||
|
|
@ -657,17 +605,13 @@ def test_get_provider_specific_params():
|
|||
== "Severity threshold for the Azure Content Safety Text Moderation guardrail across all categories"
|
||||
)
|
||||
assert nested_fields["severity_threshold"]["required"] == False
|
||||
assert (
|
||||
nested_fields["severity_threshold"]["type"] == "number"
|
||||
) # Should be number, not None
|
||||
assert nested_fields["severity_threshold"]["type"] == "number" # Should be number, not None
|
||||
|
||||
# Check other field types
|
||||
assert nested_fields["categories"]["type"] == "multiselect"
|
||||
assert nested_fields["blocklistNames"]["type"] == "array"
|
||||
assert nested_fields["haltOnBlocklistHit"]["type"] == "boolean"
|
||||
assert (
|
||||
nested_fields["outputType"]["type"] == "select"
|
||||
) # Literal type should be select
|
||||
assert nested_fields["outputType"]["type"] == "select" # Literal type should be select
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -769,17 +713,11 @@ def test_optional_params_returned_when_properly_overridden():
|
|||
|
||||
# Create specific optional params model
|
||||
class SpecificOptionalParams(BaseModel):
|
||||
threshold: Optional[float] = Field(
|
||||
default=0.5, description="Detection threshold"
|
||||
)
|
||||
categories: Optional[List[str]] = Field(
|
||||
default=None, description="Categories to check"
|
||||
)
|
||||
threshold: Optional[float] = Field(default=0.5, description="Detection threshold")
|
||||
categories: Optional[List[str]] = Field(default=None, description="Categories to check")
|
||||
|
||||
# Create a config model that DOES override optional_params with a specific type
|
||||
class TestGuardrailConfigWithOptionalParams(
|
||||
GuardrailConfigModel[SpecificOptionalParams]
|
||||
):
|
||||
class TestGuardrailConfigWithOptionalParams(GuardrailConfigModel[SpecificOptionalParams]):
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Test API key",
|
||||
|
|
@ -806,9 +744,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key():
|
|||
)
|
||||
|
||||
# Setup guardrail hook
|
||||
guardrail_hook = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
|
||||
)
|
||||
guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1")
|
||||
mock_credentials = Mock()
|
||||
test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
|
||||
|
||||
|
|
@ -839,9 +775,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch):
|
|||
)
|
||||
|
||||
# Setup guardrail hook
|
||||
guardrail_hook = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
|
||||
)
|
||||
guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1")
|
||||
|
||||
# Mock credentials
|
||||
mock_credentials = Mock()
|
||||
|
|
@ -854,7 +788,6 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch):
|
|||
patch("botocore.auth.SigV4Auth") as mock_sigv4_auth,
|
||||
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
|
||||
):
|
||||
|
||||
# Mock SigV4Auth
|
||||
mock_sigv4_instance = Mock()
|
||||
mock_sigv4_auth.return_value = mock_sigv4_instance
|
||||
|
|
@ -873,9 +806,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch):
|
|||
)
|
||||
|
||||
# Verify SigV4 auth was used
|
||||
mock_sigv4_auth.assert_called_once_with(
|
||||
mock_credentials, "bedrock", "us-east-1"
|
||||
)
|
||||
mock_sigv4_auth.assert_called_once_with(mock_credentials, "bedrock", "us-east-1")
|
||||
mock_sigv4_instance.add_auth.assert_called_once()
|
||||
|
||||
|
||||
|
|
@ -889,9 +820,7 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypat
|
|||
)
|
||||
|
||||
# Setup guardrail hook
|
||||
guardrail_hook = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
|
||||
)
|
||||
guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1")
|
||||
|
||||
# Mock credentials
|
||||
mock_credentials = Mock()
|
||||
|
|
@ -928,9 +857,7 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key():
|
|||
BedrockGuardrail,
|
||||
)
|
||||
|
||||
guardrail_hook = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail-id", guardrailVersion="1"
|
||||
)
|
||||
guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1")
|
||||
|
||||
guardrail_hook.async_handler = Mock()
|
||||
mock_response = Mock()
|
||||
|
|
@ -940,20 +867,13 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key():
|
|||
test_request_data = {"api_key": "test-api-key-789"}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)
|
||||
),
|
||||
patch.object(guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)),
|
||||
patch.object(guardrail_hook, "_load_credentials") as mock_load_creds,
|
||||
patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert,
|
||||
patch.object(
|
||||
guardrail_hook, "get_guardrail_dynamic_request_body_params"
|
||||
) as mock_get_params,
|
||||
patch.object(
|
||||
guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"
|
||||
),
|
||||
patch.object(guardrail_hook, "get_guardrail_dynamic_request_body_params") as mock_get_params,
|
||||
patch.object(guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"),
|
||||
patch("botocore.awsrequest.AWSRequest") as mock_aws_request,
|
||||
):
|
||||
|
||||
mock_load_creds.return_value = (Mock(), "us-east-1")
|
||||
mock_convert.return_value = {"source": "INPUT", "content": [{"text": {"text": "test"}}]}
|
||||
mock_get_params.return_value = {}
|
||||
|
|
@ -965,9 +885,7 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key():
|
|||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer test-api-key-789",
|
||||
}
|
||||
mock_request_instance.prepare.return_value = Mock(
|
||||
headers=mock_request_instance.headers
|
||||
)
|
||||
mock_request_instance.prepare.return_value = Mock(headers=mock_request_instance.headers)
|
||||
mock_aws_request.return_value = mock_request_instance
|
||||
|
||||
await guardrail_hook.make_bedrock_api_request(
|
||||
|
|
@ -1025,12 +943,8 @@ async def test_create_guardrail_endpoint(
|
|||
|
||||
elif scenario == "success_sync_fails":
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_in_memory_handler.initialize_guardrail.side_effect = Exception(
|
||||
"Sync failed"
|
||||
)
|
||||
mock_logger = mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
|
||||
)
|
||||
mock_in_memory_handler.initialize_guardrail.side_effect = Exception("Sync failed")
|
||||
mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1044,9 +958,7 @@ async def test_create_guardrail_endpoint(
|
|||
|
||||
elif scenario == "database_failure":
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception(
|
||||
"Database error"
|
||||
)
|
||||
mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception("Database error")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1060,9 +972,7 @@ async def test_create_guardrail_endpoint(
|
|||
# Run the test
|
||||
if expected_exception:
|
||||
with pytest.raises(expected_exception) as exc_info:
|
||||
await create_guardrail(
|
||||
MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
if scenario == "database_failure":
|
||||
assert "Database error" in str(exc_info.value.detail)
|
||||
|
|
@ -1070,9 +980,7 @@ async def test_create_guardrail_endpoint(
|
|||
assert "Prisma client not initialized" in str(exc_info.value.detail)
|
||||
|
||||
else:
|
||||
result = await create_guardrail(
|
||||
MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
result = await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert result["guardrail_id"] == expected_result
|
||||
assert result["guardrail_name"] == "Test DB Guardrail"
|
||||
|
|
@ -1086,9 +994,7 @@ async def test_create_guardrail_endpoint(
|
|||
if scenario == "success_sync_fails":
|
||||
assert mock_logger is not None
|
||||
mock_logger.warning.assert_called_once()
|
||||
assert "Failed to initialize guardrail" in str(
|
||||
mock_logger.warning.call_args
|
||||
)
|
||||
assert "Failed to initialize guardrail" in str(mock_logger.warning.call_args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -1139,12 +1045,8 @@ async def test_update_guardrail_endpoint(
|
|||
# so it keeps the pre-existing swallow-and-warn behavior rather than
|
||||
# rolling back the DB write.
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
|
||||
side_effect=Exception("Sync failed")
|
||||
)
|
||||
mock_logger = mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
|
||||
)
|
||||
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed"))
|
||||
mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1177,9 +1079,7 @@ async def test_update_guardrail_endpoint(
|
|||
|
||||
elif scenario == "database_failure":
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
|
||||
"Database error"
|
||||
)
|
||||
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1209,15 +1109,10 @@ async def test_update_guardrail_endpoint(
|
|||
# Rolled back: update_guardrail_in_db is called once for the
|
||||
# rejected write and once more to restore the previous config.
|
||||
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
|
||||
assert (
|
||||
mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
|
||||
== MOCK_DB_GUARDRAIL
|
||||
)
|
||||
assert mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] == MOCK_DB_GUARDRAIL
|
||||
|
||||
else:
|
||||
result = await update_guardrail(
|
||||
"test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
result = await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert result["guardrail_id"] == expected_result
|
||||
assert result["guardrail_name"] == "Test DB Guardrail"
|
||||
|
|
@ -1228,9 +1123,7 @@ async def test_update_guardrail_endpoint(
|
|||
prisma_client=mocker.ANY,
|
||||
)
|
||||
|
||||
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
|
||||
guardrail=mocker.ANY
|
||||
)
|
||||
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(guardrail=mocker.ANY)
|
||||
|
||||
if scenario == "success_sync_fails_unexpected_error":
|
||||
assert mock_logger is not None
|
||||
|
|
@ -1286,12 +1179,8 @@ async def test_patch_guardrail_endpoint(
|
|||
# config-rejection signal, so it keeps the pre-existing swallow-and-warn
|
||||
# behavior rather than rolling back the DB write.
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
|
||||
side_effect=Exception("Sync failed")
|
||||
)
|
||||
mock_logger = mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
|
||||
)
|
||||
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed"))
|
||||
mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1324,9 +1213,7 @@ async def test_patch_guardrail_endpoint(
|
|||
|
||||
elif scenario == "database_failure":
|
||||
mock_prisma_client = mocker.Mock()
|
||||
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
|
||||
"Database error"
|
||||
)
|
||||
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error")
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
|
|
@ -1358,18 +1245,14 @@ async def test_patch_guardrail_endpoint(
|
|||
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
|
||||
|
||||
else:
|
||||
result = await patch_guardrail(
|
||||
"test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
result = await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert result["guardrail_id"] == expected_result
|
||||
assert result["guardrail_name"] == "Test DB Guardrail"
|
||||
|
||||
mock_guardrail_registry.update_guardrail_in_db.assert_called_once()
|
||||
|
||||
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
|
||||
guardrail=mocker.ANY
|
||||
)
|
||||
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(guardrail=mocker.ANY)
|
||||
|
||||
if scenario == "success_sync_fails_unexpected_error":
|
||||
assert mock_logger is not None
|
||||
|
|
@ -1428,12 +1311,8 @@ async def test_delete_guardrail_endpoint(
|
|||
)
|
||||
|
||||
elif scenario == "success_sync_fails":
|
||||
mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception(
|
||||
"Sync failed"
|
||||
)
|
||||
mock_logger = mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
|
||||
)
|
||||
mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception("Sync failed")
|
||||
mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger")
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
|
||||
|
|
@ -1446,13 +1325,9 @@ async def test_delete_guardrail_endpoint(
|
|||
|
||||
if expected_exception:
|
||||
with pytest.raises(expected_exception):
|
||||
await delete_guardrail(
|
||||
guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
else:
|
||||
result = await delete_guardrail(
|
||||
guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER
|
||||
)
|
||||
result = await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert result == MOCK_DB_GUARDRAIL
|
||||
|
||||
|
|
@ -1463,9 +1338,7 @@ async def test_delete_guardrail_endpoint(
|
|||
guardrail_id=expected_result, prisma_client=mock_prisma_client
|
||||
)
|
||||
|
||||
mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with(
|
||||
guardrail_id=expected_result
|
||||
)
|
||||
mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with(guardrail_id=expected_result)
|
||||
|
||||
if scenario == "success_sync_fails":
|
||||
assert mock_logger is not None
|
||||
|
|
@ -1483,9 +1356,7 @@ async def test_apply_guardrail_not_found(mocker):
|
|||
# Mock the GUARDRAIL_REGISTRY to return None (guardrail not found)
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = None
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
mock_proxy_logging = mocker.Mock()
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
|
|
@ -1495,9 +1366,7 @@ async def test_apply_guardrail_not_found(mocker):
|
|||
mocker.patch("litellm.proxy.proxy_server.version", "test")
|
||||
|
||||
# Create request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="non-existent-guardrail", text="Test input text"
|
||||
)
|
||||
request = ApplyGuardrailRequest(guardrail_name="non-existent-guardrail", text="Test input text")
|
||||
|
||||
# Mock user auth
|
||||
mock_user_auth = UserAPIKeyAuth()
|
||||
|
|
@ -1531,9 +1400,7 @@ async def test_apply_guardrail_execution_error(mocker):
|
|||
# Mock the GUARDRAIL_REGISTRY
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
mock_logging_obj = mocker.Mock()
|
||||
mock_logging_obj.async_failure_handler = AsyncMock()
|
||||
|
|
@ -1555,9 +1422,7 @@ async def test_apply_guardrail_execution_error(mocker):
|
|||
mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor")
|
||||
|
||||
# Create request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail", text="Test input text with forbidden content"
|
||||
)
|
||||
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="Test input text with forbidden content")
|
||||
|
||||
# Mock user auth
|
||||
mock_user_auth = UserAPIKeyAuth()
|
||||
|
|
@ -1581,9 +1446,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
|
|||
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
mock_logging_obj = mocker.Mock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
|
|
@ -1604,13 +1467,9 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
|
|||
mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock())
|
||||
mocker.patch("litellm.proxy.proxy_server.version", "test")
|
||||
mock_executor = mocker.Mock()
|
||||
mocker.patch(
|
||||
"litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor
|
||||
)
|
||||
mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor)
|
||||
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail", text="hello@example.com"
|
||||
)
|
||||
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello@example.com")
|
||||
response = await apply_guardrail(
|
||||
fastapi_request=mocker.Mock(),
|
||||
request=request,
|
||||
|
|
@ -1634,9 +1493,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result):
|
|||
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
mock_logging_obj = mocker.Mock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
|
|
@ -1772,9 +1629,7 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker):
|
|||
# Mock the GUARDRAIL_REGISTRY to return None from DB (so it checks config)
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=None)
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
# Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
|
|
@ -1814,12 +1669,8 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker):
|
|||
|
||||
# Mock the GUARDRAIL_REGISTRY to return a guardrail from DB
|
||||
mock_registry = mocker.Mock()
|
||||
mock_registry.get_guardrail_by_id_from_db = AsyncMock(
|
||||
return_value=MOCK_DB_GUARDRAIL
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry
|
||||
)
|
||||
mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry)
|
||||
|
||||
# Mock IN_MEMORY_GUARDRAIL_HANDLER to return None
|
||||
mock_in_memory_handler = mocker.Mock()
|
||||
|
|
@ -1978,9 +1829,7 @@ async def test_register_guardrail_non_admin_cross_team_allowed(mocker):
|
|||
team_id="team-beta",
|
||||
litellm_params=MOCK_REGISTER_REQUEST.litellm_params,
|
||||
)
|
||||
user = UserAPIKeyAuth(
|
||||
user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha"
|
||||
)
|
||||
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha")
|
||||
|
||||
result = await register_guardrail(req, user)
|
||||
|
||||
|
|
@ -2000,9 +1849,7 @@ async def test_register_guardrail_non_admin_cross_team_forbidden(mocker):
|
|||
team_id="team-other",
|
||||
litellm_params=MOCK_REGISTER_REQUEST.litellm_params,
|
||||
)
|
||||
user = UserAPIKeyAuth(
|
||||
user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha"
|
||||
)
|
||||
user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await register_guardrail(req, user)
|
||||
|
|
@ -2184,9 +2031,7 @@ async def test_list_guardrail_submissions_team_id_filter(mocker):
|
|||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
result = await list_guardrail_submissions(
|
||||
user_api_key_dict=user, team_id="team-abc"
|
||||
)
|
||||
result = await list_guardrail_submissions(user_api_key_dict=user, team_id="team-abc")
|
||||
|
||||
assert len(result.submissions) == 1
|
||||
assert result.submissions[0].guardrail_id == "team-1"
|
||||
|
|
@ -2288,9 +2133,7 @@ async def test_get_guardrail_submission_admin_viewer_other_team_allowed(mocker):
|
|||
"litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids",
|
||||
AsyncMock(return_value=[]),
|
||||
)
|
||||
user = UserAPIKeyAuth(
|
||||
user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
)
|
||||
user = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
|
||||
|
||||
result = await get_guardrail_submission("sub-1", user)
|
||||
|
||||
|
|
@ -2370,9 +2213,7 @@ async def test_reject_guardrail_submission_success(mocker):
|
|||
async def test_reject_guardrail_submission_not_pending(mocker):
|
||||
"""Reject returns 400 when status is not pending_review (e.g. already active)."""
|
||||
mock_prisma = mocker.Mock()
|
||||
row = mocker.Mock(
|
||||
guardrail_id="already-active", guardrail_name="g", status="active"
|
||||
)
|
||||
row = mocker.Mock(guardrail_id="already-active", guardrail_name="g", status="active")
|
||||
mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
|
@ -2404,9 +2245,7 @@ async def test_reject_guardrail_submission_not_pending(mocker):
|
|||
"no_hostname",
|
||||
],
|
||||
)
|
||||
async def test_register_guardrail_rejects_bad_api_base(
|
||||
mocker, api_base, expected_detail
|
||||
):
|
||||
async def test_register_guardrail_rejects_bad_api_base(mocker, api_base, expected_detail):
|
||||
"""Register returns 400 when api_base has invalid scheme or missing hostname."""
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
|
||||
req = RegisterGuardrailRequest(
|
||||
|
|
@ -2474,9 +2313,7 @@ async def test_approve_guardrail_init_failure_returns_warning(mocker):
|
|||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
mock_handler = mocker.Mock()
|
||||
mock_handler.initialize_guardrail = mocker.Mock(
|
||||
side_effect=Exception("missing dependency")
|
||||
)
|
||||
mock_handler.initialize_guardrail = mocker.Mock(side_effect=Exception("missing dependency"))
|
||||
mocker.patch(
|
||||
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
|
||||
mock_handler,
|
||||
|
|
@ -2572,9 +2409,7 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker):
|
|||
user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
# Filter to only pending, but summary should still show both
|
||||
result = await list_guardrail_submissions(
|
||||
status="pending_review", user_api_key_dict=user
|
||||
)
|
||||
result = await list_guardrail_submissions(status="pending_review", user_api_key_dict=user)
|
||||
|
||||
assert len(result.submissions) == 1 # filtered
|
||||
assert result.summary.total == 2 # unfiltered
|
||||
|
|
@ -2630,15 +2465,13 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks():
|
|||
for provider, guardrail_class in guardrail_class_registry.items():
|
||||
declared = guardrail_class.get_supported_event_hooks()
|
||||
if declared is None:
|
||||
assert (
|
||||
provider not in result.supported_modes_by_provider
|
||||
), f"{provider} returned None from classmethod but appears in map"
|
||||
assert provider not in result.supported_modes_by_provider, (
|
||||
f"{provider} returned None from classmethod but appears in map"
|
||||
)
|
||||
continue
|
||||
|
||||
assert provider in result.supported_modes_by_provider, provider
|
||||
assert result.supported_modes_by_provider[provider] == [
|
||||
hook.value for hook in declared
|
||||
], provider
|
||||
assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider
|
||||
|
||||
|
||||
def test_content_filter_runtime_rejects_unsupported_mcp_hook():
|
||||
|
|
@ -2725,3 +2558,115 @@ def test_field_type_inference_handles_pep604_unions():
|
|||
assert _get_field_type_from_annotation(list[str] | None) == "array"
|
||||
assert _get_field_type_from_annotation(bool | None) == "boolean"
|
||||
assert _unwrap_optional_type(str | None) is str
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(20)
|
||||
async def test_test_custom_code_endpoint_returns_a_timeout_for_an_infinite_loop():
|
||||
"""The endpoint used to join the worker thread after its timeout fired, so an infinite
|
||||
loop hung the request forever."""
|
||||
request = TestCustomCodeGuardrailRequest(
|
||||
custom_code="def apply_guardrail(inputs, request_data, input_type):\n n = 0\n while True:\n n += 1\n",
|
||||
test_input={"texts": ["x"]},
|
||||
)
|
||||
started = time.monotonic()
|
||||
|
||||
response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert response.success is False
|
||||
assert response.error_type == "execution"
|
||||
assert response.error is not None
|
||||
assert response.error.startswith("Execution timeout: code took longer than 5 seconds")
|
||||
assert time.monotonic() - started < 8.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(20)
|
||||
async def test_test_custom_code_endpoint_reports_a_module_level_infinite_loop_as_a_timeout():
|
||||
"""Module-level code that outran the load deadline was reported as a compile failure, as if the
|
||||
source were invalid."""
|
||||
request = TestCustomCodeGuardrailRequest(
|
||||
custom_code=(
|
||||
"n = 0\nwhile True:\n n += 1\n\n"
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
|
||||
),
|
||||
test_input={"texts": ["x"]},
|
||||
)
|
||||
started = time.monotonic()
|
||||
|
||||
response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert response.success is False
|
||||
assert response.error_type == "execution"
|
||||
assert response.error is not None
|
||||
assert response.error.startswith("Execution timeout: code took longer than 5 seconds")
|
||||
assert time.monotonic() - started < 8.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_custom_code_endpoint_awaits_an_async_guardrail():
|
||||
request = TestCustomCodeGuardrailRequest(
|
||||
custom_code=(
|
||||
'async def apply_guardrail(inputs, request_data, input_type):\n return block("async said no")\n'
|
||||
),
|
||||
test_input={"texts": ["x"]},
|
||||
)
|
||||
|
||||
response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert response.success is True
|
||||
assert response.result is not None
|
||||
assert response.result["action"] == "block"
|
||||
assert response.result["reason"] == "async said no"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_custom_code_endpoint_returns_a_sync_guardrails_result():
|
||||
request = TestCustomCodeGuardrailRequest(
|
||||
custom_code='def apply_guardrail(inputs, request_data, input_type):\n return block("sync said no")\n',
|
||||
test_input={"texts": ["x"]},
|
||||
)
|
||||
|
||||
response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert response.success is True
|
||||
assert response.result is not None
|
||||
assert response.result["action"] == "block"
|
||||
assert response.result["reason"] == "sync said no"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrail_rolls_back_a_custom_code_guardrail_that_fails_to_compile(mocker, mock_guardrail_registry):
|
||||
stored = {
|
||||
"guardrail_id": "custom-code-broken",
|
||||
"guardrail_name": "custom-code-broken",
|
||||
"litellm_params": {"guardrail": "custom_code", "mode": "pre_call", "custom_code": "x = 1\n"},
|
||||
"guardrail_info": {},
|
||||
}
|
||||
mock_guardrail_registry.add_guardrail_to_db = AsyncMock(return_value=stored)
|
||||
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
|
||||
delete_row = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints._delete_guardrail_row", AsyncMock())
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await create_guardrail(CreateGuardrailRequest(guardrail=stored), user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "apply_guardrail" in exc_info.value.detail
|
||||
delete_row.assert_awaited_once_with(mocker.ANY, where={"guardrail_id": "custom-code-broken"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_custom_code_endpoint_reports_a_system_exit_as_an_execution_error():
|
||||
request = TestCustomCodeGuardrailRequest(
|
||||
custom_code="def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n",
|
||||
test_input={"texts": ["x"]},
|
||||
)
|
||||
started = time.monotonic()
|
||||
|
||||
response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER)
|
||||
|
||||
assert response.success is False
|
||||
assert response.error == "Execution error: SystemExit: bye"
|
||||
assert response.error_type == "execution"
|
||||
assert time.monotonic() - started < 2.0
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch
|
|||
import pytest
|
||||
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeCompilationError
|
||||
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
||||
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
|
@ -335,6 +336,27 @@ def test_init_guardrails_v2_skips_invalid_guardrail_instead_of_crashing_boot():
|
|||
assert "healthy_presidio" in guardrail_names
|
||||
|
||||
|
||||
def test_init_guardrails_v2_stops_boot_when_a_custom_code_guardrail_does_not_compile():
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.clear()
|
||||
IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.clear()
|
||||
|
||||
all_guardrails = [
|
||||
{
|
||||
"guardrail_name": "custom-code-without-apply-guardrail",
|
||||
"litellm_params": {
|
||||
"guardrail": SupportedGuardrailIntegrations.CUSTOM_CODE.value,
|
||||
"mode": "pre_call",
|
||||
"custom_code": "x = 1\n",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"):
|
||||
init_guardrails_v2(all_guardrails=all_guardrails)
|
||||
|
||||
|
||||
def test_init_guardrails_v2_accepts_during_call_advisory_mode():
|
||||
"""
|
||||
Maintainer finding on BerriAI/litellm#34940: on_flagged='inject_system_message'
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue