From f467136fa8b9fcea4c84d41c0c0a13e812a92672 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sun, 14 Jun 2026 21:30:34 +0530 Subject: [PATCH] feat(sdk): PydanticAI adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bind a PydanticAI agent to a governed run by wrapping its model: `Agent(govern(model, run))` — no other change to the agent. One model request counts as one governed step, so the deterministic loop/time budget halts a runaway agent, and with `gate_tools=True` each tool call the model proposes routes through the human-approval gate before the agent executes it. I built this on PydanticAI's `WrapperModel`, which forwards every model method to the wrapped model, so I only override `request`/`request_stream`. That keeps it to one stable surface (the `Model.request` contract, which hasn't changed across the post-1.0 line) and covers both loop/time enforcement and tool gating: the tool calls a step will make are present in the model's response before the agent runs them, so I gate them right after the response comes back. The halt propagates: PydanticAI only retries on its own `ModelRetry` signal and re-raises every other model-request error by default, so I raise plain `BudgetExceeded`/`ApprovalDenied` (never `ModelRetry`) and they bubble out of `agent.run()`/`run_sync()` and stop the agent rather than being retried into another paid request. I verified this against a real pydantic-ai install with a FunctionModel that loops forever. `pydantic-ai` is lazily imported and stays an optional extra, so the SDK still installs and imports with no third-party deps. Supported against pydantic-ai (pydantic-ai-slim) >= 1, < 2. --- CHANGELOG.md | 14 + sdks/python/README.md | 5 + sdks/python/pyproject.toml | 1 + sdks/python/riskkernel/adapters/__init__.py | 1 + .../python/riskkernel/adapters/pydantic_ai.py | 170 +++++++++++ sdks/python/tests/test_pydantic_ai.py | 284 ++++++++++++++++++ 6 files changed, 475 insertions(+) create mode 100644 sdks/python/riskkernel/adapters/pydantic_ai.py create mode 100644 sdks/python/tests/test_pydantic_ai.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 86f532e..49e534f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,20 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). `step_callback` is synchronous and re-raised by the executor, unlike its fire-and-forget event bus, which would drop it). `crewai` is lazily imported, so it stays an optional dependency; supported against `crewai` >= 0.80, < 2. +- **PydanticAI adapter (Python SDK).** `from riskkernel.adapters.pydantic_ai import + govern` — wrap your model with `Agent(govern(model, run))` to bind a PydanticAI + agent to a governed run with no other code change. One model request counts as one + governed step, so the deterministic loop/time budget halts a runaway agent; with + `gate_tools=True` each tool call the model proposes routes through the human-approval + gate before the agent executes it. The halt is raised from the model wrapper and + propagates out of `agent.run()` / `agent.run_sync()` — PydanticAI only retries its + own `ModelRetry` signal, so the deterministic `BudgetExceeded`/`ApprovalDenied` + surfaces to the caller and stops the agent rather than being retried into another + paid request. Built on PydanticAI's `WrapperModel` contract, which forwards every + model method to the wrapped model and is stable across the post-1.0 line; both the + non-streaming and streaming request paths are governed. `pydantic-ai` is lazily + imported, so it stays an optional dependency; supported against `pydantic-ai` + (`pydantic-ai-slim`) >= 1, < 2. - **Streaming proxy.** Both `POST /v1/chat/completions` and `POST /v1/messages` now support `stream:true`: the budget is enforced before the stream opens, the provider's SSE is forwarded to the client verbatim (authentic OpenAI or Anthropic diff --git a/sdks/python/README.md b/sdks/python/README.md index 4d903d5..155c9a0 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -103,6 +103,11 @@ crew = Crew(agents=[...], tasks=[...], step_callback=RiskKernelStepCallback(run) from riskkernel.adapters.autogen import GovernedChatCompletionClient client = GovernedChatCompletionClient(model_client, run) # drop-in for the real client agent = AssistantAgent("assistant", model_client=client) # halts the agent at budget + +# PydanticAI — wrap your model (one governed step per model request; halts the agent at budget) +from pydantic_ai import Agent +from riskkernel.adapters.pydantic_ai import govern +agent = Agent(govern("anthropic:claude-sonnet-4-5", run, gate_tools=True)) ``` > AutoGen halts the run either way, but a *team* (`RoundRobinGroupChat`, …) re-raises diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index 59b5384..5ae2bc9 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -28,6 +28,7 @@ Source = "https://github.com/prashar32/riskkernel" langchain = ["langchain-core"] crewai = ["crewai>=0.80,<2"] autogen = ["autogen-agentchat>=0.4,<1", "autogen-core>=0.4,<1"] +pydantic-ai = ["pydantic-ai-slim>=1,<2"] dev = ["pytest"] # This project lives in a subdirectory (sdks/python) of the repo. Hatchling's diff --git a/sdks/python/riskkernel/adapters/__init__.py b/sdks/python/riskkernel/adapters/__init__.py index 06f800b..3085382 100644 --- a/sdks/python/riskkernel/adapters/__init__.py +++ b/sdks/python/riskkernel/adapters/__init__.py @@ -8,4 +8,5 @@ - ``openai_agents`` — RunHooks for the OpenAI Agents SDK (steps + approval gate). - ``crewai`` — a step_callback for CrewAI (steps + tool approval gate). - ``autogen`` — a model-client wrapper for AutoGen v0.4+ (steps + tool gate). +- ``pydantic_ai`` — a model wrapper for PydanticAI (steps + tool approval gate). """ diff --git a/sdks/python/riskkernel/adapters/pydantic_ai.py b/sdks/python/riskkernel/adapters/pydantic_ai.py new file mode 100644 index 0000000..338c2fc --- /dev/null +++ b/sdks/python/riskkernel/adapters/pydantic_ai.py @@ -0,0 +1,170 @@ +"""PydanticAI adapter: a model wrapper that enforces a governed run's loop and time +budgets, ticking one governed step per model request, and (optionally) gates the +tool calls a model proposes through the approval gate. Wrap your real model with it +and hand the wrapper to your ``Agent`` — no other change to your agent code: + + from pydantic_ai import Agent + from riskkernel.adapters.pydantic_ai import govern + + agent = Agent(govern("anthropic:claude-sonnet-4-5", run)) + agent.run_sync("...") # halts with BudgetExceeded when the budget is spent + +``govern`` accepts either a model name (resolved by PydanticAI) or a ``Model`` +instance, so it slots in wherever you already build your model: + + base = AnthropicModel("claude-sonnet-4-5") + agent = Agent(govern(base, run, gate_tools=True)) + +A BudgetExceeded (or ApprovalDenied) raised inside the wrapped request propagates +out of ``agent.run()`` / ``agent.run_sync()`` and halts the agent. PydanticAI only +retries a model request on its own ``ModelRetry`` signal — every other exception is +terminal and bubbles out of the run (its model-request error handler re-raises by +default). So the deterministic halt surfaces to the caller and stops the agent; we +deliberately raise plain SDK exceptions (NOT ``ModelRetry``) so they are never +retried into another paid request. + +Why wrap the model rather than hook tool execution: one model request maps cleanly +to one governed step (the agent's outer loop is request → tool calls → request …), +and the tool calls a step will make are present in the model's *response* before the +agent executes them — so a single, stable surface (the ``Model.request`` contract, +unchanged across the 1.x line) covers both loop/time enforcement and tool gating +without depending on a newer hooks API. We gate the proposed tool calls right after +the response comes back and before the agent runs them; a denial raises and halts. + +Supported API: PydanticAI's ``Model`` / ``WrapperModel`` contract — pinned and +tested against ``pydantic-ai`` (``pydantic-ai-slim``) >= 1, < 2 (the post-1.0 line, +which commits to no breaking changes before 2.0; ``Model.request`` takes +``(messages, model_settings, model_request_parameters)`` and returns a +``ModelResponse``). PydanticAI is lazily imported and is NOT a dependency of the +SDK: this module imports cleanly without it, and the wrapper is only constructed +when you actually wire it. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..approval import ApprovalGate +from ..runtime import Run + + +def _wrapper_base(): + # Inherit PydanticAI's WrapperModel when it is installed (it forwards every + # Model method to the wrapped model, so we override only request/request_stream); + # otherwise fall back to object so the module still imports without pydantic-ai. + try: + from pydantic_ai.models.wrapper import WrapperModel # type: ignore + return WrapperModel + except Exception: + return object + + +class GovernedModel(_wrapper_base()): # type: ignore[misc] + """A PydanticAI model wrapper that binds a governed run to an agent. + + It delegates every request to the wrapped model, but first ticks one governed + step (enforcing the loop/time budget) and — when ``gate_tools`` is set — routes + each tool call the model proposes through the approval gate before the agent + executes it. + + Build it via :func:`govern` (which resolves a model name or instance), or + construct it directly with an already-built ``Model``. + + Args: + wrapped: the real PydanticAI ``Model`` (or a model name) to govern. + run: the governed Run. + gate_tools: if True, every tool call the model proposes must pass the + approval gate before the agent runs it; a denial raises and halts. + tool_side_effect: side-effect label reported for gated tools. + timeout: max seconds to await a human decision on a gated tool. + """ + + def __init__(self, wrapped: Any, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool", timeout: Optional[float] = None): + # WrapperModel.__init__ resolves a model name to a Model and stores it on + # self.wrapped. Only call it when we actually inherit it (not the object + # fallback), so the module is importable without pydantic-ai present. + base = type(self).__mro__[1] + if base is not object: + base.__init__(self, wrapped) + else: + self.wrapped = wrapped + self.run = run + self.gate_tools = gate_tools + self.tool_side_effect = tool_side_effect + self.timeout = timeout + self._gate = ApprovalGate(run) + + # One model request == one governed step. We tick the step first (so the loop/ + # time budget halts before we spend on the request), then delegate, then gate the + # tool calls in the response before the agent executes them. + async def request(self, messages: Any, model_settings: Any, + model_request_parameters: Any) -> Any: + self.run.step() # raises BudgetExceeded when the loop/time budget is spent + response = await self.wrapped.request( + messages, model_settings, model_request_parameters) + if self.gate_tools: + self._gate_response_tools(response) + return response + + # Streaming takes the same path: tick a step, then stream from the wrapped model. + # We don't gate tools here — the parts aren't known until the stream is consumed, + # and the non-streaming request path is where tool calls are gated. Decorated as + # an async context manager only when WrapperModel is present (matching the base). + def request_stream(self, messages: Any, model_settings: Any, + model_request_parameters: Any, + run_context: Any = None) -> Any: + self.run.step() # raises BudgetExceeded when the loop/time budget is spent + return self.wrapped.request_stream( + messages, model_settings, model_request_parameters, run_context) + + def _gate_response_tools(self, response: Any) -> None: + """Route each tool call the model proposed through the approval gate. A denial + raises ApprovalDenied, which propagates out of the run and halts the agent + before the tool executes.""" + for call in _tool_calls(response): + name = getattr(call, "tool_name", None) or "tool" + self._gate.require(str(name), side_effect=self.tool_side_effect, + arguments={"args": _stringify(getattr(call, "args", None))}, + timeout=self.timeout) + + +def govern(model: Any, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool", + timeout: Optional[float] = None) -> GovernedModel: + """Wrap a PydanticAI model (or model name) so its agent run is governed. + + Pass the result wherever you'd pass a model:: + + agent = Agent(govern("anthropic:claude-sonnet-4-5", run)) + + Args: + model: a PydanticAI ``Model`` instance or a model name string. + run: the governed Run. + gate_tools: gate proposed tool calls through the approval gate (default off). + tool_side_effect: side-effect label reported for gated tools. + timeout: max seconds to await a human decision on a gated tool. + """ + return GovernedModel(model, run, gate_tools=gate_tools, + tool_side_effect=tool_side_effect, timeout=timeout) + + +def _tool_calls(response: Any): + """Yield the ToolCallPart-like parts of a PydanticAI ModelResponse. Duck-typed + (a tool call has a ``tool_name`` attribute) so it works across versions and + without importing pydantic-ai for the unit tests.""" + parts = getattr(response, "parts", None) + if not parts: + return + for part in parts: + if getattr(part, "tool_name", None) is not None: + yield part + + +def _stringify(v: Any) -> Any: + try: + import json + json.dumps(v) + return v + except Exception: + return repr(v) diff --git a/sdks/python/tests/test_pydantic_ai.py b/sdks/python/tests/test_pydantic_ai.py new file mode 100644 index 0000000..7ba2073 --- /dev/null +++ b/sdks/python/tests/test_pydantic_ai.py @@ -0,0 +1,284 @@ +"""PydanticAI adapter tests — governance behavior against a fake Run, no daemon, no +third-party deps. The wrapper imports even without pydantic-ai installed (lazy base +class), so these exercise the enforcement path on stdlib alone; a real PydanticAI +integration test is gated behind skipUnless. +""" + +import asyncio +import unittest + +from riskkernel.adapters.pydantic_ai import GovernedModel, govern, _tool_calls +from riskkernel.errors import ApprovalDenied, BudgetExceeded + + +def _has_pydantic_ai() -> bool: + try: + import pydantic_ai # noqa: F401 + + return True + except Exception: + return False + + +class FakeRun: + """A stand-in for runtime.Run: counts steps and halts past a loop budget, + mirroring how the daemon's BeginStep raises BudgetExceeded when the budget is + spent. No HTTP, no daemon.""" + + def __init__(self, loop_budget=None): + self.steps = 0 + self.loop_budget = loop_budget + self.id = "run-1" + + def step(self): + self.steps += 1 + if self.loop_budget is not None and self.steps > self.loop_budget: + raise BudgetExceeded("loop_budget_exceeded") + return self.steps + + +def _model_base(): + # When pydantic-ai is installed, GovernedModel inherits WrapperModel, whose + # __init__ runs infer_model(wrapped) — which only accepts a real Model instance + # (or a model-name string). So the fake model must BE a real Model for these unit + # tests to construct GovernedModel. When pydantic-ai is absent, the adapter falls + # back to object and stores the wrapped model directly, so a plain object works. + try: + from pydantic_ai.models import Model # type: ignore + + return Model + except Exception: + return object + + +class _FakeModel(_model_base()): # type: ignore[misc] + """A stand-in for a PydanticAI Model: records that it was asked, and returns a + canned response. Lets us exercise GovernedModel.request without (and with) + pydantic-ai installed.""" + + def __init__(self, response): + self.response = response + self.requests = 0 + + async def request(self, messages, model_settings, model_request_parameters): + self.requests += 1 + return self.response + + # Model is an ABC with abstract model_name/system properties; provide them so the + # subclass is instantiable when pydantic-ai is present. + @property + def model_name(self) -> str: + return "fake" + + @property + def system(self) -> str: + return "fake" + + +class _Resp: + """A minimal ModelResponse stand-in: has `.parts` like the real one.""" + + def __init__(self, parts): + self.parts = parts + + +class _ToolCallPart: + """A minimal ToolCallPart stand-in: duck-typed on `.tool_name` / `.args`, which + is exactly what the adapter inspects to find proposed tool calls.""" + + def __init__(self, tool_name, args=None): + self.tool_name = tool_name + self.args = args + + +class _TextPart: + """A non-tool response part (final text): has no `.tool_name`, so it is never + gated.""" + + def __init__(self, content): + self.content = content + + +class _FakeGate: + """Records gate calls; can be set to deny like a human denying a tool.""" + + def __init__(self, deny=False): + self.deny = deny + self.calls = [] + + def require(self, tool, side_effect="", arguments=None, timeout=None): + self.calls.append({"tool": tool, "side_effect": side_effect, + "arguments": arguments, "timeout": timeout}) + if self.deny: + raise ApprovalDenied(tool, "no") + + +def _run(coro): + return asyncio.run(coro) + + +class PydanticAIAdapterTest(unittest.TestCase): + def test_module_imports_without_pydantic_ai(self): + # The lazy base-class fallback must let the SDK import the adapter even with + # no pydantic-ai installed (it is not a dependency). + self.assertTrue(callable(GovernedModel)) + self.assertTrue(callable(govern)) + + def test_request_ticks_one_step_per_model_request(self): + run = FakeRun() + inner = _FakeModel(_Resp([_TextPart("done")])) + gm = govern(inner, run) + _run(gm.request([], None, None)) + self.assertEqual(run.steps, 1) # one governed step per model request + self.assertEqual(inner.requests, 1) # and it delegated to the real model + + def test_request_enforces_loop_budget(self): + run = FakeRun(loop_budget=2) + inner = _FakeModel(_Resp([_TextPart("ok")])) + gm = govern(inner, run) + _run(gm.request([], None, None)) # step 1 + _run(gm.request([], None, None)) # step 2 + with self.assertRaises(BudgetExceeded) as cm: + _run(gm.request([], None, None)) # step 3 -> over budget + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + + def test_budget_halt_is_not_swallowed_and_skips_the_request(self): + # The propagation contract: the wrapper must let BudgetExceeded escape (so + # PydanticAI, which only retries its own ModelRetry, halts the agent), and it + # must tick the step BEFORE delegating, so an over-budget run never spends on + # the underlying model request. + run = FakeRun(loop_budget=0) # over budget from the very first step + inner = _FakeModel(_Resp([_TextPart("ok")])) + gm = govern(inner, run) + with self.assertRaises(BudgetExceeded): + _run(gm.request([], None, None)) + self.assertEqual(inner.requests, 0) # the model was never called + + def test_tool_gating_off_by_default(self): + # Default: gate_tools is False, so proposed tool calls are never gated. + run = FakeRun() + inner = _FakeModel(_Resp([_ToolCallPart("deploy", {"x": 1})])) + gm = govern(inner, run) + gate = _FakeGate() + gm._gate = gate + _run(gm.request([], None, None)) + self.assertEqual(gate.calls, []) # gate_tools defaults off + self.assertEqual(run.steps, 1) # the step still ticks + + def test_tool_gating_on_passes_tool_name_and_side_effect(self): + run = FakeRun() + inner = _FakeModel(_Resp([_ToolCallPart("deploy", {"env": "prod"})])) + gm = govern(inner, run, gate_tools=True, tool_side_effect="exec", timeout=5) + gate = _FakeGate() + gm._gate = gate + _run(gm.request([], None, None)) + self.assertEqual(len(gate.calls), 1) + self.assertEqual(gate.calls[0]["tool"], "deploy") + self.assertEqual(gate.calls[0]["side_effect"], "exec") + self.assertEqual(gate.calls[0]["timeout"], 5) + self.assertEqual(run.steps, 1) + + def test_tool_gating_gates_every_proposed_tool_call(self): + # A single model response may propose several tool calls; each is gated. + run = FakeRun() + inner = _FakeModel(_Resp([ + _ToolCallPart("search", {"q": "a"}), + _ToolCallPart("write", {"path": "/x"}), + ])) + gm = govern(inner, run, gate_tools=True) + gate = _FakeGate() + gm._gate = gate + _run(gm.request([], None, None)) + self.assertEqual([c["tool"] for c in gate.calls], ["search", "write"]) + + def test_tool_gating_does_not_gate_a_final_text_answer(self): + # A text-only response proposes no tool calls: gate_tools=True must not gate + # it, even though the request still ticks a governed step. + run = FakeRun() + inner = _FakeModel(_Resp([_TextPart("final answer")])) + gm = govern(inner, run, gate_tools=True) + gate = _FakeGate() + gm._gate = gate + _run(gm.request([], None, None)) + self.assertEqual(gate.calls, []) + self.assertEqual(run.steps, 1) + + def test_tool_gating_denied_raises_and_halts(self): + # A denied tool must raise ApprovalDenied, which propagates out of the run and + # halts the agent before the tool executes. + run = FakeRun() + inner = _FakeModel(_Resp([_ToolCallPart("deploy", {})])) + gm = govern(inner, run, gate_tools=True) + gm._gate = _FakeGate(deny=True) + with self.assertRaises(ApprovalDenied): + _run(gm.request([], None, None)) + + def test_tool_calls_helper_is_duck_typed(self): + # The adapter duck-types: a part with `.tool_name` is a tool call; a text part + # (no `.tool_name`) is not. Works without importing pydantic-ai. + parts = [_ToolCallPart("x"), _TextPart("y"), _ToolCallPart("z")] + found = [p.tool_name for p in _tool_calls(_Resp(parts))] + self.assertEqual(found, ["x", "z"]) + self.assertEqual(list(_tool_calls(_Resp([]))), []) + self.assertEqual(list(_tool_calls(_Resp(None))), []) + + @unittest.skipUnless(_has_pydantic_ai(), "pydantic-ai not installed") + def test_pydantic_ai_integration_halts_runaway_agent(self): + # The real proof: a PydanticAI agent that would loop forever is hard-stopped + # at its loop budget. We drive a real Agent with a FunctionModel that always + # proposes a tool call (so the agent never produces a final answer) and assert + # the governor's BudgetExceeded propagates out of agent.run_sync — PydanticAI + # only retries its own ModelRetry, so a plain exception halts the agent. + from pydantic_ai import Agent + from pydantic_ai.messages import ModelResponse, ToolCallPart + from pydantic_ai.models.function import FunctionModel + + def always_call_tool(messages, info): + # Never emit a text part -> the run never reaches a final output -> the + # agent loops on tool calls forever without the governor. + return ModelResponse(parts=[ToolCallPart(tool_name="spin", args={})]) + + run = FakeRun(loop_budget=3) + agent = Agent(govern(FunctionModel(always_call_tool), run)) + + @agent.tool_plain + def spin() -> str: + return "still going" + + with self.assertRaises(BudgetExceeded) as cm: + agent.run_sync("go") + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + # The governor capped the loop at budget + 1 (the request that tripped it). + self.assertEqual(run.steps, 4) + + @unittest.skipUnless(_has_pydantic_ai(), "pydantic-ai not installed") + def test_pydantic_ai_integration_tool_gating_denied_halts_agent(self): + # gate_tools=True with a denial: a real agent's proposed tool call is routed + # through the gate, the denial raises ApprovalDenied, and it propagates out of + # agent.run_sync — the tool never executes. + from pydantic_ai import Agent + from pydantic_ai.messages import ModelResponse, ToolCallPart + from pydantic_ai.models.function import FunctionModel + + executed = {"deploy": False} + + def call_deploy(messages, info): + return ModelResponse(parts=[ToolCallPart(tool_name="deploy", args={})]) + + run = FakeRun(loop_budget=10) + gm = govern(FunctionModel(call_deploy), run, gate_tools=True) + gm._gate = _FakeGate(deny=True) + agent = Agent(gm) + + @agent.tool_plain + def deploy() -> str: + executed["deploy"] = True + return "deployed" + + with self.assertRaises(ApprovalDenied): + agent.run_sync("go") + self.assertFalse(executed["deploy"]) # the denied tool never ran + + +if __name__ == "__main__": + unittest.main()