railtracks.prebuilt.guardrails

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

11class InputLengthGuard(InputGuard):
12    """Blocks LLM input (the full message history) that exceeds ``max_chars`` characters.
13
14    Character counting is used as the simplest, dependency-free unit.  A future
15    implementation may add word- or token-based counting via an optional parameter.
16
17    Example::
18
19        guard = InputLengthGuard(max_chars=4000)
20
21    Args:
22        max_chars: Maximum number of characters allowed across all messages in the
23            input history.  Defaults to ``4096``.
24        name: Optional display name for the guardrail instance.
25        fail_open: Whether to allow the request to continue when this guard raises
26            an unexpected exception.
27    """
28
29    def __init__(
30        self, max_chars: int = 4096, name: str | None = None, fail_open: bool = False
31    ) -> None:
32        super().__init__(name=name, fail_open=fail_open)
33        if max_chars <= 0:
34            raise ValueError(f"max_chars must be a positive integer, got {max_chars!r}")
35        self.max_chars = max_chars
36
37    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
38        total_chars = sum(len(m.content or "") for m in event.messages)
39        if total_chars > self.max_chars:
40            return GuardrailDecision.block(
41                reason=(
42                    f"Input length {total_chars} characters exceeds the maximum of "
43                    f"{self.max_chars} characters."
44                ),
45                user_facing_message=(
46                    "Your message is too long. Please shorten your input and try again."
47                ),
48                meta={"total_chars": total_chars, "max_chars": self.max_chars},
49            )
50        return GuardrailDecision.allow(
51            reason=f"Input length {total_chars} chars is within the {self.max_chars}-char limit.",
52            meta={"total_chars": total_chars, "max_chars": self.max_chars},
53        )

11class OutputLengthGuard(OutputGuard):
12    """Blocks LLM output that exceeds ``max_chars`` characters.
13
14    Inspects ``event.output_message`` (the assistant reply produced this turn).
15
16    Example::
17
18        guard = OutputLengthGuard(max_chars=2000)
19
20    Args:
21        max_chars: Maximum number of characters allowed in the assistant reply.
22            Defaults to ``2048``.
23        name: Optional display name for the guardrail instance.
24    """
25
26    def __init__(
27        self, max_chars: int = 2048, name: str | None = None, fail_open: bool = False
28    ) -> None:
29        super().__init__(name=name, fail_open=fail_open)
30        if max_chars <= 0:
31            raise ValueError(f"max_chars must be a positive integer, got {max_chars!r}")
32        self.max_chars = max_chars
33
34    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
35        if event.output_message is None:
36            return GuardrailDecision.allow(reason="No output message to evaluate.")
37
38        content = event.output_message.content or ""
39        total_chars = len(content)
40        if total_chars > self.max_chars:
41            return GuardrailDecision.block(
42                reason=(
43                    f"Output length {total_chars} characters exceeds the maximum of "
44                    f"{self.max_chars} characters."
45                ),
46                user_facing_message=(
47                    "The response was too long and has been blocked. "
48                    "Please try a more specific question."
49                ),
50                meta={"total_chars": total_chars, "max_chars": self.max_chars},
51            )
52        return GuardrailDecision.allow(
53            reason=f"Output length {total_chars} chars is within the {self.max_chars}-char limit.",
54            meta={"total_chars": total_chars, "max_chars": self.max_chars},
55        )

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

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