railtracks.guardrails.llm.decorators

  1"""Decorators for authoring guardrails from a plain function.
  2
  3Mirror the ``@pre_llm`` / ``@post_llm`` middleware decorators: wrap a function
  4that maps an event to a :class:`GuardrailDecision` and get back a ready-to-use
  5:class:`InputGuard` / :class:`OutputGuard` instance.
  6
  7Example::
  8
  9    @rt.input_guard
 10    def block_secrets(
 11        event: rt.guardrails.LLMGuardrailEvent,
 12    ) -> rt.guardrails.GuardrailDecision:
 13        for msg in event.messages:
 14            if isinstance(msg.content, str) and "SECRET" in msg.content:
 15                return rt.guardrails.GuardrailDecision.block(reason="secret leaked")
 16        return rt.guardrails.GuardrailDecision.allow()
 17
 18
 19    @rt.output_guard(fail_open=True)
 20    def no_profanity(event) -> rt.guardrails.GuardrailDecision: ...
 21
 22The decorated guard is callable both with an :class:`LLMGuardrailEvent` (how the
 23guard's middleware invokes it) and with a raw ``str`` / :class:`~railtracks.llm.message.Message`
 24/ :class:`~railtracks.llm.history.MessageHistory`, which is coerced to the correct
 25event for the phase via :meth:`convert` before the function sees it.
 26
 27The wrapped function may be sync or async. An ``async def`` guard is awaited while
 28its middleware evaluates the rail (``_eval_one_rail``), so it can ``await`` anything
 29-- including ``rt.call`` into another agent, which is how you build an LLM-judge
 30rail::
 31
 32    @rt.input_guard
 33    async def llm_judge(event):
 34        verdict = await rt.call(Judge, event.messages[-1].content)
 35        if "UNSAFE" in str(verdict):
 36            return rt.guardrails.GuardrailDecision.block(reason="judge flagged input")
 37        return rt.guardrails.GuardrailDecision.allow()
 38
 39Note that a rail runs once per *model round-trip*, not once per agent call, so an
 40async rail on a tool-calling agent fires on every iteration of the tool loop.
 41"""
 42
 43from __future__ import annotations
 44
 45import inspect
 46from typing import Awaitable, Callable, TypeVar, cast, overload
 47
 48from railtracks.guardrails.core.decision import GuardrailDecision
 49from railtracks.guardrails.core.event import LLMGuardrailEvent
 50from railtracks.guardrails.llm.concrete import InputGuard, OutputGuard
 51from railtracks.guardrails.llm.llm_guard import BaseLLMGuardrail
 52
 53_GuardFn = Callable[
 54    [LLMGuardrailEvent], GuardrailDecision | Awaitable[GuardrailDecision]
 55]
 56_GuardT = TypeVar("_GuardT", bound=BaseLLMGuardrail)
 57
 58
 59def _make_guard(
 60    base: type[_GuardT],
 61    fn: _GuardFn,
 62    *,
 63    name: str | None,
 64    fail_open: bool,
 65) -> _GuardT:
 66    """Build an ``InputGuard``/``OutputGuard`` instance that delegates to ``fn``.
 67
 68    The generated guard coerces any non-event input to an event via the base's
 69    phase-aware :meth:`convert`, so ``fn`` always receives an
 70    :class:`LLMGuardrailEvent`.
 71
 72    An ``async def`` ``fn`` produces a guard with an ``async def __call__``, which the
 73    rail evaluator awaits. Everything downstream of the decision is identical either
 74    way.
 75    """
 76    guard_name = name or fn.__name__
 77
 78    if inspect.iscoroutinefunction(fn):
 79
 80        class _FunctionGuard(base):  # type: ignore[valid-type, misc]
 81            async def __call__(self, event) -> GuardrailDecision:
 82                if not isinstance(event, LLMGuardrailEvent):
 83                    event = self.convert(event)
 84                return await fn(event)
 85
 86    else:
 87
 88        class _FunctionGuard(base):  # type: ignore[valid-type, misc, no-redef]
 89            def __call__(self, event) -> GuardrailDecision:
 90                if not isinstance(event, LLMGuardrailEvent):
 91                    event = self.convert(event)
 92                return cast(GuardrailDecision, fn(event))
 93
 94    _FunctionGuard.__name__ = f"{base.__name__}[{guard_name}]"
 95    _FunctionGuard.__qualname__ = _FunctionGuard.__name__
 96    _FunctionGuard.__doc__ = fn.__doc__
 97
 98    guard_cls = cast("type[_GuardT]", _FunctionGuard)
 99    return guard_cls(name=guard_name, fail_open=fail_open)
