railtracks.guardrails

12class GuardrailAction(str, Enum):
13    """What the runner should do after a guardrail returns.
14
15    Members:
16        ALLOW: Keep the current input or output unchanged.
17        TRANSFORM: Replace input messages or the output message from the decision.
18        BLOCK: Stop and treat the interaction as blocked.
19    """
20
21    ALLOW = "allow"
22    TRANSFORM = "transform"
23    BLOCK = "block"

12class GuardrailBlockedError(NodeInvocationError):
13    """
14    Raised when guardrails deterministically block an operation (e.g. LLM input).
15
16    This error is intended to remain distinguishable from `LLMError` so callers/tests can
17    assert guardrail rejection explicitly.
18    """
19
20    def __init__(
21        self,
22        *,
23        rail_name: str | None = None,
24        reason: str,
25        user_facing_message: str | None = None,
26        traces: list["GuardrailTrace"] | None = None,
27        meta: dict[str, Any] | None = None,
28        notes: list[str] | None = None,
29        fatal: bool = False,
30    ):
31        """Create a block error with optional trace and user-facing context.
32
33        Args:
34            rail_name: Name of the rail that blocked, if known.
35            reason: Machine-oriented explanation (also embedded in the base message).
36            user_facing_message: Optional short text for clients or UIs.
37            traces: Optional list of :class:`~railtracks.guardrails.core.trace.GuardrailTrace`
38                from the failed run.
39            meta: Optional structured details copied from the blocking decision.
40            notes: Extra debug lines forwarded to :class:`NodeInvocationError`.
41            fatal: Passed through to :class:`NodeInvocationError` (whether the run is
42                considered unrecoverable).
43        """
44        self.rail_name = rail_name
45        self.reason = reason
46        self.user_facing_message = user_facing_message
47        self.traces = traces
48        self.meta = meta
49
50        base_message = "Blocked by guardrails"
51        if rail_name:
52            base_message += f" ({rail_name})"
53        base_message += f": {reason}"
54
55        derived_notes: list[str] = []
56        if user_facing_message:
57            derived_notes.append(f"user_message={user_facing_message!r}")
58        if meta:
59            derived_notes.append("meta attached (see exception.meta)")
60
61        super().__init__(
62            message=base_message,
63            notes=[*(notes or []), *derived_notes],
64            fatal=fatal,
65        )

 26class GuardrailDecision(BaseModel):
 27    """Result of one guardrail invocation.
 28
 29    Which fields are set depends on :attr:`action`:
 30
 31    * ``ALLOW``: only :attr:`reason` and optional :attr:`meta` are typically used.
 32    * ``TRANSFORM``: for input phase, :attr:`messages` holds the new history; for
 33      output phase, :attr:`output_message` holds the new assistant message.
 34    * ``BLOCK``: :attr:`user_facing_message` and :attr:`meta` may carry details for
 35      callers or UIs.
 36    """
 37
 38    model_config = ConfigDict(arbitrary_types_allowed=True)
 39
 40    action: GuardrailAction
 41    reason: str
 42    messages: MessageHistory | None = None
 43    output_message: Message | None = None
 44    user_facing_message: str | None = None
 45    meta: dict[str, Any] | None = None
 46
 47    @classmethod
 48    def allow(
 49        cls, reason: str = "Allowed by guardrail.", meta: dict[str, Any] | None = None
 50    ) -> "GuardrailDecision":
 51        """Build an ``ALLOW`` decision with no content changes.
 52
 53        Args:
 54            reason: Explanation for traces and debugging.
 55            meta: Optional extra fields for observability.
 56
 57        Returns:
 58            A decision with :attr:`action` ``ALLOW``.
 59        """
 60        return cls(action=GuardrailAction.ALLOW, reason=reason, meta=meta)
 61
 62    @classmethod
 63    def block(
 64        cls,
 65        reason: str,
 66        user_facing_message: str | None = None,
 67        meta: dict[str, Any] | None = None,
 68    ) -> "GuardrailDecision":
 69        """Build a ``BLOCK`` decision.
 70
 71        Args:
 72            reason: Explanation for logs, traces, and raised errors.
 73            user_facing_message: Optional message safe to show to end users.
 74            meta: Optional extra fields (e.g. exception details).
 75
 76        Returns:
 77            A decision with :attr:`action` ``BLOCK``.
 78        """
 79        return cls(
 80            action=GuardrailAction.BLOCK,
 81            reason=reason,
 82            user_facing_message=user_facing_message,
 83            meta=meta,
 84        )
 85
 86    @classmethod
 87    def transform_messages(
 88        cls,
 89        messages: MessageHistory,
 90        reason: str,
 91        meta: dict[str, Any] | None = None,
 92    ) -> "GuardrailDecision":
 93        """Build a ``TRANSFORM`` decision for LLM input (message history).
 94
 95        Args:
 96            messages: Replacement conversation history for the model call.
 97            reason: Explanation for traces and debugging.
 98            meta: Optional extra fields (e.g. redaction counts).
 99
100        Returns:
101            A decision with :attr:`action` ``TRANSFORM`` and :attr:`messages` set.
102        """
103        return cls(
104            action=GuardrailAction.TRANSFORM,
105            reason=reason,
106            messages=messages,
107            meta=meta,
108        )
109
110    @classmethod
111    def transform_output(
112        cls,
113        output_message: Message,
114        reason: str,
115        meta: dict[str, Any] | None = None,
116    ) -> "GuardrailDecision":
117        """Build a ``TRANSFORM`` decision for LLM output (assistant message).
118
119        Args:
120            output_message: Replacement assistant message to return.
121            reason: Explanation for traces and debugging.
122            meta: Optional extra fields (e.g. redaction counts).
123
124        Returns:
125            A decision with :attr:`action` ``TRANSFORM`` and :attr:`output_message`
126            set.
127        """
128        return cls(
129            action=GuardrailAction.TRANSFORM,
130            reason=reason,
131            output_message=output_message,
132            meta=meta,
133        )

 9class GuardrailTrace(BaseModel):
