Merge pull request #21 from Azure-Samples/pydanticai-upgrade

Upgrade samples to use Pydantic AI v1

Author: pamelafox

AGENTS.md

@@ -0,0 +1,9 @@
+# Python AI Agents Demos
+
+This repository contains examples of several popular Python AI agent frameworks,
+in both english and spanish.
+
+All examples should be the same functionally across english and spanish,
+with the prompts and data in the respective language.
+
+All examples should support Azure OpenAI, GitHub Models, and Ollama.

README.md

@@ -98,7 +98,11 @@ You can run the examples in this repository by executing the scripts in the `exa
 | openai_agents_basic.py | Uses the OpenAI Agents framework to build a single agent. |
 | openai_agents.py | Uses the OpenAI Agents framework to handoff between several agents with tools. |
 | openai_functioncalling.py | Uses OpenAI Function Calling to call functions based on LLM output. |
-| pydanticai.py | Uses PydanticAI to build a two-agent sequential workflow for flight planning. |
+| pydanticai_basic.py | Uses PydanticAI to build a basic single agent (Spanish tutor). |
+| pydanticai_multiagent.py | Uses PydanticAI to build a two-agent sequential workflow (flight + seat selection). |
+| pydanticai_graph.py | Uses PydanticAI with pydantic-graph to build a small question/answer evaluation graph. |
+| pydanticai_tools.py | Uses PydanticAI with multiple Python tools for weekend activity planning. |
+| pydanticai_mcp_http.py | Uses PydanticAI with an MCP HTTP server toolset for travel planning (hotel search). |
 | semantickernel.py | Uses Semantic Kernel to build a writer/editor two-agent workflow. |
 | smolagents_codeagent.py | Uses SmolAgents to build a question-answering agent that can search the web and run code. |
 

examples/langgraph_agent.py

@@ -1,5 +1,4 @@
 # https://github.com/JRAlexander/IntroToAgents1-Oxford/blob/main/intro-langgraph/time-travel.ipynb
-
 import os
 
 import azure.identity
@@ -41,13 +40,13 @@ def play_song_on_apple(song: str):
         openai_api_version=os.environ["AZURE_OPENAI_VERSION"],
         azure_ad_token_provider=token_provider,
     )
-else:
+elif API_HOST == "github":
     model = ChatOpenAI(model=os.getenv("GITHUB_MODEL", "gpt-4o"), base_url="https://models.inference.ai.azure.com", api_key=os.environ["GITHUB_TOKEN"])
+elif API_HOST == "ollama":
+    model = ChatOpenAI(model=os.environ["OLLAMA_MODEL"], base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
 
 model = model.bind_tools(tools, parallel_tool_calls=False)
 
-# Define nodes and conditional edges
-
 
 # Define the function that determines whether to continue or not
 def should_continue(state):

examples/openai_agents_tools.py

@@ -6,7 +6,7 @@
 
 import azure.identity
 import openai
-from agents import Agent, Runner, function_tool, set_tracing_disabled
+from agents import Agent, OpenAIChatCompletionsModel, Runner, function_tool, set_tracing_disabled
 from dotenv import load_dotenv
 from rich.logging import RichHandler
 
@@ -32,8 +32,8 @@
     )
     MODEL_NAME = os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"]
 elif API_HOST == "ollama":
-    client = openai.AsyncOpenAI(base_url="http://localhost:11434/v1", api_key="none")
-    MODEL_NAME = "llama3.1:latest"
+    client = openai.AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    MODEL_NAME = os.environ["OLLAMA_MODEL"]
 
 
 @function_tool
@@ -70,7 +70,12 @@ def get_current_date() -> str:
     return datetime.now().strftime("%Y-%m-%d")
 
 
-agent = Agent(name="Weekend Planner", instructions="You help users plan their weekends and choose the best activities for the given weather. If an activity would be unpleasant in the weather, don't suggest it. Include the date of the weekend in your response.", tools=[get_weather, get_activities, get_current_date])
+agent = Agent(
+    name="Weekend Planner",
+    instructions="You help users plan their weekends and choose the best activities for the given weather. If an activity would be unpleasant in the weather, don't suggest it. Include the date of the weekend in your response.",
+    tools=[get_weather, get_activities, get_current_date],
+    model=OpenAIChatCompletionsModel(model=MODEL_NAME, openai_client=client),
+)
 
 
 async def main():

examples/pydanticai_basic.py

@@ -5,7 +5,7 @@
 from dotenv import load_dotenv
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 
 # Setup the OpenAI client to use either Azure OpenAI or GitHub Models
@@ -14,26 +14,25 @@
 
 if API_HOST == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "azure":
     token_provider = azure.identity.get_bearer_token_provider(azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
-    client = AsyncAzureOpenAI(
-        api_version=os.environ["AZURE_OPENAI_VERSION"],
-        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
-        azure_ad_token_provider=token_provider,
-    )
-    model = OpenAIModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+    client = AsyncAzureOpenAI(api_version=os.environ["AZURE_OPENAI_VERSION"], azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"], azure_ad_token_provider=token_provider)
+    model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "ollama":
+    client = AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    model = OpenAIChatModel(os.environ["OLLAMA_MODEL"], provider=OpenAIProvider(openai_client=client))
 
 agent: Agent[None, str] = Agent(
     model,
     system_prompt="You are a Spanish tutor. Help the user learn Spanish. ONLY respond in Spanish.",
-    result_type=str,
+    output_type=str,
 )
 
 
 async def main():
     result = await agent.run("oh hey how are you?")
-    print(result.data)
+    print(result.output)
 
 
 if __name__ == "__main__":

examples/pydanticai_graph.py

@@ -8,10 +8,9 @@
 from dotenv import load_dotenv
 from groq import BaseModel
 from openai import AsyncAzureOpenAI, AsyncOpenAI
-from pydantic_ai import Agent
-from pydantic_ai.format_as_xml import format_as_xml
+from pydantic_ai import Agent, format_as_xml
 from pydantic_ai.messages import ModelMessage
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 from pydantic_graph import (
     BaseNode,
@@ -26,24 +25,24 @@
 
 if API_HOST == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "azure":
     token_provider = azure.identity.get_bearer_token_provider(azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
     client = AsyncAzureOpenAI(
         api_version=os.environ["AZURE_OPENAI_VERSION"],
         azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
         azure_ad_token_provider=token_provider,
     )
-    model = OpenAIModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "ollama":
-    client = AsyncOpenAI(base_url="http://localhost:11434/v1", api_key="none")
-    model = OpenAIModel("llama3.1:latest", provider=OpenAIProvider(openai_client=client))
+    client = AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    model = OpenAIChatModel(os.environ["OLLAMA_MODEL"], provider=OpenAIProvider(openai_client=client))
 
 """
 Agent definitions
 """
 
-ask_agent = Agent(model, result_type=str, instrument=True)
+ask_agent = Agent(model, output_type=str)
 
 
 class EvaluationResult(BaseModel, use_attribute_docstrings=True):
@@ -55,7 +54,7 @@ class EvaluationResult(BaseModel, use_attribute_docstrings=True):
 
 evaluate_agent = Agent(
     model,
-    result_type=EvaluationResult,
+    output_type=EvaluationResult,
     system_prompt="Given a question and answer, evaluate if the answer is correct.",
 )
 
@@ -79,8 +78,8 @@ async def run(self, ctx: GraphRunContext[QuestionState]) -> Answer:
             message_history=ctx.state.ask_agent_messages,
         )
         ctx.state.ask_agent_messages += result.all_messages()
-        ctx.state.question = result.data
-        return Answer(result.data)
+        ctx.state.question = result.output
+        return Answer(result.output)
 
 
 @dataclass
@@ -106,10 +105,10 @@ async def run(
             message_history=ctx.state.evaluate_agent_messages,
         )
         ctx.state.evaluate_agent_messages += result.all_messages()
-        if result.data.correct:
-            return End(result.data.comment)
+        if result.output.correct:
+            return End(result.output.comment)
         else:
-            return Reprimand(result.data.comment)
+            return Reprimand(result.output.comment)
 
 
 @dataclass

examples/pydanticai_mcp_http.py

@@ -7,7 +7,7 @@
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 from pydantic_ai import Agent
 from pydantic_ai.mcp import MCPServerStreamableHTTP
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 
 # Setup the OpenAI client to use either Azure OpenAI or GitHub Models
@@ -16,19 +16,20 @@
 
 if API_HOST == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "azure":
     token_provider = azure.identity.get_bearer_token_provider(azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
     client = AsyncAzureOpenAI(
         api_version=os.environ["AZURE_OPENAI_VERSION"],
         azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
         azure_ad_token_provider=token_provider,
     )
-    model = OpenAIModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
-else:
-    raise ValueError(f"Unknown API_HOST: {API_HOST}. Set it to either 'github' or 'azure'.")
+    model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "ollama":
+    client = AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    model = OpenAIChatModel(os.environ["OLLAMA_MODEL"], provider=OpenAIProvider(openai_client=client))
 
-server = MCPServerStreamableHTTP("http://localhost:8000/mcp")
+server = MCPServerStreamableHTTP(url="http://localhost:8000/mcp")
 
 agent: Agent[None, str] = Agent(model, system_prompt="You are a travel planning agent. You can help users find hotels.", output_type=str, toolsets=[server])
 
@@ -39,5 +40,5 @@ async def main():
 
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.DEBUG)
+    logging.basicConfig(level=logging.WARNING)
     asyncio.run(main())

examples/pydanticai_multiagent.py

@@ -8,7 +8,7 @@
 from pydantic import BaseModel, Field
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.messages import ModelMessage
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 from rich.prompt import Prompt
 
@@ -18,15 +18,18 @@
 
 if API_HOST == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "azure":
     token_provider = azure.identity.get_bearer_token_provider(azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
     client = AsyncAzureOpenAI(
         api_version=os.environ["AZURE_OPENAI_VERSION"],
         azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
         azure_ad_token_provider=token_provider,
     )
-    model = OpenAIModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "ollama":
+    client = AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    model = OpenAIChatModel(os.environ["OLLAMA_MODEL"], provider=OpenAIProvider(openai_client=client))
 
 
 class Flight(BaseModel):
@@ -39,7 +42,7 @@ class Failed(BaseModel):
 
 flight_search_agent = Agent(
     model,
-    result_type=Flight | Failed,
+    output_type=Flight | Failed,
     system_prompt=('Use the "flight_search" tool to find a flight from the given origin to the given destination.'),
 )
 
@@ -58,10 +61,10 @@ async def find_flight() -> Flight | None:
             "Where would you like to fly from and to?",
         )
         result = await flight_search_agent.run(prompt, message_history=message_history)
-        if isinstance(result.data, Flight):
-            return result.data
+        if isinstance(result.output, Flight):
+            return result.output
         else:
-            message_history = result.all_messages(result_tool_return_content="Please try again.")
+            message_history = result.all_messages()
 
 
 class Seat(BaseModel):
@@ -72,7 +75,7 @@ class Seat(BaseModel):
 # This agent is responsible for extracting the user's seat selection
 seat_preference_agent = Agent(
     model,
-    result_type=Seat | Failed,
+    output_type=Seat | Failed,
     system_prompt=("Extract the user's seat preference. " "Seats A and F are window seats. " "Row 1 is the front row and has extra leg room. " "Rows 14, and 20 also have extra leg room. "),
 )
 
@@ -83,8 +86,8 @@ async def find_seat() -> Seat:
         answer = Prompt.ask("What seat would you like?")
 
         result = await seat_preference_agent.run(answer, message_history=message_history)
-        if isinstance(result.data, Seat):
-            return result.data
+        if isinstance(result.output, Seat):
+            return result.output
         else:
             print("Could not understand seat preference. Please try again.")
             message_history = result.all_messages()

examples/pydanticai_tools.py

@@ -8,7 +8,7 @@
 from dotenv import load_dotenv
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 from rich.logging import RichHandler
 
@@ -23,18 +23,21 @@
 
 if API_HOST == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "azure":
     token_provider = azure.identity.get_bearer_token_provider(azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
     client = AsyncAzureOpenAI(
         api_version=os.environ["AZURE_OPENAI_VERSION"],
         azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
         azure_ad_token_provider=token_provider,
     )
-    model = OpenAIModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "ollama":
+    client = AsyncOpenAI(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none")
+    model = OpenAIChatModel(os.environ["OLLAMA_MODEL"], provider=OpenAIProvider(openai_client=client))
 
 
-def get_weather(city: str) -> str:
+def get_weather(city: str) -> dict:
     logger.info(f"Getting weather for {city}")
     if random.random() < 0.05:
         return {
@@ -58,6 +61,7 @@ def get_activities(city: str, date: str) -> list:
         {"name": "Museum", "location": city},
     ]
 
+
 def get_current_date() -> str:
     logger.info("Getting current date")
     return datetime.now().strftime("%Y-%m-%d")
@@ -69,6 +73,7 @@ def get_current_date() -> str:
     tools=[get_weather, get_activities, get_current_date],
 )
 
+
 async def main():
     result = await agent.run("what can I do for funzies this weekend in Seattle?")
     print(result.output)

examples/spanish/README.md

@@ -102,9 +102,11 @@ Podés ejecutar los ejemplos en este repositorio ejecutando los scripts en el di
 | openai_agents_tools.py | Usa el framework de Agentes de OpenAI para crear un planificador de fin de semana. |
 | openai_functioncalling.py | Usa OpenAI Function Calling para llamar funciones basadas en la salida del LLM. |
 | openai_githubmodels.py | Configuración básica para usar modelos de GitHub con la API de OpenAI. |
-| pydanticai_basic.py | Implementación básica usando PydanticAI para crear un agente estructurado. |
-| pydanticai_graph.py | Usa PydanticAI para construir un grafo de agentes para hacer preguntas y evaluar respuestas. |
-| pydanticai_multiagent.py | Usa PydanticAI para construir un flujo de trabajo secuencial de dos agentes para planificación de vuelos. |
+| pydanticai_basic.py | Usa PydanticAI para construir un agente básico (tutor de español). |
+| pydanticai_multiagent.py | Usa PydanticAI para un flujo secuencial de dos agentes (vuelo + selección de asiento). |
+| pydanticai_graph.py | Usa PydanticAI con pydantic-graph para un pequeño grafo de evaluación pregunta/respuesta. |
+| pydanticai_tools.py | Usa PydanticAI con varias herramientas de Python para planificar actividades de fin de semana. |
+| pydanticai_mcp_http.py | Usa PydanticAI con un servidor MCP HTTP como conjunto de herramientas para planificación de viajes (búsqueda de hoteles). |
 | semantickernel_basic.py | Usa Semantic Kernel para construir un agente simple que enseña español. |
 | semantickernel_groupchat.py | Usa Semantic Kernel para construir un flujo de trabajo de dos agentes escritor/editor. |
 | smolagents_codeagent.py | Usa SmolAgents para construir un agente de respuesta a preguntas que puede buscar en la web y ejecutar código. |
@@ -173,4 +175,4 @@ Este proyecto incluye infraestructura como código (IaC) para provisionar despli
 * [Documentación de OpenAI Function Calling](https://platform.openai.com/docs/guides/function-calling?api-mode=chat)
 * [Documentación de PydanticAI](https://ai.pydantic.dev/multi-agent-applications/)
 * [Documentación de Semantic Kernel](https://learn.microsoft.com/semantic-kernel/overview/)
-* [Documentación de SmolAgents](https://huggingface.co/docs/smolagents/index)
+* [Documentación de SmolAgents](https://huggingface.co/docs/smolagents/index)