100
101
102@overload
103def input_guard(fn: _GuardFn, /) -> InputGuard: ...
104@overload
105def input_guard(
106    *, name: str | None = ..., fail_open: bool = ...
107) -> Callable[[_GuardFn], InputGuard]: ...
108def input_guard(
109    fn: _GuardFn | None = None,
110    *,
111    name: str | None = None,
112    fail_open: bool = False,
113):
114    """Turn a function into an :class:`InputGuard` instance.
115
116    The function receives an :class:`LLMGuardrailEvent` (INPUT phase; inspect
117    ``event.messages``) and returns a :class:`GuardrailDecision`. It may be sync or
118    ``async def``; an async rail is awaited, so it can ``await rt.call(...)``.
119
120    Usable bare or parameterized::
121
122        @rt.input_guard
123        def guard(event): ...
124
125
126        @rt.input_guard(name="my_rail", fail_open=True)
127        async def guard(event): ...
128
129    Args:
130        fn: The guard function (supplied automatically in the bare form).
131        name: Rail name for traces; defaults to the function name.
132        fail_open: Allow the request through if the guard raises unexpectedly.
133
134    Returns:
135        An :class:`InputGuard` instance in the bare form, or a decorator in the
136        parameterized form.
137    """
138
139    def decorate(func: _GuardFn, /) -> InputGuard:
140        return _make_guard(InputGuard, func, name=name, fail_open=fail_open)
141
142    if fn is not None:
143        return decorate(fn)
144    return decorate
145
146
147@overload
148def output_guard(fn: _GuardFn, /) -> OutputGuard: ...
149@overload
150def output_guard(
151    *, name: str | None = ..., fail_open: bool = ...
152) -> Callable[[_GuardFn], OutputGuard]: ...
153def output_guard(
154    fn: _GuardFn | None = None,
155    *,
156    name: str | None = None,
157    fail_open: bool = False,
158):
159    """Turn a function into an :class:`OutputGuard` instance.
160
161    The function receives an :class:`LLMGuardrailEvent` (OUTPUT phase; inspect
162    ``event.output_message``) and returns a :class:`GuardrailDecision`. It may be
163    sync or ``async def``; an async rail is awaited, so it can ``await rt.call(...)``.
164    Intermediate tool-call turns are skipped by :class:`OutputGuard`, so the
165    function fires only on the final reply.
166
167    Usable bare or parameterized::
168
169        @rt.output_guard
170        def guard(event): ...
171
172
173        @rt.output_guard(name="my_rail", fail_open=True)
174        async def guard(event): ...
175
176    Args:
177        fn: The guard function (supplied automatically in the bare form).
178        name: Rail name for traces; defaults to the function name.
179        fail_open: Allow the response through if the guard raises unexpectedly.
180
181    Returns:
182        An :class:`OutputGuard` instance in the bare form, or a decorator in the
183        parameterized form.
184    """
185
186    def decorate(func: _GuardFn, /) -> OutputGuard:
187        return _make_guard(OutputGuard, func, name=name, fail_open=fail_open)
188
189    if fn is not None:
190        return decorate(fn)
191    return decorate