Merge pull request #139 from mistralai/harizo/chore-improve-gcp-auth

GCP client improvements

packages/mistralai_gcp/.speakeasy/gen.yaml

@@ -18,7 +18,7 @@ python:
       pytest: ^8.2.2
       pytest-asyncio: ^0.23.7
     main:
-      google-auth: ^2.31.0
+      google-auth: 2.27.0
       requests: ^2.32.3
   authors:
     - Mistral

packages/mistralai_gcp/src/mistralai_gcp/sdk.py

@@ -32,6 +32,7 @@ def __init__(
         self,
         region: str = "europe-west4",
         project_id: Optional[str] = None,
+        access_token: Optional[str] = None,
         client: Optional[HttpClient] = None,
         async_client: Optional[AsyncHttpClient] = None,
         retry_config: Optional[Nullable[RetryConfig]] = None,
@@ -46,26 +47,30 @@ def __init__(
         :param retry_config: The retry configuration to use for all supported methods
         """
 
-        credentials, loaded_project_id = google.auth.default(
-            scopes=["https://www.googleapis.com/auth/cloud-platform"],
-        )
-
-        if not isinstance(credentials, google.auth.credentials.Credentials):
-            raise models.SDKError(
-                "credentials must be an instance of google.auth.credentials.Credentials"
+        if not access_token:
+            credentials, loaded_project_id = google.auth.default(
+                scopes=["https://www.googleapis.com/auth/cloud-platform"],
             )
+            credentials.refresh(google.auth.transport.requests.Request())
+
+            if not isinstance(credentials, google.auth.credentials.Credentials):
+                raise models.SDKError(
+                    "credentials must be an instance of google.auth.credentials.Credentials"
+                )
 
         project_id = project_id or loaded_project_id
         if project_id is None:
             raise models.SDKError("project_id must be provided")
 
         def auth_token() -> str:
-            if credentials.expired:
+            if access_token:
+                return access_token
+            else:
                 credentials.refresh(google.auth.transport.requests.Request())
-            token = credentials.token
-            if not token:
-                raise models.SDKError("Failed to get token from credentials")
-            return token
+                token = credentials.token
+                if not token:
+                    raise models.SDKError("Failed to get token from credentials")
+                return token
 
         if client is None:
             client = httpx.Client()

pyproject.toml

@@ -21,7 +21,7 @@ jsonpath-python = "^1.0.6"
 pydantic = "~2.8.2"
 python-dateutil = "^2.9.0.post0"
 typing-inspect = "^0.9.0"
-google-auth = { version = "^2.31.0", optional = true }
+google-auth = { version = "2.27.0", optional = true }
 requests = { version = "^2.32.3", optional = true }
 
 [tool.poetry.group.dev.dependencies]