Skip to content

Commit

Permalink
Add support for resumable uploads for File API (google-gemini#275)
Browse files Browse the repository at this point in the history
* add support for resumable uploads

* run formatter

* update docs

* simplify resumable flow and set as default
  • Loading branch information
TYMichaelChen authored May 13, 2024
1 parent 855f96f commit efead6b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
5 changes: 4 additions & 1 deletion google/generativeai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ def create_file(
mime_type: str | None = None,
name: str | None = None,
display_name: str | None = None,
resumable: bool = True,
) -> glm.File:
if self._discovery_api is None:
self._setup_discovery_api()
Expand All @@ -69,7 +70,9 @@ def create_file(
if display_name is not None:
file["displayName"] = display_name

media = googleapiclient.http.MediaFileUpload(filename=path, mimetype=mime_type)
media = googleapiclient.http.MediaFileUpload(
filename=path, mimetype=mime_type, resumable=resumable
)
request = self._discovery_api.media().upload(body={"file": file}, media_body=media)
result = request.execute()

Expand Down
19 changes: 18 additions & 1 deletion google/generativeai/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,24 @@ def upload_file(
mime_type: str | None = None,
name: str | None = None,
display_name: str | None = None,
resumable: bool = True,
) -> file_types.File:
"""Uploads a file using a supported file service.
Args:
path: The path to the file 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').
If not provided, a system generated ID will be created.
display_name: Optional display name of the file.
resumable: Whether to use the resumable upload protocol. By default, this is enabled.
See details at
https://googleapis.github.io/google-api-python-client/docs/epy/googleapiclient.http.MediaFileUpload-class.html#resumable
Returns:
file_types.File: The response of the uploaded file.
"""
client = get_default_file_client()

path = pathlib.Path(os.fspath(path))
Expand All @@ -50,7 +67,7 @@ def upload_file(
display_name = path.name

response = client.create_file(
path=path, mime_type=mime_type, name=name, display_name=display_name
path=path, mime_type=mime_type, name=name, display_name=display_name, resumable=resumable
)
return file_types.File(response)

Expand Down

0 comments on commit efead6b

Please sign in to comment.