Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ jobs:
run: uv run pyright

- name: MyPy
run: uv run mypy src/pydantic_ai_middleware tests
run: uv run mypy src/pydantic_ai_shields tests

test:
name: Test Python ${{ matrix.python-version }}
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/publish.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ jobs:
runs-on: ubuntu-latest
environment:
name: pypi
url: https://pypi.org/project/pydantic-ai-middleware/
url: https://pypi.org/project/pydantic-ai-shields/
permissions:
id-token: write

Expand Down
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
notes/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
Expand Down
81 changes: 24 additions & 57 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,83 +4,50 @@ Instructions for AI coding assistants working on this repository.

## Project Overview

**pydantic-ai-middleware** is a simple middleware library for [pydantic-ai](https://ai.pydantic.dev/) agents. It provides clean before/after hooks at every lifecycle stage without imposing a guardrails structure - you decide what to do (logging, guardrails, metrics, transformations).
**pydantic-ai-shields** provides guardrail capabilities for [pydantic-ai](https://ai.pydantic.dev/) agents. Built on pydantic-ai's native capabilities API (v1.71+). No middleware wrappers — pure capabilities.

## Quick Reference

| Task | Command |
|------|---------|
| Install | `make install` |
| Test | `make test` |
| Test + Coverage | `uv run coverage run -m pytest && uv run coverage report` |
| Lint | `uv run ruff check .` |
| Format | `uv run ruff format .` |
| Typecheck | `uv run pyright` and `uv run mypy .` |
| All checks | `make all` |
| Build docs | `make docs-serve` |
| Test | `uv run pytest tests/ -v` |
| Test + Coverage | `uv run coverage run -m pytest tests/ && uv run coverage report --fail-under=100` |
| Lint | `uv run ruff check src/ tests/` |
| Typecheck | `uv run pyright src/` |

## Architecture

```
pydantic_ai_middleware/
├── base.py # AgentMiddleware - base class with lifecycle hooks
├── agent.py # MiddlewareAgent - wraps agents with middleware
├── toolset.py # MiddlewareToolset - wraps toolsets for tool call interception
├── decorators.py # @before_run, @after_run, etc. - function decorators
├── exceptions.py # InputBlocked, ToolBlocked, OutputBlocked
└── __init__.py # Public API exports
src/pydantic_ai_shields/
__init__.py — Package exports
guardrails.py — 5 capability implementations + exceptions + CostInfo
tests/
test_guardrails.py — All tests
```

### Key Design: Middleware Chain
### Capabilities

Middleware executes in order for `before_*` hooks and reverse order for `after_*` hooks:

```python
agent = MiddlewareAgent(
agent=base_agent,
middleware=[mw1, mw2, mw3],
)
# before_run: mw1 -> mw2 -> mw3 -> [Agent]
# after_run: [Agent] -> mw3 -> mw2 -> mw1
```

### Lifecycle Hooks

| Hook | When Called | Can Modify |
|------|-------------|------------|
| `before_run` | Before agent starts | Prompt |
| `after_run` | After agent finishes | Output |
| `before_model_request` | Before each model call | Messages |
| `before_tool_call` | Before tool execution | Tool arguments |
| `after_tool_call` | After tool execution | Tool result |
| `on_error` | When error occurs | Exception |
| Capability | Hooks Used | Purpose |
|------------|-----------|---------|
| `CostTracking` | `before_run`, `after_run` | Token/USD tracking, budget enforcement |
| `ToolGuard` | `prepare_tools`, `before_tool_execute` | Block tools, require approval |
| `InputGuard` | `before_run` | Validate user input |
| `OutputGuard` | `after_run` | Validate model output |
| `AsyncGuardrail` | `wrap_run` | Concurrent guardrail + LLM |

## Code Standards

- **Coverage**: 100% required - check with `uv run coverage report`
- **Types**: Pyright and mypy - all functions need type annotations
- **Style**: ruff handles formatting and linting
- **Coverage**: 100% required
- **Types**: Pyright strict on src/
- **Style**: ruff for formatting and linting

## Testing

Tests are in `tests/` directory. Use pytest-asyncio for async tests:

```python
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai_middleware import MiddlewareAgent, AgentMiddleware
from pydantic_ai_shields import CostTracking, ToolGuard

async def test_middleware():
model = TestModel()
model.custom_output_text = "test"
agent = Agent(model, output_type=str)
middleware_agent = MiddlewareAgent(agent, middleware=[...])
result = await middleware_agent.run("test")
agent = Agent(TestModel(), capabilities=[CostTracking()])
result = await agent.run("test")
```

## When Modifying

1. Run tests after changes: `make test`
2. Check coverage stays at 100%
3. Run `uv run pyright` for type errors
4. Format with `uv run ruff format .`
Loading
Loading