Use Pydantic AI properly

Author: pamelafox

examples/agentframework_basic.py

@@ -15,11 +15,23 @@
 if API_HOST == "azure":
     async_credential = DefaultAzureCredential()
     token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
-    client = OpenAIChatClient(base_url=f"{os.environ['AZURE_OPENAI_ENDPOINT']}/openai/v1/", api_key=token_provider, model_id=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"])
+    client = OpenAIChatClient(
+        base_url=f"{os.environ['AZURE_OPENAI_ENDPOINT']}/openai/v1/",
+        api_key=token_provider,
+        model_id=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"],
+    )
 elif API_HOST == "github":
-    client = OpenAIChatClient(base_url="https://models.github.ai/inference", api_key=os.environ["GITHUB_TOKEN"], model_id=os.getenv("GITHUB_MODEL", "openai/gpt-4o"))
+    client = OpenAIChatClient(
+        base_url="https://models.github.ai/inference",
+        api_key=os.environ["GITHUB_TOKEN"],
+        model_id=os.getenv("GITHUB_MODEL", "openai/gpt-4o"),
+    )
 elif API_HOST == "ollama":
-    client = OpenAIChatClient(base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"), api_key="none", model_id=os.environ.get("OLLAMA_MODEL", "llama3.1:latest"))
+    client = OpenAIChatClient(
+        base_url=os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434/v1"),
+        api_key="none",
+        model_id=os.environ.get("OLLAMA_MODEL", "llama3.1:latest"),
+    )
 else:
     client = OpenAIChatClient(api_key=os.environ["OPENAI_API_KEY"], model_id=os.environ.get("OPENAI_MODEL", "gpt-4o"))
 

examples/pydanticai_basic.py

@@ -1,8 +1,7 @@
 import asyncio
 import os
 
-from azure.identity import DefaultAzureCredential
-from azure.identity.aio import get_bearer_token_provider
+from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 from pydantic_ai import Agent
@@ -13,19 +12,24 @@
 load_dotenv(override=True)
 API_HOST = os.getenv("API_HOST", "github")
 
-if API_HOST == "github":
-    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
-elif API_HOST == "azure":
-    token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
+async_credential = None
+if API_HOST == "azure":
+    async_credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
     client = AsyncOpenAI(
         base_url=os.environ["AZURE_OPENAI_ENDPOINT"] + "/openai/v1",
         api_key=token_provider,
     )
     model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "github":
+    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), 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))
+else:
+    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
+    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 
 agent: Agent[None, str] = Agent(
     model,
@@ -38,6 +42,9 @@ async def main():
     result = await agent.run("oh hey how are you?")
     print(result.output)
 
+    if async_credential:
+        await async_credential.close()
+
 
 if __name__ == "__main__":
     asyncio.run(main())

examples/pydanticai_graph.py

@@ -4,8 +4,7 @@
 import os
 from dataclasses import dataclass, field
 
-from azure.identity import DefaultAzureCredential
-from azure.identity.aio import get_bearer_token_provider
+from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from dotenv import load_dotenv
 from groq import BaseModel
 from openai import AsyncOpenAI
@@ -24,19 +23,24 @@
 load_dotenv(override=True)
 API_HOST = os.getenv("API_HOST", "github")
 
-if API_HOST == "github":
-    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
-elif API_HOST == "azure":
-    token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
+async_credential = None
+if API_HOST == "azure":
+    async_credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
     client = AsyncOpenAI(
         base_url=os.environ["AZURE_OPENAI_ENDPOINT"] + "/openai/v1",
         api_key=token_provider,
     )
     model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "github":
+    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), 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))
+else:
+    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
+    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 
 """
 Agent definitions
@@ -130,6 +134,9 @@ async def main():
     end = await question_graph.run(node, state=state)
     print("END:", end.output)
 
+    if async_credential:
+        await async_credential.close()
+
 
 if __name__ == "__main__":
     asyncio.run(main())

examples/pydanticai_mcp_github.py

@@ -18,8 +18,7 @@
 import logging
 import os
 
-from azure.identity import DefaultAzureCredential
-from azure.identity.aio import get_bearer_token_provider
+from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 from pydantic import BaseModel, Field
@@ -40,26 +39,24 @@
 load_dotenv(override=True)
 API_HOST = os.getenv("API_HOST", "github")
 
-
+async_credential = None
 if API_HOST == "azure":
-    token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
+    async_credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
     client = AsyncOpenAI(
         base_url=os.environ["AZURE_OPENAI_ENDPOINT"] + "/openai/v1",
         api_key=token_provider,
     )
-    model = OpenAIChatModel(
-        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 == "github":
     client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIChatModel(os.environ.get("GITHUB_MODEL", "gpt-4o-mini"), provider=OpenAIProvider(openai_client=client))
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 elif API_HOST == "ollama":
-    client = AsyncOpenAI(base_url=os.environ["OLLAMA_ENDPOINT"], api_key="none")
+    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))
 else:
-    client = AsyncOpenAI()
-    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o-mini"), provider=OpenAIProvider(openai_client=client))
+    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
+    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 
 
 class IssueProposal(BaseModel):
@@ -73,13 +70,18 @@ class IssueProposal(BaseModel):
 
 
 async def main():
-    server = MCPServerStreamableHTTP(url="https://api.githubcopilot.com/mcp/", headers={"Authorization": f"Bearer {os.getenv('GITHUB_TOKEN', '')}"})
+    server = MCPServerStreamableHTTP(
+        url="https://api.githubcopilot.com/mcp/", headers={"Authorization": f"Bearer {os.getenv('GITHUB_TOKEN', '')}"}
+    )
     desired_tool_names = ("list_issues", "search_code", "search_issues", "search_pull_requests")
     filtered_tools = server.filtered(lambda ctx, tool_def: tool_def.name in desired_tool_names)
 
     agent: Agent[None, IssueProposal] = Agent(
         model,
-        system_prompt=("You are an issue triage assistant. Use the provided tools to find an issue that can be closed " "and produce an IssueProposal."),
+        system_prompt=(
+            "You are an issue triage assistant. Use the provided tools to find an issue that can be closed "
+            "and produce an IssueProposal."
+        ),
         output_type=IssueProposal,
         toolsets=[filtered_tools],
     )
@@ -96,6 +98,9 @@ async def main():
 
     print(agent_run.result.output)
 
+    if async_credential:
+        await async_credential.close()
+
 
 if __name__ == "__main__":
     logger.setLevel(logging.INFO)

examples/pydanticai_mcp_http.py

@@ -9,8 +9,7 @@
 import logging
 import os
 
-from azure.identity import DefaultAzureCredential
-from azure.identity.aio import get_bearer_token_provider
+from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 from pydantic_ai import Agent
@@ -22,19 +21,24 @@
 load_dotenv(override=True)
 API_HOST = os.getenv("API_HOST", "github")
 
-if API_HOST == "github":
-    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
-elif API_HOST == "azure":
-    token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
+async_credential = None
+if API_HOST == "azure":
+    async_credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
     client = AsyncOpenAI(
         base_url=os.environ["AZURE_OPENAI_ENDPOINT"] + "/openai/v1",
         api_key=token_provider,
     )
     model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "github":
+    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), 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))
+else:
+    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
+    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 
 server = MCPServerStreamableHTTP(url="http://localhost:8000/mcp")
 
@@ -47,9 +51,14 @@
 
 
 async def main():
-    result = await agent.run("Find me a hotel in San Francisco for 2 nights starting from 2024-01-01. I need a hotel with free WiFi and a pool.")
+    result = await agent.run(
+        "Find me a hotel in San Francisco for 2 nights starting from 2024-01-01. I need free WiFi and a pool."
+    )
     print(result.output)
 
+    if async_credential:
+        await async_credential.close()
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.WARNING)

examples/pydanticai_multiagent.py

@@ -2,8 +2,7 @@
 import os
 from typing import Literal
 
-from azure.identity import DefaultAzureCredential
-from azure.identity.aio import get_bearer_token_provider
+from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
 from dotenv import load_dotenv
 from openai import AsyncOpenAI
 from pydantic import BaseModel, Field
@@ -17,19 +16,24 @@
 load_dotenv(override=True)
 API_HOST = os.getenv("API_HOST", "github")
 
-if API_HOST == "github":
-    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
-    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
-elif API_HOST == "azure":
-    token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")
+async_credential = None
+if API_HOST == "azure":
+    async_credential = DefaultAzureCredential()
+    token_provider = get_bearer_token_provider(async_credential, "https://cognitiveservices.azure.com/.default")
     client = AsyncOpenAI(
         base_url=os.environ["AZURE_OPENAI_ENDPOINT"] + "/openai/v1",
         api_key=token_provider,
     )
     model = OpenAIChatModel(os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=client))
+elif API_HOST == "github":
+    client = AsyncOpenAI(api_key=os.environ["GITHUB_TOKEN"], base_url="https://models.inference.ai.azure.com")
+    model = OpenAIChatModel(os.getenv("GITHUB_MODEL", "gpt-4o"), 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))
+else:
+    client = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
+    model = OpenAIChatModel(os.environ.get("OPENAI_MODEL", "gpt-4o"), provider=OpenAIProvider(openai_client=client))
 
 
 class Flight(BaseModel):
@@ -76,7 +80,12 @@ class Seat(BaseModel):
 seat_preference_agent = Agent(
     model,
     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. "),
+    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. "
+    ),
 )
 
 
@@ -100,6 +109,9 @@ async def main():
         seat_preference = await find_seat()
         print(f"Seat preference: {seat_preference}")
 
+    if async_credential:
+        await async_credential.close()
+
 
 if __name__ == "__main__":
     asyncio.run(main())