test structured responses

tests/test_live.py

@@ -7,6 +7,7 @@
 
 import httpx
 import pytest
+from pydantic import BaseModel
 
 from pydantic_ai import Agent
 from pydantic_ai.models.gemini import GeminiModel
@@ -45,6 +46,20 @@ async def test_openai_stream(http_client: httpx.AsyncClient):
     assert cost.total_tokens is not None and cost.total_tokens > 0
 
 
+class MyModel(BaseModel):
+    city: str
+
+
+async def test_openai_structured(http_client: httpx.AsyncClient):
+    agent = Agent(OpenAIModel('gpt-4o-mini', http_client=http_client), result_type=MyModel)
+    result = await agent.run('What is the capital of the UK?')
+    print('OpenAI structured response:', result.data)
+    assert result.data.city.lower() == 'london'
+    print('OpenAI structured cost:', result.cost())
+    cost = result.cost()
+    assert cost.total_tokens is not None and cost.total_tokens > 0
+
+
 async def test_gemini(http_client: httpx.AsyncClient):
     agent = Agent(GeminiModel('gemini-1.5-flash', http_client=http_client))
     result = await agent.run('What is the capital of France?')
@@ -56,11 +71,21 @@ async def test_gemini(http_client: httpx.AsyncClient):
 
 
 async def test_gemini_stream(http_client: httpx.AsyncClient):
-    agent = Agent(GeminiModel('gemini-1.5-flash', http_client=http_client))
+    agent = Agent(GeminiModel('gemini-1.5-pro', http_client=http_client))
     async with agent.run_stream('What is the capital of France?') as result:
         data = await result.get_data()
     print('Gemini stream response:', data)
     assert 'paris' in data.lower()
     print('Gemini stream cost:', result.cost())
     cost = result.cost()
     assert cost.total_tokens is not None and cost.total_tokens > 0
+
+
+async def test_gemini_structured(http_client: httpx.AsyncClient):
+    agent = Agent(GeminiModel('gemini-1.5-pro', http_client=http_client), result_type=MyModel)
+    result = await agent.run('What is the capital of the UK?')
+    print('Gemini structured response:', result.data)
+    assert result.data.city.lower() == 'london'
+    print('Gemini structured cost:', result.cost())
+    cost = result.cost()
+    assert cost.total_tokens is not None and cost.total_tokens > 0