Skip to content

Pydanticai

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, LLMConfig
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.

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

with llm_config:
    chatbot = AssistantAgent(name="chatbot")

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