PydanticAI Tools Integration

Introduction

PydanticAI is a newer framework that brings powerful features for working with LLMs. Although it doesn’t yet have a collection of pre-built tools like other frameworks, it offers useful capabilities such as dependency injection. This feature allows you to inject a “Context” into tools, which can help pass parameters or manage state without relying on LLMs. Though it’s still evolving, you can easily integrate PydanticAI tools into AG2 to boost agent capabilities, particularly for tasks that involve structured data and context-driven logic.

Installation

To get PydanticAI tools working with AG2, install the necessary dependencies:

pip install ag2[openai,interop-pydantic-ai]

Imports

Import necessary modules and tools.

• BaseModel: Used to define data structures for tool inputs and outputs.
• RunContext: Provides context during the execution of tools.
• PydanticAITool: Represents a tool in the PydanticAI framework.
• AssistantAgent and UserProxyAgent: Agents that facilitate communication in the AG2 framework.
• Interoperability: This module acts as a bridge, making it easier to integrate PydanticAI tools with AG2’s architecture.

import os
from typing import Optional

from pydantic import BaseModel
from pydantic_ai import RunContext
from pydantic_ai.tools import Tool as PydanticAITool

from autogen import AssistantAgent, UserProxyAgent
from autogen.interop import Interoperability

Agent Configuration

Configure the agents for the interaction.

• config_list defines the LLM configurations, including the model and API key.
• UserProxyAgent simulates user inputs without requiring actual human interaction (set to NEVER).
• AssistantAgent represents the AI agent, configured with the LLM settings.

config_list = [{"api_type": "openai", "model": "gpt-4o", "api_key": os.environ["OPENAI_API_KEY"]}]
user_proxy = UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
)

chatbot = AssistantAgent(
    name="chatbot",
    llm_config={"config_list": config_list},
)

Tool Integration

To integrate a PydanticAI tool into AG2:

• First, define a Player model using BaseModel to structure the input data.
• Use RunContext to inject dependencies (like the Player instance) securely into the tool.
• The get_player function defines the tool’s functionality, retrieving injected data through ctx.deps.
• Then, convert the tool into an AG2-compatible format with Interoperability.
• Register the tool for execution and interaction with both the user_proxy and chatbot.

class Player(BaseModel):
    name: str
    age: int


def get_player(ctx: RunContext[Player], additional_info: Optional[str] = None) -> str:  # type: ignore[valid-type]
    """Get the player's name.

    Args:
        additional_info: Additional information which can be used.
    """
    return f"Name: {ctx.deps.name}, Age: {ctx.deps.age}, Additional info: {additional_info}"  # type: ignore[attr-defined]


interop = Interoperability()
pydantic_ai_tool = PydanticAITool(get_player, takes_ctx=True)

# player will be injected as a dependency
player = Player(name="Luka", age=25)
ag2_tool = interop.convert_tool(tool=pydantic_ai_tool, type="pydanticai", deps=player)

ag2_tool.register_for_execution(user_proxy)
ag2_tool.register_for_llm(chatbot)

Initiating the chat

Now that everything is set up, you can initiate a chat between the UserProxyAgent and the AssistantAgent:

• The user_proxy sends a message to the chatbot.
• The user requests player information, and includes “goal keeper” as additional context.
• The Player data is securely injected into the tool, and the chatbot can access and use it during the chat.

user_proxy.initiate_chat(
    recipient=chatbot, message="Get player, for additional information use 'goal keeper'", max_turns=3
)

Output

User (to chatbot):

Get player, for additional information use 'goal keeper'

--------------------------------------------------------------------------------
chatbot (to User):

***** Suggested tool call (call_lPXIohFiJfnjmgwDnNFPQCzc): get_player *****
Arguments:
{"additional_info":"goal keeper"}
***************************************************************************

--------------------------------------------------------------------------------

>>>>>>>> EXECUTING FUNCTION get_player...
User (to chatbot):

***** Response from calling tool (call_lPXIohFiJfnjmgwDnNFPQCzc) *****
Name: Luka, Age: 25, Additional info: goal keeper
**********************************************************************

--------------------------------------------------------------------------------
chatbot (to User):

The player's name is Luka, who is a 25-year-old goalkeeper. TERMINATE

Edit this page