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        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.")

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    """
26
27    def __init__(self, max_chars: int = 4096, name: str | None = None) -> None:
28        super().__init__(name=name)
29        if max_chars <= 0:
30            raise ValueError(f"max_chars must be a positive integer, got {max_chars!r}")
31        self.max_chars = max_chars
32
33    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
34        total_chars = sum(len(m.content or "") for m in event.messages)
35        if total_chars > self.max_chars:
36            return GuardrailDecision.block(
37                reason=(
38                    f"Input length {total_chars} characters exceeds the maximum of "
39                    f"{self.max_chars} characters."
40                ),
41                user_facing_message=(
42                    "Your message is too long. Please shorten your input and try again."
43                ),
44                meta={"total_chars": total_chars, "max_chars": self.max_chars},
45            )
46        return GuardrailDecision.allow(
47            reason=f"Input length {total_chars} chars is within the {self.max_chars}-char limit.",
48            meta={"total_chars": total_chars, "max_chars": self.max_chars},
49        )

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__(self, max_chars: int = 2048, name: str | None = None) -> None:
27        super().__init__(name=name)
28        if max_chars <= 0:
29            raise ValueError(f"max_chars must be a positive integer, got {max_chars!r}")
30        self.max_chars = max_chars
31
32    def __call__(self, event: LLMGuardrailEvent) -> GuardrailDecision:
33        if event.output_message is None:
34            return GuardrailDecision.allow(reason="No output message to evaluate.")
35
36        content = event.output_message.content or ""
37        total_chars = len(content)
38        if total_chars > self.max_chars:
39            return GuardrailDecision.block(
40                reason=(
41                    f"Output length {total_chars} characters exceeds the maximum of "
42                    f"{self.max_chars} characters."
43                ),
44                user_facing_message=(
45                    "The response was too long and has been blocked. "
46                    "Please try a more specific question."
47                ),
48                meta={"total_chars": total_chars, "max_chars": self.max_chars},
49            )
50        return GuardrailDecision.allow(
51            reason=f"Output length {total_chars} chars is within the {self.max_chars}-char limit.",
52            meta={"total_chars": total_chars, "max_chars": self.max_chars},
53        )

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        )