10    """One guardrail step recorded during a run (for logging or debugging).
11
12    Attributes:
13        rail_name: The guard's :attr:`~railtracks.guardrails.core.interfaces.BaseGuardrail.name`
14            or class name if unset.
15        phase: ``LLMGuardrailPhase`` value string (e.g. ``llm_input``).
16        action: ``allow``, ``transform``, ``block``, or ``error`` when the runner
17            caught an exception, invalid return type, or unknown action.
18        reason: Short explanation, or a fixed message for error traces.
19        meta: Optional details (e.g. :attr:`GuardrailDecision.meta` or exception info).
20    """
21
22    rail_name: str
23    phase: str
24    action: str
25    reason: str
26    meta: dict[str, Any] | None = None

 35class InputGuard(BaseLLMGuardrail[MessageHistory]):
 36    """Base for guardrails that run on LLM input (e.g. prompt / message history).
 37
 38    Attributes:
 39        phase: Always :attr:`LLMGuardrailPhase.INPUT`.
 40    """
 41
 42    phase = LLMGuardrailPhase.INPUT
 43
 44    async def _middleware_fn(
 45        self,
 46        call: Callable[
 47            [MessageHistory, type[BaseModel] | None, list[Tool] | None],
 48            Awaitable[Response],
 49        ],
 50        message_history: MessageHistory,
 51        schema: type[BaseModel] | None,
 52        tools: list[Tool] | None,
 53    ):
 54        """Run this guard on the message history, then call onward with the result."""
 55        message_history, schema, tools = await self._input_wrapper(
 56            message_history, schema, tools
 57        )
 58        return await call(message_history, schema, tools)
 59
 60    async def _input_wrapper(
 61        self,
 62        message_history: MessageHistory,
 63        schema: type[BaseModel] | None,
 64        tools: list[Tool] | None,
 65    ):
 66        """Build the input event, run this guard, and raise if it blocks."""
 67        node_uuid, run_id = self._node_metadata()
 68        event = LLMGuardrailEvent(
 69            phase=LLMGuardrailPhase.INPUT,
 70            messages=message_history,
 71            node_uuid=node_uuid,
 72            run_id=run_id,
 73        )
 74        input_event = MiddlewareGuardInputInvocationEvent(
 75            message_history=message_history,
 76        )
 77
 78        await emit(input_event)
 79
 80        try:
 81            new_messages, traces, decision = await self.run(
 82                event=event, value=message_history
 83            )
 84        except Exception as e:
 85            failure_event = MiddlewareGuardInputFailureEvent.from_exception(e)
 86            await emit(failure_event)
 87            raise e
 88
 89        result_event = MiddlewareGuardInputResponseEvent(
 90            decision=decision,
 91            message_history=new_messages,
 92        )
 93
 94        await emit(result_event)
 95
 96        self._raise_if_blocked(decision, traces)
 97
 98        return new_messages, schema, tools
 99
