mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
5ab8dd3d81
* 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>
26 lines
504 B
Python
26 lines
504 B
Python
"""Input and output guardrails for Pydantic AI agents."""
|
|
|
|
from pydantic_ai_harness.guardrails._capability import (
|
|
GuardResult,
|
|
InputGuard,
|
|
InputGuardFunc,
|
|
OutputGuard,
|
|
OutputGuardFunc,
|
|
)
|
|
from pydantic_ai_harness.guardrails._exceptions import (
|
|
GuardrailError,
|
|
InputBlocked,
|
|
OutputBlocked,
|
|
)
|
|
|
|
__all__ = [
|
|
'GuardResult',
|
|
'GuardrailError',
|
|
'InputBlocked',
|
|
'InputGuard',
|
|
'InputGuardFunc',
|
|
'OutputBlocked',
|
|
'OutputGuard',
|
|
'OutputGuardFunc',
|
|
]
|