Skip to content

Commit db8221a

Browse files
authored
fix: thread-safe discovery API setup (Issues#327) (#583)
1 parent 3d91916 commit db8221a

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

google/generativeai/client.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import inspect
66
import dataclasses
77
import pathlib
8+
import threading
89
from typing import Any, cast
910
from collections.abc import Sequence
1011
import httplib2
@@ -64,6 +65,7 @@ def patch_colab_gce_credentials():
6465
class FileServiceClient(glm.FileServiceClient):
6566
def __init__(self, *args, **kwargs):
6667
self._discovery_api = None
68+
self._local = threading.local()
6769
super().__init__(*args, **kwargs)
6870

6971
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]] = ()):
8385
request.http.close()
8486

8587
discovery_doc = content.decode("utf-8")
86-
self._discovery_api = googleapiclient.discovery.build_from_document(
88+
self._local.discovery_api = googleapiclient.discovery.build_from_document(
8789
discovery_doc, developerKey=api_key
8890
)
8991

@@ -115,7 +117,7 @@ def create_file(
115117
filename=path, mimetype=mime_type, resumable=resumable
116118
)
117119

118-
request = self._discovery_api.media().upload(body={"file": file}, media_body=media)
120+
request = self._local.discovery_api.media().upload(body={"file": file}, media_body=media)
119121
for key, value in metadata:
120122
request.headers[key] = value
121123
result = request.execute()

0 commit comments

Comments
 (0)