100    def convert(
101        self,
102        value: str | Any | MessageHistory | LLMGuardrailEvent,
103        /,
104    ):
105        """Run this guard without building an :class:`LLMGuardrailEvent` by hand.
106
107        Args:
108            value: A :class:`LLMGuardrailEvent` (passed through), a ``str`` (treated
109                as a single user message), a :class:`~railtracks.llm.message.Message`,
110                or a :class:`~railtracks.llm.history.MessageHistory`.
111
112        Returns:
113            The :class:`GuardrailDecision` from :meth:`__call__`.
114
115        Raises:
116            TypeError: If ``value`` is not a ``str``, ``Message``, ``MessageHistory``,
117                or :class:`LLMGuardrailEvent`.
118        """
119        if isinstance(value, LLMGuardrailEvent):
120            return value
121
122        messages = self._coerce_to_message_history(value)
123        event = LLMGuardrailEvent(
124            phase=LLMGuardrailPhase.INPUT,
125            messages=messages,
126        )
127        return event
128
129    def _extract_transform_value(self, decision: GuardrailDecision) -> Any:
130        """Return the replacement messages from a TRANSFORM decision."""
131        if decision.messages is None:
132            raise ValueError(
133                "Input guardrail returned TRANSFORM without decision.messages."
134            )
135        return decision.messages
136
137    def _sync_event_after_transform(
138        self,
139        event: LLMGuardrailEvent,
140        value: MessageHistory,
141    ) -> LLMGuardrailEvent:
142        """Return a copy of event with its messages replaced."""
143        return event.model_copy(update={"messages": value})

26class LLMGuardrailEvent(BaseModel):
27    """Payload for LLM input and output guardrails.
28
29    Attributes:
30        phase: Whether this is an input-phase or output-phase check.
31        messages: Conversation context, usually the history before the current
32            assistant reply is appended by the node.
33        output_message: For ``OUTPUT`` phase, the assistant :class:`~railtracks.llm.message.Message`
34            under inspection; ``None`` for input phase or when not applicable.
35        node_name: Optional node label for observability.
36        node_uuid: Optional stable node id for observability.
37        run_id: Optional run correlation id.
38        model_name: Optional resolved model name for observability.
39        model_provider: Optional provider string for observability.
40        tags: Optional key/value metadata (e.g. ``agent_kind`` from the mixin).
41    """
42
43    model_config = ConfigDict(arbitrary_types_allowed=True)
44
45    phase: LLMGuardrailPhase
46    messages: MessageHistory
47    output_message: Message | None = None
48
49    node_name: str | None = None
50    node_uuid: str | None = None
51    run_id: str | None = None
52    model_name: str | None = None
53    model_provider: str | None = None
54    tags: dict[str, str] | None = None

