Merge branch 'main' into alex/operation.cost

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -12,7 +12,7 @@
 import anyio
 from opentelemetry.trace import NoOpTracer, use_span
 from pydantic.json_schema import GenerateJsonSchema
-from typing_extensions import TypeVar, deprecated
+from typing_extensions import Self, TypeVar, deprecated
 
 from pydantic_graph import Graph
 
@@ -1355,7 +1355,7 @@ def _prepare_output_schema(
 
         return schema  # pyright: ignore[reportReturnType]
 
-    async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
+    async def __aenter__(self) -> Self:
         """Enter the agent context.
 
         This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered as `toolsets` so they are ready to be used.

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -20,6 +20,7 @@
 
 from pydantic_ai.tools import RunContext, ToolDefinition
 
+from .direct import model_request
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
 try:
@@ -329,11 +330,7 @@ async def _sampling_callback(
         if stop_sequences := params.stopSequences:  # pragma: no branch
             model_settings['stop_sequences'] = stop_sequences
 
-        model_response = await self.sampling_model.request(
-            pai_messages,
-            model_settings,
-            models.ModelRequestParameters(),
-        )
+        model_response = await model_request(self.sampling_model, pai_messages, model_settings=model_settings)
         return mcp_types.CreateMessageResult(
             role='assistant',
             content=_mcp.map_from_model_response(model_response),

tests/test_live.py

@@ -34,23 +34,30 @@ def gemini(_: httpx.AsyncClient, _tmp_path: Path) -> Model:
     return GoogleModel('gemini-1.5-pro')
 
 
-def vertexai(_: httpx.AsyncClient, tmp_path: Path) -> Model:
+def vertexai(http_client: httpx.AsyncClient, tmp_path: Path) -> Model:
     from google.oauth2 import service_account
 
     from pydantic_ai.models.google import GoogleModel
     from pydantic_ai.providers.google import GoogleProvider
 
-    service_account_content = os.environ['GOOGLE_SERVICE_ACCOUNT_CONTENT']
-    project_id = json.loads(service_account_content)['project_id']
-    service_account_path = tmp_path / 'service_account.json'
-    service_account_path.write_text(service_account_content)
+    if service_account_path := os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'):
+        project_id = json.loads(Path(service_account_path).read_text())['project_id']
+    elif service_account_content := os.environ.get('GOOGLE_SERVICE_ACCOUNT_CONTENT'):
+        project_id = json.loads(service_account_content)['project_id']
+        service_account_path = tmp_path / 'service_account.json'
+        service_account_path.write_text(service_account_content)
+    else:
+        pytest.skip(
+            'VertexAI live test requires GOOGLE_APPLICATION_CREDENTIALS or GOOGLE_SERVICE_ACCOUNT_CONTENT to be set'
+        )
 
     credentials = service_account.Credentials.from_service_account_file(  # type: ignore[reportUnknownReturnType]
         service_account_path,
         scopes=['https://www.googleapis.com/auth/cloud-platform'],
     )
     provider = GoogleProvider(credentials=credentials, project=project_id)
-    return GoogleModel('gemini-1.5-flash', provider=provider)
+    provider.client.aio._api_client._async_httpx_client = http_client  # type: ignore
+    return GoogleModel('gemini-2.0-flash', provider=provider)
 
 
 def groq(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model: