# PydanticAI-Native Integration Plan

## Goal
Make rgltr a drop-in governance layer for PydanticAI such that existing apps using any PydanticAI tool registration pattern can enable governance with minimal changes.

**Target experience:**
- Add `RegulatedAgent` (or a helper) and everything works, regardless of whether tools were registered via `@agent.tool`, `@agent.tool_plain`, `Agent(..., tools=[...])`, `FunctionToolset`, dynamic toolsets, or MCP servers.

## Scope
- Govern all **function tools** registered by the agent (including `@agent.tool` / `@agent.tool_plain` and `tools=[...]`).
- Govern **toolsets** passed to the agent (static and dynamic).
- Govern **MCP servers** via `RegulatedMCPAdapter` (already done).
- Add a config flag for builtin tools but default to **not** governing them.

## Reference Sources
- PydanticAI repo (examples + docs): `../pydantic-ai`
  - Examples: `../pydantic-ai/examples/pydantic_ai_examples/`
  - Docs: `../pydantic-ai/docs/`

## Common PydanticAI Usage Patterns (from examples/docs)

1. **`@agent.tool` with `RunContext`**
   - Example: `weather_agent`, `bank_support`, `roulette_wheel`, `data_analyst`.
   - Governance should apply automatically without refactoring.

2. **`@agent.tool_plain` without `RunContext`**
   - Example: AG-UI samples (`agentic_chat`, `shared_state`, `predictive_state_updates`).
   - Governance should apply automatically even though the tool does not accept `RunContext`.

3. **`Agent(..., tools=[...])`**
   - Example: `slack_lead_qualifier` uses `tools=[duckduckgo_search_tool()]`.
   - Governance must apply to any function or `Tool(...)` passed here.

4. **`FunctionToolset` and `.tool(...)` decorator**
   - Documented in `toolsets.md` and `tools-advanced.md`.
   - Governance should work when a toolset is provided at init or per-run.

5. **Dynamic toolsets** (`@agent.toolset` or `toolsets=[callable]`)
   - Documented in `toolsets.md`.
   - Governance should wrap the toolset produced at runtime.

6. **Streaming (`run_stream`)**
   - Example: `chat_app`, `roulette_wheel`.
   - Governance should work identically and preserve tag restoration across turns.

7. **Message history**
   - Example: `chat_app` and `flight_booking` use `message_history`.
   - Tag restoration should re-activate tags based on past tool calls.

## Design Overview

### 1) Wrap the Agent’s Internal FunctionToolset (default on)
PydanticAI stores tools registered via `@agent.tool`, `@agent.tool_plain`, and `Agent(..., tools=[...])` inside its internal `FunctionToolset` (private `_function_toolset`).

**Design:**
- Add `wrap_agent_toolset=True` (default) on `RegulatedAgent`.
- On run, wrap the internal toolset with `RegulatedToolset(toolset=agent._function_toolset)`.
- Use `Agent.override(toolsets=...)` to inject the regulated toolset without mutating internal state.

**Alternative / opt-out:**
- Allow `wrap_agent_toolset=False` for strict users who want to control toolsets manually.

**Benefits:**
- Works with all existing registrations, including `@agent.tool` and `tools=[...]`.
- Requires no changes to user tool definitions.

### 2) Wrap Additional Toolsets Provided by the App
PydanticAI accepts `toolsets=[...]` at agent creation or per-run, including `FunctionToolset`, `FilteredToolset`, `PreparedToolset`, or dynamic toolset functions.

**Design:**
- When `RegulatedAgent.run*` sees `toolsets=...`, map each toolset through a wrapper:
  - `FunctionToolset` → `RegulatedToolset(toolset=...)`
  - `WrapperToolset` or composites (Filtered/Prepared/Prefixed/Renamed) → wrap the leaf `FunctionToolset` by visiting toolsets and replacing leaves (use `visit_and_replace` if available).
  - Dynamic toolset function (callable) → wrap returned toolset at runtime.
- Preserve non-function toolsets (e.g., MCP toolsets) unless they are wrapped by `RegulatedMCPAdapter`.

**Benefits:**
- Keeps toolset-based patterns working with governance.

### 3) Builtin Tools (flagged, default off)
PydanticAI builtin tools include WebSearchTool, CodeExecutionTool, MCPServerTool, etc.