14class LLMGuardrailPhase(str, Enum):
15    """Which side of the LLM call a guardrail observes.
16
17    Members:
18        INPUT: Before the model (prompt / history), value ``llm_input``.
19        OUTPUT: After the model (assistant message), value ``llm_output``.
20    """
21
22    INPUT = "llm_input"
23    OUTPUT = "llm_output"

154class OutputGuard(BaseLLMGuardrail[Message]):
155    """Base for guardrails that run on LLM output (e.g. model response).
156
157    Inspect ``event.output_message`` for the assistant message produced this turn.
158    ``event.messages`` is conversation context and may not yet include that reply.
159
160    Intermediate tool-call turns pass through untouched, so output rails fire only
161    on the final reply.
162
163    Attributes:
164        phase: Always :attr:`LLMGuardrailPhase.OUTPUT`.
165    """
166
167    phase = LLMGuardrailPhase.OUTPUT
168
169    async def _middleware_fn(
170        self,
171        call: Callable[
172            [MessageHistory, type[BaseModel] | None, list[Tool] | None],
173            Awaitable[Response],
174        ],
175        message_history: MessageHistory,
176        schema: type[BaseModel] | None,
177        tools: list[Tool] | None,
178    ):
179        """Call onward for the response, then run this guard on the final reply.
180
181        Responses that request tools are intermediate steps of the tool-calling
182        loop and pass through unguarded.
183        """
184        result = await call(message_history, schema, tools)
185        if _is_intermediate_tool_call(result):
186            return result
187
188        return await self._output_wrapper(result=result)
189
190    async def _output_wrapper(
191        self,
192        result: Response,
193    ):
194        """Build the output event, run this guard, and rebuild the response if the message changed."""
195        node_uuid, run_id = self._node_metadata()
196        event = LLMGuardrailEvent(
197            phase=LLMGuardrailPhase.OUTPUT,
198            messages=MessageHistory([]),
199            output_message=result.message,
200            node_uuid=node_uuid,
201            run_id=run_id,
202        )
203
204        input_event = MiddlewareGuardOutputInvocationEvent(
205            response=result.message,
206        )
207
208        await emit(input_event)
209        try:
210            new_message, traces, decision = await self.run(
211                event=event, value=result.message
212            )
213        except Exception as e:
214            failure_event = MiddlewareGuardOutputFailureEvent.from_exception(e)
215            await emit(failure_event)
216            raise e
217
218        output_event = MiddlewareGuardOutputResponseEvent(
219            decision=decision,
220            response=new_message,
221        )
222        await emit(output_event)
223
224        self._raise_if_blocked(decision, traces)
225
226        if new_message is result.message:
227            return result
228
229        return Response(message=new_message, message_info=result.message_info)
230
231    def convert(self, output: str | Any | MessageHistory | LLMGuardrailEvent, /):
232        """Run this guard without building an :class:`LLMGuardrailEvent` by hand.
233
234        Args:
235            output: A :class:`LLMGuardrailEvent` (passed through), a ``str`` (becomes
236                the assistant message with empty prior history), a
237                :class:`~railtracks.llm.message.Message`, or a non-empty
238                :class:`~railtracks.llm.history.MessageHistory` (last message is the
239                output under test; earlier entries become ``event.messages``).
240
241        Returns:
242            The :class:`GuardrailDecision` from :meth:`__call__`.
243
244        Raises:
245            ValueError: If ``output`` is an empty :class:`~railtracks.llm.history.MessageHistory`.
246            TypeError: If ``output`` is not a ``str``, ``Message``, ``MessageHistory``,
247                or :class:`LLMGuardrailEvent`.
248        """
249        if isinstance(output, LLMGuardrailEvent):
250            return output
251
252        if isinstance(output, str):
253            output_message = AssistantMessage(output)
254            messages = MessageHistory()
255        elif isinstance(output, Message):
256            output_message = output
257            messages = MessageHistory()
258        elif isinstance(output, MessageHistory):
259            if not output:
260                raise ValueError("Cannot decide with an empty MessageHistory.")
261            output_message = output[-1]
262            messages = MessageHistory(output[:-1])
263        else:
264            raise TypeError(
265                f"Expected str, Message, MessageHistory, or LLMGuardrailEvent, "
266                f"got {type(output).__name__}"
267            )
268
269        event = LLMGuardrailEvent(
270            phase=LLMGuardrailPhase.OUTPUT,
271            messages=messages,
272            output_message=output_message,
273        )
274        return event
275
276    def _extract_transform_value(self, decision: GuardrailDecision) -> Message:
277        """Return the replacement output message from a TRANSFORM decision."""
278        if decision.output_message is None:
279            raise ValueError(
280                "Output guardrail returned TRANSFORM without decision.output_message."
281            )
282        return decision.output_message
283
284    def _sync_event_after_transform(
285        self,
286        event: LLMGuardrailEvent,
287        value: Message,
288    ) -> LLMGuardrailEvent:
289        """Return a copy of event with its output message replaced."""
290        return event.model_copy(update={"output_message": cast(Message, value)})

