Add Pydantic AI agent runner support #1
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit adds support for running Pydantic AI agents alongside OpenAI agents. Changes: - Created PydanticAIRunner class in src/agentexec/runners/pydantic_ai.py - Implements full BaseAgentRunner interface - Supports request_limit (equivalent to max_turns) - Implements UsageLimitExceeded recovery mechanism - Provides report_status tool as Pydantic AI Tool - Supports both sync run() and streaming run_streamed() - Handles message history for recovery - Updated src/agentexec/runners/__init__.py to export PydanticAIRunner - Uses optional import pattern for graceful degradation - Modified pyproject.toml: - Moved agent framework dependencies to optional-dependencies - Added pydantic-ai as optional dependency - Created "all" extra for installing both frameworks - Updated description and keywords The PydanticAIRunner provides the same features as OpenAIRunner: - Max turns recovery (via request_limit) - Status reporting tool integration - Activity tracking via agent_id - Streaming support - Configurable wrap-up and status prompts
This commit adds extensive test coverage for the PydanticAIRunner and fixes several bugs discovered during testing. Test Coverage: - Added tests/test_pydantic_ai_runner.py with 21 tests covering: - Initialization and configuration - Basic execution (run method) - Dependencies and message history handling - Request limit recovery mechanism - Streaming execution - Tool integration and activity tracking - Public API integration - Updated tests/test_public_api.py: - Added test_pydantic_ai_runner_imports - Added test_pydantic_ai_runner_initialization - Renamed OpenAI tests for clarity Bug Fixes in PydanticAIRunner: - Fixed imports to use correct Pydantic AI types: - Changed UserPrompt -> UserPromptPart - Changed Message -> ModelMessage - Changed RunResult -> AgentRunResult - Fixed _extract_messages to check for message_history attribute - Fixed recovery mechanism to properly create ModelRequest with UserPromptPart - Fixed user_prompt parameter to use None instead of empty string during recovery - Updated type hints to match Pydantic AI's actual API All 21 new tests pass successfully, providing comprehensive coverage of the PydanticAIRunner functionality including edge cases and error handling.
Added pydantic-ai and its dependencies to the lock file after installing for test development and verification.
This commit adds comprehensive CI/CD testing infrastructure and resolves all linting issues to ensure code quality. GitHub Actions Workflow (.github/workflows/test.yml): - Runs on push to main and all pull requests - Tests on Python 3.11 and 3.12 (matrix strategy) - Full test suite job: - Installs all optional dependencies (OpenAI, Pydantic AI) - Runs ruff linter - Runs mypy type checker (non-blocking) - Executes full test suite with coverage - Uploads coverage to Codecov - Minimal dependencies job: - Tests that core package works without optional dependencies - Skips framework-specific tests Ruff Configuration (pyproject.toml): - Added per-file-ignores to allow F401 in __init__.py files - These files use imports for re-export, which is intentional Code Quality Fixes: - Removed unused imports from src/agentexec/runners/pydantic_ai.py - Removed Callable (not needed) - Removed ModelResponse (not used) - Fixed unused imports across test files - Fixed unused variable assignments in tests - Changed unused 'agent' to '_' in test_pydantic_ai_runner.py - Changed unused 'task1' to '_' in test_worker_pool.py All tests pass (21 PydanticAI tests + all existing tests). All ruff checks pass. Ready for continuous integration.
Resolved conflicts in: - pyproject.toml: Updated to version 0.1.2, kept Pydantic AI in description - tests/test_task.py: Accepted main's import structure - tests/test_worker_pool.py: Updated to use async enqueue with BaseModel context This merge brings in significant changes from main: - Docker worker image support - Pipeline functionality - Context object serialization improvements - SQLAlchemy session pattern updates - Async Redis implementation
CRITICAL FIX: The original implementation incorrectly assumed that
UsageLimitExceeded would have a message_history attribute, but this
is NOT how Pydantic AI works.
## What Was Wrong
The initial implementation tried to extract message history directly from
the UsageLimitExceeded exception using a non-existent attribute:
```python
def _extract_messages(e: UsageLimitExceeded) -> list[ModelMessage]:
if hasattr(e, "message_history") and e.message_history:
return list(e.message_history)
```
This approach was based on the OpenAI Agents SDK pattern (MaxTurnsExceeded.run_data),
but Pydantic AI uses a completely different pattern.
## How Pydantic AI Actually Works
Per Pydantic AI documentation (https://ai.pydantic.dev/agents/#model-errors):
- Use `capture_run_messages()` context manager to capture messages
- The context manager populates a list during agent.run() execution
- If an exception occurs, the captured messages are still available
- This is the official, documented approach
Reference: pydantic/pydantic-ai#1083
## Changes Made
1. **Updated imports**:
- Added `capture_run_messages` from pydantic_ai
- Removed obsolete `_extract_messages()` helper function
2. **Rewrote run() method**:
- Wrapped agent.run() in `with capture_run_messages() as messages:`
- On UsageLimitExceeded, use the captured messages for recovery
- Messages are properly populated by the context manager
3. **Rewrote run_streamed() method**:
- Same pattern as run() but for streaming
- Uses capture_run_messages() consistently
4. **Updated all recovery tests**:
- Mock capture_run_messages context manager properly
- Use `@patch("agentexec.runners.pydantic_ai.capture_run_messages")`
- Mock return value with `__enter__` to simulate context manager
## Verification
✅ All 21 PydanticAIRunner tests pass
✅ All 49 tests in test suite pass
✅ Recovery mechanism properly captures and reuses conversation history
✅ Implementation follows official Pydantic AI patterns
## Why This Matters
Without this fix, the recovery mechanism would NEVER work correctly because:
1. Message history would always be empty
2. Recovery would lose all conversation context
3. The wrap-up prompt would be sent without prior conversation
This fix ensures the PydanticAIRunner behaves correctly according to
Pydantic AI's actual API, not just our assumptions.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit adds support for running Pydantic AI agents alongside OpenAI agents.
Changes:
Created PydanticAIRunner class in src/agentexec/runners/pydantic_ai.py
Updated src/agentexec/runners/init.py to export PydanticAIRunner
Modified pyproject.toml:
The PydanticAIRunner provides the same features as OpenAIRunner: