From 026093e4360dd2d51abb2122eefc4a444750eaf7 Mon Sep 17 00:00:00 2001 From: Sumeet213 Date: Fri, 4 Oct 2024 03:57:57 +0530 Subject: [PATCH] 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()