Testing LLM Agents: PydanticAI DI Strategies

Introduction: The Challenge of Testing AI

Building intelligent agents with Large Language Models (LLMs) is exciting, but it introduces significant challenges in software engineering. Unlike traditional deterministic code, LLM agents rely on external APIs, stochastic outputs, and complex state management. This makes testing difficult. How do you verify logic when the "brain" of your application is a non-deterministic API call? The solution lies in strict software design patterns, specifically Dependency Injection (DI). PydanticAI, a modern Python framework for building LLM applications, excels at this by leveraging type hints and pydantic models to manage dependencies cleanly. By isolating the LLM client from your business logic, you can write fast, reliable unit tests that don't cost you money in API tokens or suffer from flaky network issues.

Why Dependency Injection Matters for Agents

In a typical agent workflow, you might have an agent that searches the web, processes the results, and then generates a summary. If the web search tool is hard-coded into the agent, you cannot test the summarization logic without actually hitting a search engine. Dependency Injection allows you to swap out these dependencies. In production, you inject a real `WebSearchTool`. In tests, you inject a `MockWebSearchTool` that returns predefined data. This pattern ensures that your unit tests focus strictly on the logic of your agent, not the reliability of external services.

Setting Up a Testable Agent

First, let's define a simple agent that uses a dependency. We will use PydanticAI's system for managing tool dependencies.

from pydantic_ai import Agent
from pydantic_ai.models import Model
from pydantic_ai.toolsets import Toolset
from pydantic import BaseModel

class Fact(BaseModel):
    fact: str
    source: str

# Define a custom Toolset to manage dependencies
class FactToolset(Toolset):
    model: Model
    facts: list[Fact]

    def __init__(self, model: Model, facts: list[Fact]):
        self.model = model
        self.facts = facts

    @property
    def get_facts(self):
        """Return a list of known facts."""
        return self.facts

# Define the Agent
agent = Agent[FactToolset](
    model='openai:gpt-4o',
    system_prompt='You are a fact-checker. Use the provided facts to answer questions.',
)

async def main():
    # In production, you would inject real facts
    toolset = FactToolset(model=agent.model, facts=[Fact(fact="Water is wet", source="Physics")])
    
    try:
        result = await agent.run('Is water wet?', toolset=toolset)
        print(result.output)
    finally:
        await toolset.aclose()

Unit Testing with Mocks

Now, we can write a unit test. The key is to mock the `Model` itself so that we don't make actual API calls. PydanticAI provides a `MockModel` for this purpose.

import pytest
from pydantic_ai import Agent
from pydantic_ai.models import MockModel
from pydantic_ai.toolsets import Toolset
from pydantic import BaseModel
from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart

class Fact(BaseModel):
    fact: str
    source: str

class FactToolset(Toolset):
    def __init__(self, facts: list[Fact]):
        self.facts = facts

    def get_facts(self):
        return self.facts

# Create the agent with a placeholder model. The real model will be injected later.
agent = Agent[FactToolset](
    model='mock', 
    system_prompt='You are a fact-checker.'
)

@pytest.mark.anyio
async def test_agent_facts():
    # 1. Define the expected response from the LLM
    # The LLM should call the 'get_facts' tool
    mock_response = ModelResponse(
        parts=[
            ToolCallPart(tool_name='get_facts', args={}),
            TextPart(content='Based on the facts, water is wet.')
        ]
    )
    
    # 2. Create a MockModel that returns our predefined response
    mock_model = MockModel(messages=[mock_response])
    
    # 3. Create the toolset with test data
    test_facts = [Fact(fact="Water is wet", source="Test Source")]
    toolset = FactToolset(facts=test_facts)
    
    # 4. Run the agent with the mock model and toolset
    # Note: We inject the mock model into the agent context
    async with agent.run_with_tools('Is water wet?', toolset=toolset, model=mock_model) as result:
        assert result.output == 'Based on the facts, water is wet.'

Best Practices for Robust Tests

1. Isolate State: Ensure that each test creates a fresh instance of the toolset. Avoid shared state between tests to prevent contamination. 2. Validate Tools: In your tests, verify that the correct tools were called. You can inspect the `result.all_messages` to assert that `ToolCallPart` objects with specific names exist. 3. Use Structured Outputs: PydanticAI supports structured outputs (pydantic models). Test that the LLM's response adheres to your schema. This is crucial for downstream data processing.

Conclusion

Implementing LLM agents requires a shift in thinking from "scripting" to "engineering." By leveraging PydanticAI's dependency injection capabilities, you can build agents that are not only functional but also testable and maintainable. The combination of `MockModel` for LLM responses and custom `Toolsets` for data dependencies allows you to decouple your logic from the non-deterministic nature of AI. Start by isolating your tools, mock your models, and you will build a foundation for reliable AI applications that can scale with confidence.