railtracks.guardrails.llm

14class BlockTextInputGuard(InputGuard):
15    """Blocks LLM input when any user or system message matches a regex pattern."""
16
17    def __init__(
18        self,
19        pattern: str,
20        *,
21        name: str | None = None,
22        user_facing_message: str | None = None,
23    ) -> None:
24        """Initialize the input block-text guard.
25
26        Args:
27            pattern: Regex pattern; if it matches any scannable message content
28                the guard returns ``BLOCK``.
29            name: Optional rail name for traces (see :class:`InputGuard`).
30            user_facing_message: Optional message surfaced to UIs and
31                visualizers when the guard blocks.
32
33        Raises:
34            re.error: If *pattern* is not a valid regular expression.
35        """
36        super().__init__(name=name)
37        self._pattern = re.compile(pattern)
38        self._user_facing_message = user_facing_message
39
40    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
41        """Block if any user/system string message matches the pattern.
42
43        Returns:
44            ``BLOCK`` when the pattern is found, ``ALLOW`` otherwise.
45        """
46        for msg in event.messages:
47            if msg.role not in _SCANNABLE_ROLES or not isinstance(msg.content, str):
48                continue
49            if self._pattern.search(msg.content):
50                return GuardrailDecision.block(
51                    reason=("Input blocked: prohibited content detected."),
52                    user_facing_message=self._user_facing_message,
53                )
54        return GuardrailDecision.allow(reason="No blocked patterns detected in input.")

11class BlockTextOutputGuard(OutputGuard):
12    """Blocks LLM output when the assistant message matches a regex pattern."""
13
14    def __init__(
15        self,
16        pattern: str,
17        *,
18        name: str | None = None,
19        user_facing_message: str | None = None,
20    ) -> None:
21        """Initialize the output block-text guard.
22
23        Args:
24            pattern: Regex pattern; if it matches the output message content
25                the guard returns ``BLOCK``.
26            name: Optional rail name for traces (see :class:`OutputGuard`).
27            user_facing_message: Optional message surfaced to UIs and
28                visualizers when the guard blocks.
29
30        Raises:
31            re.error: If *pattern* is not a valid regular expression.
32        """
33        super().__init__(name=name)
34        self._pattern = re.compile(pattern)
35        self._user_facing_message = user_facing_message
36
37    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
38        """Block if the output message matches the pattern.
39
40        Returns:
41            ``BLOCK`` when the pattern is found, ``ALLOW`` otherwise.
42        """
43        msg = event.output_message
44        if msg is None or not isinstance(msg.content, str):
45            return GuardrailDecision.allow(reason="No string output to scan.")
46
47        if self._pattern.search(msg.content):
48            return GuardrailDecision.block(
49                reason=("Output blocked: prohibited content detected."),
50                user_facing_message=self._user_facing_message,
51            )
52        return GuardrailDecision.allow(reason="No blocked patterns detected in output.")

 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)

43class PIICustomPattern(BaseModel):
44    """
45    User-defined PII pattern.
46
47    ``name`` becomes the placeholder label: e.g. ``"EMPLOYEE_ID"`` yields
48    ``[EMPLOYEE_ID]`` in redacted text.
49
50    Attributes:
51        name: Label used in placeholders and metadata.
52        regex: Pattern passed to :func:`re.compile` for matching.
53    """
54
55    model_config = ConfigDict(frozen=True)
56
57    name: str
58    regex: str

 9class PIIEntity(str, Enum):
10    """Built-in PII entity types with reliable regex detection."""
11
12    EMAIL_ADDRESS = "EMAIL_ADDRESS"
13    PHONE_NUMBER = "PHONE_NUMBER"
14    CREDIT_CARD = "CREDIT_CARD"
15    US_SSN = "US_SSN"
16    CA_SIN = "CA_SIN"
17    IP_ADDRESS = "IP_ADDRESS"
18    URL = "URL"
19    IBAN_CODE = "IBAN_CODE"
20
21    @classmethod
22    def available(cls) -> dict[str, str]:
23        """Return built-in entity codes and short descriptions for UI or docs.
24
25        Returns:
26            Mapping from entity value string (e.g. ``EMAIL_ADDRESS``) to description.
27        """
28        return {e.value: _ENTITY_DESCRIPTIONS[e] for e in cls}

