From f0a4ef9896073f02f176ce5999bd2f0a11f22484 Mon Sep 17 00:00:00 2001 From: Shilpa Kancharla Date: Sat, 5 Oct 2024 09:10:28 -0600 Subject: [PATCH 1/5] Updated files.py to include asyncio --- google/generativeai/files.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index b2581bdcd..9465b9c76 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -22,6 +22,7 @@ from google.generativeai import protos from itertools import islice from io import IOBase +import asyncio from google.generativeai.types import file_types @@ -87,6 +88,8 @@ def upload_file( ) return file_types.File(response) +async def upload_file_async(*args, **kwargs): + return await asyncio.to_thread(upload_file, *args, **kwargs) def list_files(page_size=100) -> Iterable[file_types.File]: """Calls the API to list files using a supported file service.""" @@ -96,6 +99,8 @@ def list_files(page_size=100) -> Iterable[file_types.File]: for proto in response: yield file_types.File(proto) +async def list_files_async(*args, **kwargs): + return await asyncio.to_thread(list_files, *args, **kwargs) def get_file(name: str) -> file_types.File: """Calls the API to retrieve a specified file using a supported file service.""" @@ -104,6 +109,8 @@ def get_file(name: str) -> file_types.File: client = get_default_file_client() return file_types.File(client.get_file(name=name)) +async def get_file_async(*args, **kwargs): + return await asyncio.to_thread(get_file, *args, **kwargs) def delete_file(name: str | file_types.File | protos.File): """Calls the API to permanently delete a specified file using a supported file service.""" @@ -114,3 +121,6 @@ def delete_file(name: str | file_types.File | protos.File): request = protos.DeleteFileRequest(name=name) client = get_default_file_client() client.delete_file(request=request) + +async def delete_file_async(*args, **kwargs): + return await asyncio.to_thread(get_file, *args, **kwargs) \ No newline at end of file From 16cf750cba2674aa6cabb015f71e64620dd92a08 Mon Sep 17 00:00:00 2001 From: Shilpa Kancharla Date: Sat, 5 Oct 2024 09:14:29 -0600 Subject: [PATCH 2/5] Updated files.py to include asyncio --- google/generativeai/files.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 9465b9c76..635bd68cb 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -88,8 +88,10 @@ def upload_file( ) return file_types.File(response) + async def upload_file_async(*args, **kwargs): - return await asyncio.to_thread(upload_file, *args, **kwargs) + return await asyncio.to_thread(upload_file, *args, **kwargs) + def list_files(page_size=100) -> Iterable[file_types.File]: """Calls the API to list files using a supported file service.""" @@ -99,9 +101,11 @@ def list_files(page_size=100) -> Iterable[file_types.File]: for proto in response: yield file_types.File(proto) + async def list_files_async(*args, **kwargs): return await asyncio.to_thread(list_files, *args, **kwargs) + def get_file(name: str) -> file_types.File: """Calls the API to retrieve a specified file using a supported file service.""" if "/" not in name: @@ -109,9 +113,11 @@ def get_file(name: str) -> file_types.File: client = get_default_file_client() return file_types.File(client.get_file(name=name)) + async def get_file_async(*args, **kwargs): return await asyncio.to_thread(get_file, *args, **kwargs) + def delete_file(name: str | file_types.File | protos.File): """Calls the API to permanently delete a specified file using a supported file service.""" if isinstance(name, (file_types.File, protos.File)): @@ -122,5 +128,6 @@ def delete_file(name: str | file_types.File | protos.File): client = get_default_file_client() client.delete_file(request=request) + async def delete_file_async(*args, **kwargs): - return await asyncio.to_thread(get_file, *args, **kwargs) \ No newline at end of file + return await asyncio.to_thread(get_file, *args, **kwargs) From 668c5208eec240b483415078c59c641619c5c4fe Mon Sep 17 00:00:00 2001 From: Shilpa Kancharla Date: Sat, 5 Oct 2024 09:42:07 -0600 Subject: [PATCH 3/5] create async test cases for asyncio file functions --- tests/test_files_async.py | 86 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 tests/test_files_async.py diff --git a/tests/test_files_async.py b/tests/test_files_async.py new file mode 100644 index 000000000..f8322be2d --- /dev/null +++ b/tests/test_files_async.py @@ -0,0 +1,86 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import collections +from collections.abc import Iterable +import io +import os +from typing import Any, Iterable, Sequence +import unittest +import asyncio + +from google.generativeai import client as client_lib +from google.generativeai import protos +from google.generativeai import files + +from absl.testing import parameterized + +class AsyncFileServiceClient(client_lib.FileServiceClient): + async def __init__(self, test): + self.test = test + self.observed_requests = [] + self.responses = collections.defaultdict(list) + + async def create_file_async( + self, + path: str | io.IOBase | os.PathLike, + *, + mime_type: str | None = None, + name: str | None = None, + display_name: str | None = None, + resumable: bool = True, + metadata: Sequence[tuple[str, str]] = (), + ) -> protos.File: + self.observed_requests.append( + dict( + path=path, + mime_type=mime_type, + name=name, + display_name=display_name, + resumable=resumable, + ) + ) + return self.responses["create_file"].pop(0) + + async def get_file_async( + self, + request: protos.GetFileRequest, + **kwargs, + ) -> protos.File: + self.observed_requests.append(request) + return self.responses["get_file"].pop(0) + + def list_files( + self, + request: protos.ListFilesRequest, + **kwargs, + ) -> Iterable[protos.File]: + self.observed_requests.append(request) + for f in self.responses["list_files"].pop(0): + yield f + + def delete_file( + self, + request: protos.DeleteFileRequest, + **kwargs, + ): + self.observed_requests.append(request) + return + +class AsyncTests(parameterized.TestCase, unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.client = FileServiceClient(self) + + client_lib._client_manager.clients["file"] = self.client \ No newline at end of file From 3c3c66144020097126267323a5b2c559de9f8168 Mon Sep 17 00:00:00 2001 From: Shilpa Kancharla Date: Mon, 7 Oct 2024 15:59:29 -0700 Subject: [PATCH 4/5] Add sample test case for asyncio file uploads --- google/generativeai/files.py | 3 +- samples/files.py | 26 +++++++++++ tests/test_files_async.py | 86 ------------------------------------ 3 files changed, 28 insertions(+), 87 deletions(-) delete mode 100644 tests/test_files_async.py diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 635bd68cb..95c3b3a99 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -28,7 +28,8 @@ from google.generativeai.client import get_default_file_client -__all__ = ["upload_file", "get_file", "list_files", "delete_file"] +__all__ = ["upload_file", "get_file", "list_files", "delete_file", + "upload_file_async", "get_file_async", "list_files_async", "delete_file_async"] mimetypes.add_type("image/webp", ".webp") diff --git a/samples/files.py b/samples/files.py index 8f98365aa..029f3bf9f 100644 --- a/samples/files.py +++ b/samples/files.py @@ -12,11 +12,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import unittest from absl.testing import absltest import google import google.generativeai as genai import pathlib +import tempfile +import asyncio media = pathlib.Path(__file__).parents[1] / "third_party" @@ -127,5 +130,28 @@ def test_files_delete(self): # [END files_delete] +class AsyncTests(absltest.TestCase, unittest.IsolatedAsyncioTestCase): + async def test_upload_file_async(self): + import google.generativeai.files as files + tempdir = pathlib.Path(tempfile.mkdtemp()) + results = [] + async def create_and_upload_file(n: int): + fname = tempdir / str(n) + fname.write_text(str(n)) + file_obj = await files.upload_file_async(fname, mime_type='text/plain') + results.append(file_obj) + + tasks = [] + for n in range(5): + tasks.append(asyncio.create_task(create_and_upload_file(n))) + + for task in tasks: + await task + + self.assertLen(results, 5) + self.assertEqual(sorted(int(f.display_name) for f in results), + list(range(5))) + + if __name__ == "__main__": absltest.main() diff --git a/tests/test_files_async.py b/tests/test_files_async.py deleted file mode 100644 index f8322be2d..000000000 --- a/tests/test_files_async.py +++ /dev/null @@ -1,86 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2023 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import collections -from collections.abc import Iterable -import io -import os -from typing import Any, Iterable, Sequence -import unittest -import asyncio - -from google.generativeai import client as client_lib -from google.generativeai import protos -from google.generativeai import files - -from absl.testing import parameterized - -class AsyncFileServiceClient(client_lib.FileServiceClient): - async def __init__(self, test): - self.test = test - self.observed_requests = [] - self.responses = collections.defaultdict(list) - - async def create_file_async( - self, - path: str | io.IOBase | os.PathLike, - *, - mime_type: str | None = None, - name: str | None = None, - display_name: str | None = None, - resumable: bool = True, - metadata: Sequence[tuple[str, str]] = (), - ) -> protos.File: - self.observed_requests.append( - dict( - path=path, - mime_type=mime_type, - name=name, - display_name=display_name, - resumable=resumable, - ) - ) - return self.responses["create_file"].pop(0) - - async def get_file_async( - self, - request: protos.GetFileRequest, - **kwargs, - ) -> protos.File: - self.observed_requests.append(request) - return self.responses["get_file"].pop(0) - - def list_files( - self, - request: protos.ListFilesRequest, - **kwargs, - ) -> Iterable[protos.File]: - self.observed_requests.append(request) - for f in self.responses["list_files"].pop(0): - yield f - - def delete_file( - self, - request: protos.DeleteFileRequest, - **kwargs, - ): - self.observed_requests.append(request) - return - -class AsyncTests(parameterized.TestCase, unittest.IsolatedAsyncioTestCase): - def setUp(self): - self.client = FileServiceClient(self) - - client_lib._client_manager.clients["file"] = self.client \ No newline at end of file From 97f0b46aa3d31c115dd13609827b07689dbc0652 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 24 Oct 2024 11:02:46 -0700 Subject: [PATCH 5/5] format Change-Id: Id63b0ab5ce27e632ef57b1156f12907428003c8a --- google/generativeai/files.py | 12 ++++++++++-- samples/files.py | 9 +++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 95c3b3a99..0d574fe25 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -28,8 +28,16 @@ from google.generativeai.client import get_default_file_client -__all__ = ["upload_file", "get_file", "list_files", "delete_file", - "upload_file_async", "get_file_async", "list_files_async", "delete_file_async"] +__all__ = [ + "upload_file", + "get_file", + "list_files", + "delete_file", + "upload_file_async", + "get_file_async", + "list_files_async", + "delete_file_async", +] mimetypes.add_type("image/webp", ".webp") diff --git a/samples/files.py b/samples/files.py index 029f3bf9f..461e293d2 100644 --- a/samples/files.py +++ b/samples/files.py @@ -133,24 +133,25 @@ def test_files_delete(self): class AsyncTests(absltest.TestCase, unittest.IsolatedAsyncioTestCase): async def test_upload_file_async(self): import google.generativeai.files as files + tempdir = pathlib.Path(tempfile.mkdtemp()) results = [] + async def create_and_upload_file(n: int): fname = tempdir / str(n) fname.write_text(str(n)) - file_obj = await files.upload_file_async(fname, mime_type='text/plain') + file_obj = await files.upload_file_async(fname, mime_type="text/plain") results.append(file_obj) tasks = [] for n in range(5): tasks.append(asyncio.create_task(create_and_upload_file(n))) - + for task in tasks: await task self.assertLen(results, 5) - self.assertEqual(sorted(int(f.display_name) for f in results), - list(range(5))) + self.assertEqual(sorted(int(f.display_name) for f in results), list(range(5))) if __name__ == "__main__":