From 2a678cd894f0905e6608f320c03fefff8d24302b Mon Sep 17 00:00:00 2001 From: unknown Date: Fri, 4 Oct 2024 03:02:09 +0530 Subject: [PATCH 1/2] fix: thread-safe discovery API setup (Issues#327) --- google/generativeai/client.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index d2eb6b1c9..876e4955c 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -5,6 +5,7 @@ import inspect import dataclasses import pathlib +import threading from typing import Any, cast from collections.abc import Sequence import httplib2 @@ -63,7 +64,7 @@ def patch_colab_gce_credentials(): class FileServiceClient(glm.FileServiceClient): def __init__(self, *args, **kwargs): - self._discovery_api = None + self._local = threading.local() super().__init__(*args, **kwargs) def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()): @@ -83,7 +84,7 @@ def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()): request.http.close() discovery_doc = content.decode("utf-8") - self._discovery_api = googleapiclient.discovery.build_from_document( + self._local.discovery_api = googleapiclient.discovery.build_from_document( discovery_doc, developerKey=api_key ) @@ -97,8 +98,7 @@ def create_file( resumable: bool = True, metadata: Sequence[tuple[str, str]] = (), ) -> protos.File: - if self._discovery_api is None: - self._setup_discovery_api(metadata) + self._setup_discovery_api(metadata) file = {} if name is not None: @@ -115,7 +115,7 @@ def create_file( filename=path, mimetype=mime_type, resumable=resumable ) - request = self._discovery_api.media().upload(body={"file": file}, media_body=media) + request = self._local.discovery_api.media().upload(body={"file": file}, media_body=media) for key, value in metadata: request.headers[key] = value result = request.execute() From 6d0490804b3506999d318bbed2f6c72f37a6a6e7 Mon Sep 17 00:00:00 2001 From: unknown Date: Fri, 4 Oct 2024 03:02:09 +0530 Subject: [PATCH 2/2] fix: thread-safe discovery API setup (Issues#327) --- google/generativeai/client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index d2eb6b1c9..c9c5c8c5b 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -5,6 +5,7 @@ import inspect import dataclasses import pathlib +import threading from typing import Any, cast from collections.abc import Sequence import httplib2 @@ -64,6 +65,7 @@ def patch_colab_gce_credentials(): class FileServiceClient(glm.FileServiceClient): def __init__(self, *args, **kwargs): self._discovery_api = None + self._local = threading.local() super().__init__(*args, **kwargs) def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()): @@ -83,7 +85,7 @@ def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()): request.http.close() discovery_doc = content.decode("utf-8") - self._discovery_api = googleapiclient.discovery.build_from_document( + self._local.discovery_api = googleapiclient.discovery.build_from_document( discovery_doc, developerKey=api_key ) @@ -115,7 +117,7 @@ def create_file( filename=path, mimetype=mime_type, resumable=resumable ) - request = self._discovery_api.media().upload(body={"file": file}, media_body=media) + request = self._local.discovery_api.media().upload(body={"file": file}, media_body=media) for key, value in metadata: request.headers[key] = value result = request.execute()