109def input_guard(
110    fn: _GuardFn | None = None,
111    *,
112    name: str | None = None,
113    fail_open: bool = False,
114):
115    """Turn a function into an :class:`InputGuard` instance.
116
117    The function receives an :class:`LLMGuardrailEvent` (INPUT phase; inspect
118    ``event.messages``) and returns a :class:`GuardrailDecision`. It may be sync or
119    ``async def``; an async rail is awaited, so it can ``await rt.call(...)``.
120
121    Usable bare or parameterized::
122
123        @rt.input_guard
124        def guard(event): ...
125
126
127        @rt.input_guard(name="my_rail", fail_open=True)
128        async def guard(event): ...
129
130    Args:
131        fn: The guard function (supplied automatically in the bare form).
132        name: Rail name for traces; defaults to the function name.
133        fail_open: Allow the request through if the guard raises unexpectedly.
134
135    Returns:
136        An :class:`InputGuard` instance in the bare form, or a decorator in the
137        parameterized form.
138    """
139
140    def decorate(func: _GuardFn, /) -> InputGuard:
141        return _make_guard(InputGuard, func, name=name, fail_open=fail_open)
142
143    if fn is not None:
144        return decorate(fn)
145    return decorate

154def output_guard(
155    fn: _GuardFn | None = None,
156    *,
157    name: str | None = None,
158    fail_open: bool = False,
159):
160    """Turn a function into an :class:`OutputGuard` instance.
161
162    The function receives an :class:`LLMGuardrailEvent` (OUTPUT phase; inspect
163    ``event.output_message``) and returns a :class:`GuardrailDecision`. It may be
164    sync or ``async def``; an async rail is awaited, so it can ``await rt.call(...)``.
165    Intermediate tool-call turns are skipped by :class:`OutputGuard`, so the
166    function fires only on the final reply.
167
168    Usable bare or parameterized::
169
170        @rt.output_guard
171        def guard(event): ...
172
173
174        @rt.output_guard(name="my_rail", fail_open=True)
175        async def guard(event): ...
176
177    Args:
178        fn: The guard function (supplied automatically in the bare form).
179        name: Rail name for traces; defaults to the function name.
180        fail_open: Allow the response through if the guard raises unexpectedly.
181
182    Returns:
183        An :class:`OutputGuard` instance in the bare form, or a decorator in the
184        parameterized form.
185    """
186
187    def decorate(func: _GuardFn, /) -> OutputGuard:
188        return _make_guard(OutputGuard, func, name=name, fail_open=fail_open)
189
190    if fn is not None:
191        return decorate(fn)
192    return decorate