railtracks.guardrails.llm

 25class LLMGuardrailsMixin:
 26    """Mixin for nodes that invoke an LLM.
 27
 28    Overrides ``_pre_invoke`` and ``_post_invoke`` to run input and output guardrails.
 29    Set ``guardrails=`` when building the node. Expects ``self._details`` to contain a
 30    ``guard_details`` list; each run extends it with :class:`~railtracks.guardrails.core.trace.GuardrailTrace`
 31    rows when rails execute.
 32
 33    Output rails run only when ``_post_invoke`` receives a :class:`~railtracks.llm.response.Response`;
 34    other result types are returned unchanged.
 35
 36    Raises:
 37        GuardrailBlockedError: When a rail returns ``BLOCK`` (runner stops the chain
 38            with a blocking decision).
 39    """
 40
 41    guardrails: Guard | None = None
 42    _details: DebugDetails
 43    llm_model: ModelBase
 44    uuid: str
 45    name: Callable[[], str]
 46
 47    def _append_guard_traces(self, traces: list[GuardrailTrace]) -> None:
 48        if not traces:
 49            return
 50        self._details["guard_details"].extend(traces)
 51
 52    def _guardrail_agent_kind(self) -> str:
 53        cls_name = self.__class__.__name__.lower()
 54        if "structured" in cls_name and "toolcall" in cls_name:
 55            return "structured_tool_call"
 56        if "toolcall" in cls_name:
 57            return "tool_call"
 58        if "structured" in cls_name:
 59            return "structured"
 60        if "terminal" in cls_name:
 61            return "terminal"
 62        return "llm"
 63
 64    def _resolve_model_metadata(self) -> tuple[str | None, str | None]:
 65        model_name = getattr(self.llm_model, "model_name", None)
 66        if callable(model_name):
 67            model_name = model_name()
 68        model_provider = getattr(self.llm_model, "model_provider", None)
 69        if callable(model_provider):
 70            model_provider = model_provider()
 71        return (
 72            cast(str | None, model_name),
 73            str(model_provider) if model_provider is not None else None,
 74        )
 75
 76    def _build_input_event(self, context: Any) -> LLMGuardrailEvent:
 77        """Build LLMGuardrailEvent for input phase from context (MessageHistory)."""
 78        model_name, model_provider = self._resolve_model_metadata()
 79        return LLMGuardrailEvent(
 80            phase=LLMGuardrailPhase.INPUT,
 81            messages=context,
 82            node_name=self.name(),
 83            node_uuid=self.uuid,
 84            model_name=model_name,
 85            model_provider=model_provider,
 86            tags={"agent_kind": self._guardrail_agent_kind()},
 87        )
 88
 89    def _build_output_event(
 90        self, context: Any, assistant_message: Message
 91    ) -> LLMGuardrailEvent:
 92        """Build LLMGuardrailEvent for output phase: context is message history; assistant_message is this turn's output."""
 93        model_name, model_provider = self._resolve_model_metadata()
 94        return LLMGuardrailEvent(
 95            phase=LLMGuardrailPhase.OUTPUT,
 96            messages=context,
 97            output_message=assistant_message,
 98            node_name=self.name(),
 99            node_uuid=self.uuid,
100            model_name=model_name,
101            model_provider=model_provider,
102            tags={"agent_kind": self._guardrail_agent_kind()},
103        )
104
105    def _pre_invoke(self, context: Any) -> Any:
106        if self.guardrails is None or not self.guardrails.input:
107            return context
108        event = self._build_input_event(context)
109        new_context, traces, decision = GuardRunner(self.guardrails).run_llm_input(
110            event
111        )
112        self._append_guard_traces(traces)
113        if decision is not None and decision.action == GuardrailAction.BLOCK:
114            rail_name = traces[-1].rail_name if traces else None
115            raise GuardrailBlockedError(
116                rail_name=rail_name,
117                reason=decision.reason,
118                user_facing_message=decision.user_facing_message,
119                traces=traces,
120                meta=decision.meta,
121            )
122
123        return new_context
124
125    def _post_invoke(self, context: Any, result: Any) -> Any:
126        if self.guardrails is None or not self.guardrails.output:
127            return result
128        if not isinstance(result, Response):
129            return result
130        event = self._build_output_event(context, result.message)
131        new_message, traces, decision = GuardRunner(self.guardrails).run_llm_output(
132            event, result.message
133        )
134        self._append_guard_traces(traces)
135        if decision is not None and decision.action == GuardrailAction.BLOCK:
136            rail_name = traces[-1].rail_name if traces else None
137            raise GuardrailBlockedError(
138                rail_name=rail_name,
139                reason=decision.reason,
140                user_facing_message=decision.user_facing_message,
141                traces=traces,
142                meta=decision.meta,
143            )
144
145        return Response(message=new_message, message_info=result.message_info)