|
| 1 | +import requests |
| 2 | +from pydantic import BaseModel |
| 3 | +from pydantic_ai import Agent, RunContext |
| 4 | + |
| 5 | + |
| 6 | +class UserDatabase: |
| 7 | + """Simulate a user database using the JSONPlaceholder users API.""" |
| 8 | + |
| 9 | + _base_url = "https://jsonplaceholder.typicode.com" |
| 10 | + |
| 11 | + def get_user_info(self, user_id: int) -> dict: |
| 12 | + response = requests.get(f"{self._base_url}/users/{user_id}") |
| 13 | + response.raise_for_status() |
| 14 | + return response.json() |
| 15 | + |
| 16 | + |
| 17 | +class UserSummary(BaseModel): |
| 18 | + name: str |
| 19 | + email: str |
| 20 | + company: str |
| 21 | + |
| 22 | + |
| 23 | +agent = Agent( |
| 24 | + "google-gla:gemini-2.5-flash", |
| 25 | + output_type=UserSummary, |
| 26 | + deps_type=UserDatabase, |
| 27 | + instructions=( |
| 28 | + "You retrieve user information from an external database. " |
| 29 | + "Use the available tools to gather user info, " |
| 30 | + "then return a structured summary." |
| 31 | + ), |
| 32 | +) |
| 33 | + |
| 34 | + |
| 35 | +@agent.tool |
| 36 | +def fetch_user(ctx: RunContext[UserDatabase], user_id: int) -> str: |
| 37 | + """Fetch user profile from the service.""" |
| 38 | + try: |
| 39 | + user = ctx.deps.get_user_info(user_id) |
| 40 | + return str(user) |
| 41 | + except requests.HTTPError: |
| 42 | + return f"User with ID {user_id} not found" |
| 43 | + |
| 44 | + |
| 45 | +db = UserDatabase() |
| 46 | +result = agent.run_sync( |
| 47 | + "Get a summary for user 7", |
| 48 | + deps=db, |
| 49 | +) # Inject the database |
| 50 | +print(f"Name: {result.output.name}") |
| 51 | +print(f"Email: {result.output.email}") |
| 52 | +print(f"Company: {result.output.company}") |
0 commit comments