Routing workflow with Pydantic AI

Author

Affiliation

Dylan Castillo

Iwana Labs

I’m trying to get more familiar with Pydantic AI, so I’ve been re-implementing typical patterns for building agentic systems.

In this post, I’ll build a routing workflow. I won’t cover the basics of agentic workflows, so if you’re not familiar with the concept, I recommend you to read this post first.

I’ve also written other TILs about Pydantic AI:

• Prompt chaining
• ReAct agent
• Evaluator-optimizer
• Parallelization and Orchestrator-workers

You can download this notebook here.

What is router?

Routing is a workflow pattern that takes the input, classifies it and then sends it to the right place for the best handling. This process can be managed by an LLM or a traditional classification model. It makes sense to use when a system needs to apply different logic to different types of queries.

It looks like this:

flowchart LR 
    In([In]) --> Router["LLM Call Router"]

    Router -->|Route 1| LLM1["LLM Call 1"]
    Router -->|Route 2| LLM2["LLM Call 2"]
    Router -->|Route 3| LLM3["LLM Call 3"]

    LLM1 --> Out([Out])
    LLM2 --> Out
    LLM3 --> Out

Examples:

• Classify complexity of question and adjust model depending on it
• Classify type of query and use specialized tools (e.g., indexes, prompts)

Let’s see how this looks like in code.

Setup

I will implement a workflow that will take a query from a user and will route it to the appropriate agent.

There will be three agents in the workflow:

• Agent TOC: Generate a table of contents for the article
• Agent Writer: Generate the content of the article
• Agent Editor: Update the content of the article if it’s too long

Because Pydantic AI uses asyncio under the hood, you need to enable nest_asyncio to use it in a notebook:

import nest_asyncio

nest_asyncio.apply()

Then, you need to import the required libraries. Logfire is part of the Pydantic ecosystem, so I thought it’d be good to use it for observability.

import os
from typing import Literal

import logfire
import requests
from dotenv import load_dotenv
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext

load_dotenv()

True

PydanticAI is compatible with OpenTelemetry (OTel). It’s straightforward to use it with Logfire or with any other OTel-compatible observability tool (e.g., Langfuse).

To enable tracking, create a project in Logfire, generate a Write token and add it to the .env file. Then, you just need to run:

logfire.configure(
    token=os.getenv("LOGFIRE_TOKEN"),
)
logfire.instrument_pydantic_ai()

The first time you run this, it will ask you to create a project in Logfire. From it, it will generate a logfire_credentials.json file in your working directory. In following runs, it will automatically use the credentials from the file.

Prompt chaining workflow

As mentioned before, the workflow will be composed of three agents. So I created three Agent instances. Each one takes care of one of the tasks

Here’s the code:

class RouterOutput(BaseModel):
    category: Literal["write_article", "generate_table_of_contents", "review_article"]


router_agent = Agent(
    "openai:gpt-4.1-mini",
    system_prompt=(
        "You are a helpful assistant. You will classify the message into one of the following categories: 'write_article', 'generate_table_of_contents', 'review_article'."
    ),
    output_type=RouterOutput,
)

agent_writer = Agent(
    "openai:gpt-4.1-mini",
    system_prompt=(
        "You are a writer. You will write an article about the topic provided."
    ),
)

agent_toc = Agent(
    "openai:gpt-4.1-mini",
    system_prompt=(
        "You are an expert writer specialized in SEO. Provided with a topic, you will generate the table of contents for a short article."
    ),
)

agent_reviewer = Agent(
    "openai:gpt-4.1-mini",
    system_prompt=(
        "You are a writer. You will review the article for the topic provided."
    ),
)

Logfire project URL: https://logfire-us.pydantic.dev/dylanjcastillo/blog

@logfire.instrument("Run workflow")
def run_workflow(topic: str) -> str:
    router_output = router_agent.run_sync(
        f"Classify the message: {topic}"
    )
    category = router_output.output.category 
    if category == "write_article":
        return agent_writer.run_sync(f"Write an article about {topic}").output
    elif category == "generate_table_of_contents":
        return agent_toc.run_sync(f"Generate the table of contents of an article about {topic}").output
    else:
        return agent_reviewer.run_sync(f"Review the article for the topic {topic}").output

You can run the workflow and it will route your message and use the appropriate agent. For example, try to generate a table of contents for an article about AI:

toc = run_workflow("Generate a table of contents for an article about AI")

20:08:05.547 Run workflow
20:08:05.548   router_agent run
20:08:05.549     chat gpt-4.1-mini
20:08:06.810   agent_toc run
20:08:06.811     chat gpt-4.1-mini

print(toc)

Table of Contents

1. Introduction to Artificial Intelligence  
2. History and Evolution of AI  
3. Types of Artificial Intelligence  
4. Key Technologies Behind AI  
5. Applications of AI in Various Industries  
6. Benefits and Challenges of AI  
7. Future Trends in Artificial Intelligence  
8. Ethical Considerations in AI Development  
9. Conclusion

Or, ask the workflow to review a social media post.

review = run_workflow("Review this post: 'There are times where there's no time, so you don't have time to write an article about it.'")

20:08:08.917 Run workflow
20:08:08.918   router_agent run
20:08:08.919     chat gpt-4.1-mini
20:08:09.691   agent_reviewer run
20:08:09.692     chat gpt-4.1-mini

print(review)

The post "There are times where there's no time, so you don't have time to write an article about it." offers a succinct reflection on the challenges of time constraints, especially in tasks like writing. Its brevity captures the irony of not having enough time to address a situation—in this case, the lack of time itself. The message resonates with anyone who has felt overwhelmed by deadlines or competing priorities.

However, as a piece intended for a broader audience or a formal article, it could benefit from expansion. Elaborating on scenarios where time scarcity impacts productivity, or providing strategies for managing pressing tasks despite limited time, would add depth and practical value. Additionally, refining the sentence for clarity and flow could enhance its impact—for example: "Sometimes, we're so pressed for time that we can't even write about the very pressure we're under."

In summary, the post effectively conveys a common frustration with time limitations in a clever and relatable way but serves better as a starting point for a more detailed discussion rather than a standalone article.

That’s all!

You can access this notebook here.

If you have any questions or feedback, please let me know in the comments below.