From 33ef09fff456fbdc017cb32167e4e5bdbdcf6dcb Mon Sep 17 00:00:00 2001 From: Hamza-nabil Date: Sat, 14 Sep 2024 13:21:38 +0100 Subject: [PATCH 1/6] Add IOBase support to FileServiceClient.create_file --- google/generativeai/client.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index 938018b96..33592bcea 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -8,6 +8,7 @@ from typing import Any, cast from collections.abc import Sequence import httplib2 +from io import IOBase import google.ai.generativelanguage as glm import google.generativeai.protos as protos @@ -88,7 +89,7 @@ def _setup_discovery_api(self, metadata: dict | Sequence[tuple[str, str]] = ()): def create_file( self, - path: str | pathlib.Path | os.PathLike, + path: str | pathlib.Path | os.PathLike | IOBase, *, mime_type: str | None = None, name: str | None = None, @@ -104,10 +105,16 @@ def create_file( file["name"] = name if display_name is not None: file["displayName"] = display_name - - media = googleapiclient.http.MediaFileUpload( - filename=path, mimetype=mime_type, resumable=resumable - ) + + if isinstance(path, IOBase): + media = googleapiclient.http.MediaIoBaseUpload( + fd=path, mimetype=mime_type, resumable=resumable + ) + else: + media = googleapiclient.http.MediaFileUpload( + filename=path, mimetype=mime_type, resumable=resumable + ) + request = self._discovery_api.media().upload(body={"file": file}, media_body=media) for key, value in metadata: request.headers[key] = value From ca1c0ed0a4c7ca13b6984deade645d32fc1dbec9 Mon Sep 17 00:00:00 2001 From: Hamza-nabil Date: Sat, 14 Sep 2024 14:09:36 +0100 Subject: [PATCH 2/6] add support to upload file-like object --- google/generativeai/files.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index c0d8e1e0a..597381686 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -21,6 +21,7 @@ import logging from google.generativeai import protos from itertools import islice +from io import IOBase from google.generativeai.types import file_types @@ -32,7 +33,7 @@ def upload_file( - path: str | pathlib.Path | os.PathLike, + path: str | pathlib.Path | os.PathLike | IOBase, *, mime_type: str | None = None, name: str | None = None, @@ -42,7 +43,7 @@ def upload_file( """Calls the API to upload a file using a supported file service. Args: - path: The path to the file to be uploaded. + path: The path to the file or a file-like object (e.g., BytesIO) to be uploaded. mime_type: The MIME type of the file. If not provided, it will be inferred from the file extension. name: The name of the file in the destination (e.g., 'files/sample-image'). @@ -57,17 +58,22 @@ def upload_file( """ client = get_default_file_client() - path = pathlib.Path(os.fspath(path)) + if not isinstance(path, IOBase): + path = pathlib.Path(os.fspath(path)) + + if display_name is None: + display_name = path.name + + if mime_type is None: + mime_type, _ = mimetypes.guess_type(path) if mime_type is None: - mime_type, _ = mimetypes.guess_type(path) + # Guess failed or IOBase, use octet-stream. + mime_type = 'application/octet-stream' if name is not None and "/" not in name: name = f"files/{name}" - if display_name is None: - display_name = path.name - response = client.create_file( path=path, mime_type=mime_type, name=name, display_name=display_name, resumable=resumable ) From 0cfc19152b9ee5a3e422d55cf6c54b2adc6dcc6c Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Fri, 27 Sep 2024 11:36:55 -0700 Subject: [PATCH 3/6] Give clear errors for 'unknown' mime-types Change-Id: Iea071c396c4cfe2b2c8eacae74dd8fb0acbc128f --- google/generativeai/client.py | 4 ++-- google/generativeai/files.py | 17 +++++++++++++---- tests/test_files.py | 3 ++- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index 33592bcea..d2eb6b1c9 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -105,7 +105,7 @@ def create_file( file["name"] = name if display_name is not None: file["displayName"] = display_name - + if isinstance(path, IOBase): media = googleapiclient.http.MediaIoBaseUpload( fd=path, mimetype=mime_type, resumable=resumable @@ -114,7 +114,7 @@ def create_file( media = googleapiclient.http.MediaFileUpload( filename=path, mimetype=mime_type, resumable=resumable ) - + request = self._discovery_api.media().upload(body={"file": file}, media_body=media) for key, value in metadata: request.headers[key] = value diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 597381686..df3bd9119 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -58,7 +58,13 @@ def upload_file( """ client = get_default_file_client() - if not isinstance(path, IOBase): + if isinstance(path, IOBase): + if mime_type is None: + raise ValueError( + "Unknown mime type: When passing a file like object to `path` (instead of a\n" + " path-like object) you must set the `mime_type` argument" + ) + else: path = pathlib.Path(os.fspath(path)) if display_name is None: @@ -67,9 +73,12 @@ def upload_file( if mime_type is None: mime_type, _ = mimetypes.guess_type(path) - if mime_type is None: - # Guess failed or IOBase, use octet-stream. - mime_type = 'application/octet-stream' + if mime_type is None: + if mime_type is None: + raise ValueError( + "Unknown mime type: Could not determine the mimetype for your file\n" + " please set the `mime_type` argument" + ) if name is not None and "/" not in name: name = f"files/{name}" diff --git a/tests/test_files.py b/tests/test_files.py index cb48316bd..d6a022aa1 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -102,12 +102,13 @@ def test_video_metadata(self): protos.File( uri="https://test", state="ACTIVE", + mime_type="video/quicktime", video_metadata=dict(video_duration=datetime.timedelta(seconds=30)), error=dict(code=7, message="ok?"), ) ) - f = genai.upload_file(path="dummy") + f = genai.upload_file(path="dummy.mov") self.assertEqual(google.rpc.status_pb2.Status(code=7, message="ok?"), f.error) self.assertEqual( protos.VideoMetadata(dict(video_duration=datetime.timedelta(seconds=30))), From 7b2c5efa79c55382109fef06c9d6c76cc3185390 Mon Sep 17 00:00:00 2001 From: Hamza-nabil Date: Fri, 27 Sep 2024 20:51:21 +0100 Subject: [PATCH 4/6] Remove duplicate check for mime_type --- google/generativeai/files.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index df3bd9119..b2581bdcd 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -74,11 +74,10 @@ def upload_file( mime_type, _ = mimetypes.guess_type(path) if mime_type is None: - if mime_type is None: - raise ValueError( - "Unknown mime type: Could not determine the mimetype for your file\n" - " please set the `mime_type` argument" - ) + raise ValueError( + "Unknown mime type: Could not determine the mimetype for your file\n" + " please set the `mime_type` argument" + ) if name is not None and "/" not in name: name = f"files/{name}" From a3a005aa873ae21e3a4ac8c624ab43c53160fcd9 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Fri, 27 Sep 2024 14:14:29 -0700 Subject: [PATCH 5/6] Add a test uploading from a file-like IO object Change-Id: I572f02ed98b9ca45299b76e7a01695fdcf917e1e --- samples/files.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/samples/files.py b/samples/files.py index cbed68a1e..8f98365aa 100644 --- a/samples/files.py +++ b/samples/files.py @@ -83,6 +83,18 @@ def test_files_create_pdf(self): print(response.text) # [END files_create_pdf] + def test_files_create_from_IO(self): + # [START files_create_io] + # You can pass a file-like object, instead of a path. + # Useful for streaming. + model = genai.GenerativeModel("gemini-1.5-flash") + fpath = media / "test.pdf" + with open(fpath, "rb") as f: + sample_pdf = genai.upload_file(f, mime_type="application/pdf") + response = model.generate_content(["Give me a summary of this pdf file.", sample_pdf]) + print(response.text) + # [END files_create_io] + def test_files_list(self): # [START files_list] print("My files:") From 8093de85dbe4b261cfe82a6b625a640b316358c8 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Fri, 27 Sep 2024 14:32:48 -0700 Subject: [PATCH 6/6] fix type check Change-Id: I220c05eee73ae76ced25254d67332f70a4069f7e --- tests/test_files.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_files.py b/tests/test_files.py index d6a022aa1..0f7ca5707 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -18,6 +18,7 @@ import collections import datetime +import io import os from typing import Iterable, Sequence import pathlib @@ -38,7 +39,7 @@ def __init__(self, test): def create_file( self, - path: str | pathlib.Path | os.PathLike, + path: str | io.IOBase | os.PathLike, *, mime_type: str | None = None, name: str | None = None,