61class PIIRedactConfig(BaseModel):
62    """
63    Configuration for PII redaction guardrails.
64
65    Frozen so a single instance can safely be shared between input and output
66    guard instances.
67
68    Attributes:
69        entities: Built-in :class:`PIIEntity` kinds to detect; defaults to all members.
70        custom_patterns: Extra :class:`PIICustomPattern` rows merged into detection.
71    """
72
73    model_config = ConfigDict(frozen=True)
74
75    entities: list[PIIEntity] = list(PIIEntity)
76    custom_patterns: list[PIICustomPattern] = []

18class PIIRedactInputGuard(InputGuard):
19    """Redacts PII from user and system string messages before they reach the LLM."""
20
21    def __init__(
22        self,
23        config: PIIRedactConfig | None = None,
24        *,
25        name: str | None = None,
26    ) -> None:
27        """Initialize the input PII redactor.
28
29        Args:
30            config: Redaction settings; defaults to all built-in entity kinds and no
31                custom patterns (see :class:`~railtracks.guardrails.llm.PIIRedactConfig`).
32            name: Optional rail name for traces (see :class:`InputGuard`).
33        """
34        super().__init__(name=name)
35        self._config = config or PIIRedactConfig()
36        self._engine = PIIEngine(self._config)
37
38    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
39        """Scan user/system string content and redact matches.
40
41        Returns:
42            ``ALLOW`` when no PII is found, or ``TRANSFORM`` with rewritten messages on
43            :attr:`~railtracks.guardrails.core.decision.GuardrailDecision.messages`
44            and redaction metadata in ``meta``.
45        """
46        all_records: list[RedactionRecord] = []
47        new_messages: list[Message] = []
48        messages_affected = 0
49
50        for msg in event.messages:
51            if msg.role not in _SCANNABLE_ROLES or not isinstance(msg.content, str):
52                new_messages.append(msg)
53                continue
54
55            redacted_text, records = self._engine.redact(msg.content)
56            if records:
57                all_records.extend(records)
58                messages_affected += 1
59                clone = deepcopy(msg)
60                clone._content = redacted_text
61                new_messages.append(clone)
62            else:
63                new_messages.append(msg)
64
65        if not all_records:
66            return GuardrailDecision.allow(reason="No PII detected in input.")
67
68        return GuardrailDecision.transform_messages(
69            messages=MessageHistory(new_messages),
70            reason=f"Redacted {len(all_records)} PII span(s) from input messages.",
71            meta=build_redaction_meta(all_records, messages_affected=messages_affected),
72        )

14class PIIRedactOutputGuard(OutputGuard):
15    """Redacts PII from the assistant string response after LLM generation."""
16
17    def __init__(
18        self,
19        config: PIIRedactConfig | None = None,
20        *,
21        name: str | None = None,
22    ) -> None:
23        """Initialize the output PII redactor.
24
25        Args:
26            config: Which built-in entities and custom patterns to apply; defaults to
27                all built-in entity kinds.
28            name: Optional rail name for traces (see :class:`OutputGuard`).
29        """
30        super().__init__(name=name)
31        self._config = config or PIIRedactConfig()
32        self._engine = PIIEngine(self._config)
33
34    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
35        """Redact PII from string assistant content on ``event.output_message``.
36
37        Returns:
38            ``ALLOW`` when there is nothing to scan or no PII, or ``TRANSFORM`` with the
39            rewritten message on
40            :attr:`~railtracks.guardrails.core.decision.GuardrailDecision.output_message`
41            and redaction metadata in ``meta``.
42        """
43        msg = event.output_message
44        if msg is None or not isinstance(msg.content, str):
45            return GuardrailDecision.allow(reason="No string output to scan.")
46
47        redacted_text, records = self._engine.redact(msg.content)
48        if not records:
49            return GuardrailDecision.allow(reason="No PII detected in output.")
50
51        clone = deepcopy(msg)
52        clone._content = redacted_text
53        return GuardrailDecision.transform_output(
54            output_message=clone,
55            reason=f"Redacted {len(records)} PII span(s) from output.",
56            meta=build_redaction_meta(records),
57        )