Recreate access token on 401 for Google Vertex provider (#1195)

Author: Kludex

pydantic_ai_slim/pydantic_ai/providers/google_vertex.py

@@ -2,7 +2,6 @@
 
 import functools
 from collections.abc import AsyncGenerator, Mapping
-from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Literal, overload
 
@@ -28,9 +27,6 @@
 
 __all__ = ('GoogleVertexProvider',)
 
-# default expiry is 3600 seconds
-MAX_TOKEN_AGE = timedelta(seconds=3000)
-
 
 class GoogleVertexProvider(Provider[httpx.AsyncClient]):
     """Provider for Vertex AI API."""
@@ -131,19 +127,21 @@ def __init__(
         self.region = region
 
         self.credentials = None
-        self.token_created: datetime | None = None
 
     async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
         if self.credentials is None:
             self.credentials = await self._get_credentials()
-        if self.credentials.token is None or self._token_expired():  # type: ignore[reportUnknownMemberType]
-            await anyio.to_thread.run_sync(self._refresh_token)
-            self.token_created = datetime.now()
+        if self.credentials.token is None:  # type: ignore[reportUnknownMemberType]
+            await self._refresh_token()
         request.headers['Authorization'] = f'Bearer {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
-
         # NOTE: This workaround is in place because we might get the project_id from the credentials.
         request.url = httpx.URL(str(request.url).replace('projects/None', f'projects/{self.project_id}'))
-        yield request
+        response = yield request
+
+        if response.status_code == 401:
+            await self._refresh_token()
+            request.headers['Authorization'] = f'Bearer {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
+            yield request
 
     async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials:
         if self.service_account_file is not None:
@@ -166,15 +164,9 @@ async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials:
             self.project_id = creds_project_id
         return creds
 
-    def _token_expired(self) -> bool:
-        if self.token_created is None:
-            return True
-        else:
-            return (datetime.now() - self.token_created) > MAX_TOKEN_AGE
-
-    def _refresh_token(self) -> str:  # pragma: no cover
+    async def _refresh_token(self) -> str:  # pragma: no cover
         assert self.credentials is not None
-        self.credentials.refresh(Request())  # type: ignore[reportUnknownMemberType]
+        await anyio.to_thread.run_sync(self.credentials.refresh, Request())  # type: ignore[reportUnknownMemberType]
         assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
         return self.credentials.token
 

tests/providers/test_google_vertex.py

@@ -57,7 +57,10 @@ async def test_google_vertex_provider_auth(allow_model_requests: None, http_clie
     await provider.client.post('/gemini-1.0-pro:generateContent')
     assert provider.region == 'us-central1'
     assert getattr(provider.client.auth, 'project_id') == 'my-project-id'
-    assert getattr(provider.client.auth, 'token_created') is not None
+
+
+async def mock_refresh_token():
+    return 'my-token'
 
 
 async def test_google_vertex_provider_service_account_file(
@@ -67,11 +70,10 @@ async def test_google_vertex_provider_service_account_file(
     save_service_account(service_account_path, 'my-project-id')
 
     provider = GoogleVertexProvider(service_account_file=service_account_path)
-    monkeypatch.setattr(provider.client.auth, '_refresh_token', lambda: 'my-token')
+    monkeypatch.setattr(provider.client.auth, '_refresh_token', mock_refresh_token)
     await provider.client.post('/gemini-1.0-pro:generateContent')
     assert provider.region == 'us-central1'
     assert getattr(provider.client.auth, 'project_id') == 'my-project-id'
-    assert getattr(provider.client.auth, 'token_created') is not None
 
 
 async def test_google_vertex_provider_service_account_file_info(
@@ -80,11 +82,10 @@ async def test_google_vertex_provider_service_account_file_info(
     account_info = prepare_service_account_contents('my-project-id')
 
     provider = GoogleVertexProvider(service_account_info=account_info)
-    monkeypatch.setattr(provider.client.auth, '_refresh_token', lambda: 'my-token')
+    monkeypatch.setattr(provider.client.auth, '_refresh_token', mock_refresh_token)
     await provider.client.post('/gemini-1.0-pro:generateContent')
     assert provider.region == 'us-central1'
     assert getattr(provider.client.auth, 'project_id') == 'my-project-id'
-    assert getattr(provider.client.auth, 'token_created') is not None
 
 
 async def test_google_vertex_provider_service_account_xor(allow_model_requests: None):