**Design:**
- Add `govern_builtin_tools: bool = False` (global + per-agent config).
- When enabled, wrap builtin toolsets with a thin adapter that applies rgltr visibility/policy rules.
- Leave off by default to avoid surprising behavior; document that builtin tools are not governed unless enabled.

### 4) Add a Helper for Quick Adoption
Provide a single helper that makes adoption trivial:

```python
regulated = rgltr.pydantic_ai.regulate(agent, agent_id="my-agent")
result = await regulated.run(...)
```

This helper should:
- Wrap internal function toolset
- Wrap provided toolsets
- Keep coverage reporting and tag restoration

### 5) Dynamic Toolsets (@agent.toolset)
PydanticAI allows toolset factories via `@agent.toolset` or `toolsets=[callable]`.

**Design:**
- Detect callables in `toolsets` and wrap the result on each run.
- If the callable returns a `FunctionToolset`, wrap it with `RegulatedToolset`.
- If it returns a composite toolset, visit and wrap leaf function toolsets.

### 6) Maintain Consistent Tag Restoration
Tag restoration from `message_history` should work for tools created through any registration path.

**Design:**
- During `RegulatedAgent.run*`, build a lookup that includes:
  - Rules from wrapped internal toolset
  - Rules from wrapped toolsets (including dynamic toolsets, after they resolve)
  - MCP tool rules (via `RegulatedMCPAdapter.rules_for_tool`)

### 7) Document Supported Patterns
Update README to state:
- `@agent.tool` / `@agent.tool_plain` are fully supported and governed when using `RegulatedAgent`.
- `tools=[...]` and `FunctionToolset` are governed automatically.
- Dynamic toolsets are supported.
- Builtin tools are not governed unless `govern_builtin_tools=True`.

## Test Plan (map to examples)

1. **`@agent.tool` + `RunContext`**
   - Register tools via `@agent.tool`, ensure boundary + tag rules apply without explicit `RegulatedToolset`.

2. **`@agent.tool_plain`**
   - Register tool via `@agent.tool_plain`, ensure it is governed and visible/hidden per tags.

3. **`Agent(..., tools=[...])`**
   - Tools list includes plain function and `Tool(...)`. Ensure governance applies.

4. **`FunctionToolset` via `.tool(...)`**
   - Build a toolset and pass it into `Agent(..., toolsets=[...])` or per-run override. Ensure governance applies.

5. **Dynamic toolset (`@agent.toolset` / callable)**
   - Provide a factory returning a `FunctionToolset`. Ensure governance applies.

6. **Streaming + message history**
   - Run stream, activate tag, ensure next run with `message_history` hides tools.

7. **MCP toolsets**
   - Ensure MCP tools still governed and filtering works.

8. **Builtin tools (optional)**
   - When `govern_builtin_tools=True`, ensure builtin tool invocations are blocked/allowed based on tags/policies.

## Implementation Steps
1. **RegulatedAgent wrapper**
   - Add `wrap_agent_toolset` flag (default `True`).
   - Add `govern_builtin_tools` flag (default `False`).
   - Build `_get_wrapped_toolsets(run_toolsets)` that:
     - wraps internal function toolset when `wrap_agent_toolset=True`
     - wraps provided toolsets (including dynamic toolset callables)
     - optionally wraps builtin toolsets when `govern_builtin_tools=True`
   - Use `Agent.override(toolsets=wrapped)` to run.

2. **Toolset wrapping utility**
   - Add `wrap_toolset(toolset)` helper:
     - If `FunctionToolset`: `RegulatedToolset(toolset=toolset)`
     - If composite toolset: `toolset.visit_and_replace` to wrap leaf `FunctionToolset`
     - If MCP adapter already regulated: leave it
   - Add `wrap_toolsets([...])` for lists

3. **Dynamic toolsets**
   - If toolset is callable: call inside run and wrap the returned toolset
   - Cache per-run to avoid double-wrapping

4. **Tests**
   - Add tests listed above

## Risks / Edge Cases
- PydanticAI private attributes (`_function_toolset`) may change; prefer `Agent.override(toolsets=...)` to avoid mutation.
- Composite toolset wrapping must preserve identity and ordering.
- Durable execution (Temporal/Prefect) may rely on stable toolset IDs; ensure wrappers preserve or forward IDs.

## Deliverables
- Updated `RegulatedAgent` to auto-govern all function tools.
- Wrapping helpers for toolsets.
- Updated README guidance.
- Tests covering agent-native tool registration paths.
