Files
5ab8dd3d81 feat: input and output guardrails (block, redact, retry) (#249)
* feat: add input and output guardrails

* chore: use single backticks

* fix: run InputGuard only on the first model request

* fix: replace list[Any] with Sequence[ModelMessage]

* fix: pass raw output to OutputGuard, not str(result.output)

* refactor: organize tests into TestCapabilityName classes

* fix: drain cancelled tasks in InputGuard parallel finally

* fix: re-raise task exceptions via await instead of .exception()

* refactor: consolidate parallel cancel/drain into single finally

* feat: support callable block_message; collapse InputGuard hooks

block_message now accepts a callable so the refusal text can reflect the
prompt/output that tripped the guard, rather than being frozen at
construction time.

InputGuard's sequential path moves from before_model_request into
wrap_model_request, so a single hook covers both sequential and parallel
modes instead of two hooks each branching on `parallel`.

Tests move to tests/guardrails/ to match the tests/<capability>/ layout.

* refactor: guards return bool | GuardResult; drop block_message

A guard now returns either a bare bool or a GuardResult carrying a
refusal message, replacing the separate block_message constructor
field. The message is produced when the guard decides, so it can
reflect the guard's own reasoning rather than a string frozen at
construction time.

Guards (and the GuardResult path) may optionally take a RunContext as
a first parameter, detected from the signature like pydantic-ai's
output validators, so deps- and history-aware guards are possible
without closing over globals. Prompt/output-only guards are unchanged.

* feat: guard outcomes — allow/block/replace/retry with OTel spans

A guard now reports one of four outcomes via GuardResult classmethods
(bool shorthand still works): allow, block, replace, retry.

- replace lets a guard redact rather than refuse — InputGuard rewrites
  the prompt sent to the model, OutputGuard substitutes the output.
- retry lets OutputGuard send a bad output back to the model; OutputGuard
  moves from after_run to after_output_process so it can raise ModelRetry
  and return a modified output.
- replace and block emit spans on the run tracer so a redaction or
  refusal is visible in Logfire; redacted content is included only when
  RunContext.trace_include_content is set.

InputGuard replace requires sequential mode and retry is rejected as a
usage error, since neither is meaningful for input.

* test: cover guardrail test helpers for the 100% gate

The coverage gate measures test files too; the _prompt_text helper had
unreached branches. Drop it and assert on message parts inline.

* docs: scope streaming behavior; add streaming tests

A pydantic-ai-correctness review surfaced two streaming points. Verified
both empirically:

- InputGuard(parallel=True) works under run_stream() — no deadlock.
- OutputGuard GuardResult.retry() is unsupported under run_stream():
  pydantic-ai does not retry output while streaming, so a retry verdict
  surfaces as UnexpectedModelBehavior.

Document the retry limitation and that OutputGuard screens only the final
output (partial chunks reach the caller first while streaming). Note that
input redaction also rewrites persisted history and targets text prompts.
Add streaming tests for both guards to lock the behavior in.

* test: make the parallel-streaming test a real regression guard

The test proving InputGuard(parallel=True) does not deadlock under
run_stream() had no timeout — a reintroduced deadlock would hang CI
instead of failing. Wrap it in asyncio.wait_for and document the
reviewed concern it guards against.

* feat: guardrail capability ordering + GuardResult hardening

Address @adtyavrdhn's review on #249.

- `InputGuard.get_ordering()` → `position='innermost'` so any
  message-morphing capability runs first and the guard sees the final
  prompt the model will receive.
- `OutputGuard.get_ordering()` → `position='outermost', wrapped_by=
  [Instrumentation]` so the guard's block/redact spans are always
  captured by an enclosing `Instrumentation` span regardless of user
  list order.
- `GuardResult` is `frozen=True, kw_only=True` with a `__post_init__`
  that rejects field combinations the four-outcome contract does not
  allow (e.g. `replace` without a replacement). `block` with no message
  stays valid — the default kicks in at the use site.
- `_run_guard` and `after_output_process` dispatch via `match action:`
  with `assert_never` exhaustiveness guards.
- `_trace_block` gates the refusal `message` attribute behind
  `ctx.trace_include_content`, matching `_trace_redaction` — the
  message can quote sensitive content from the guarded value.
- New tests: ordering declarations, `__post_init__` validation, frozen
  enforcement, outer-cancellation no-leak regression guard (which
  confirms `asyncio.shield` around the cleanup `gather` is not needed —
  the outer cancel is already consumed by `asyncio.wait`).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore: bump pydantic-ai-slim to 2.4.0 in lock

CI resolves pydantic packages to newest (uv.toml marks them exclude-newer=false);
the merged lock was stale at 2.1.0 so uv sync --locked failed. Regenerate to 2.4.0.

* test(code_mode): drop removed ToolSearchMatch.description field

pydantic-ai 2.4.0 narrowed ToolSearchMatch to carry only `name`; the full
ToolDefinition (with description) is surfaced separately. Update the test
fixtures to the new shape.

---------

Co-authored-by: Kacper Włodarczyk <kacperwlodarczyk@protonmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-07-10 14:21:23 -05:00

21 lines
668 B
Python

"""Exceptions raised by the guardrail capabilities."""
from __future__ import annotations
class GuardrailError(Exception):
"""Base exception for guardrail violations."""
class InputBlocked(GuardrailError):
"""Raised by a user-supplied input guard to hard-fail a run.
Prefer returning `False` from the guard callable to trigger a graceful
refusal via [`SkipModelRequest`][pydantic_ai.exceptions.SkipModelRequest].
Raise this explicitly when the caller should have to handle the failure.
"""
class OutputBlocked(GuardrailError):
"""Raised by [`OutputGuard`][pydantic_ai_harness.OutputGuard] when the final output fails validation."""