From 14bb2bbae14c2109435b7214545f9b8dda72c754 Mon Sep 17 00:00:00 2001 From: Michael Chen Date: Thu, 28 Mar 2024 11:14:06 -0400 Subject: [PATCH 1/5] Add additional functionalities with files SDK --- google/generativeai/__init__.py | 1 + google/generativeai/client.py | 10 ++++------ google/generativeai/files.py | 15 ++++++++++++--- google/generativeai/generative_models.py | 8 ++++---- google/generativeai/types/file_types.py | 4 ++++ 5 files changed, 25 insertions(+), 13 deletions(-) diff --git a/google/generativeai/__init__.py b/google/generativeai/__init__.py index a236cffbe..53383a1b3 100644 --- a/google/generativeai/__init__.py +++ b/google/generativeai/__init__.py @@ -56,6 +56,7 @@ from google.generativeai.files import upload_file from google.generativeai.files import get_file from google.generativeai.files import list_files +from google.generativeai.files import delete_file from google.generativeai.generative_models import GenerativeModel from google.generativeai.generative_models import ChatSession diff --git a/google/generativeai/client.py b/google/generativeai/client.py index f9217b326..4415ec8f2 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -27,7 +27,7 @@ __version__ = "0.0.0" USER_AGENT = "genai-py" - +GENAI_API_DISCOVERY_URL="https://generativelanguage.googleapis.com/$discovery/rest" class FileServiceClient(glm.FileServiceClient): def __init__(self, *args, **kwargs): @@ -37,14 +37,12 @@ def __init__(self, *args, **kwargs): def _setup_discovery_api(self): api_key = self._client_options.api_key if api_key is None: - raise ValueError("Uploading to the Files API requires an api key.") - - end_point = self.api_endpoint + raise ValueError("Uploading to the File API requires an API key.") request = googleapiclient.http.HttpRequest( http=httplib2.Http(), postproc=lambda resp, content: (resp, content), - uri=f"https://{end_point}/$discovery/rest?version=v1beta&key={api_key}", + uri=f"{GENAI_API_DISCOVERY_URL}?version=v1beta&key={api_key}", ) response, content = request.execute() @@ -84,7 +82,7 @@ def create_file( class FileServiceAsyncClient(glm.FileServiceAsyncClient): async def create_file(self, *args, **kwargs): - raise NotImplementedError("Create_file is not yet implemented for the async client.") + raise NotImplementedError("create_file is not yet implemented for the async client.") @dataclasses.dataclass diff --git a/google/generativeai/files.py b/google/generativeai/files.py index b3a8a3e42..330a5b8cc 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -18,12 +18,15 @@ import pathlib import mimetypes from typing import Iterable +import logging +import google.ai.generativelanguage as glm +from itertools import islice from google.generativeai.types import file_types from google.generativeai.client import get_default_file_client -__all__ = ["upload_file", "get_file", "list_files"] +__all__ = ["upload_file", "get_file", "list_files", "delete_file"] def upload_file( @@ -52,10 +55,10 @@ def upload_file( return file_types.File(response) -def list_files(page_size=50) -> Iterable[file_types.File]: +def list_files(max_results=50) -> Iterable[file_types.File]: client = get_default_file_client() - response = client.list_files(page_size=page_size) + response = client.list_files(glm.ListFilesRequest()) for proto in response: yield file_types.File(proto) @@ -63,3 +66,9 @@ def list_files(page_size=50) -> Iterable[file_types.File]: def get_file(name) -> file_types.File: client = get_default_file_client() return file_types.File(client.get_file(name=name)) + + +def delete_file(name): + request = glm.DeleteFileRequest(name=name) + client = get_default_file_client() + client.delete_file(request=request) diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py index 135a34ce1..ceb08ef07 100644 --- a/google/generativeai/generative_models.py +++ b/google/generativeai/generative_models.py @@ -239,8 +239,8 @@ def generate_content( except google.api_core.exceptions.InvalidArgument as e: if e.message.startswith("Request payload size exceeds the limit:"): e.message += ( - " Please upload your files with the Files api instead." - "`f = genai.create_file(path); m.generate_content(['tell me about this file:', f])`" + " Please upload your files with the File API instead." + "`f = genai.upload_file(path); m.generate_content(['tell me about this file:', f])`" ) raise @@ -284,8 +284,8 @@ async def generate_content_async( except google.api_core.exceptions.InvalidArgument as e: if e.message.startswith("Request payload size exceeds the limit:"): e.message += ( - " Please upload your files with the Files api instead." - "`f = genai.create_file(path); m.generate_content(['tell me about this file:', f])`" + " Please upload your files with the File API instead." + "`f = genai.upload_file(path); m.generate_content(['tell me about this file:', f])`" ) raise diff --git a/google/generativeai/types/file_types.py b/google/generativeai/types/file_types.py index e46288b1c..f403f5ae2 100644 --- a/google/generativeai/types/file_types.py +++ b/google/generativeai/types/file_types.py @@ -67,6 +67,10 @@ def update_time(self) -> datetime.datetime: def sha256_hash(self) -> bytes: return self._proto.sha256_hash + @property + def uri(self) -> str: + return self._proto.uri + def delete(self): client = get_default_file_client() client.delete_file(name=self.name) From 325a8d38568ae3e966959e159c92ce980b984c55 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 28 Mar 2024 09:20:03 -0700 Subject: [PATCH 2/5] Apply suggestions from code review --- google/generativeai/client.py | 2 +- google/generativeai/files.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index 4415ec8f2..5072f5ea6 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -82,7 +82,7 @@ def create_file( class FileServiceAsyncClient(glm.FileServiceAsyncClient): async def create_file(self, *args, **kwargs): - raise NotImplementedError("create_file is not yet implemented for the async client.") + raise NotImplementedError("`create_file` is not yet implemented for the async client.") @dataclasses.dataclass diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 330a5b8cc..9d892ddcb 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -55,10 +55,10 @@ def upload_file( return file_types.File(response) -def list_files(max_results=50) -> Iterable[file_types.File]: +def list_files(page_size=50) -> Iterable[file_types.File]: client = get_default_file_client() - response = client.list_files(glm.ListFilesRequest()) + response = client.list_files(glm.ListFilesRequest(page_size=page_size)) for proto in response: yield file_types.File(proto) @@ -69,6 +69,8 @@ def get_file(name) -> file_types.File: def delete_file(name): + if isinstance(name, (file_types.File, glm.File)): + name = name.name request = glm.DeleteFileRequest(name=name) client = get_default_file_client() client.delete_file(request=request) From b698f442d9d5f1abdf45615252c4fb007627c626 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 28 Mar 2024 09:47:14 -0700 Subject: [PATCH 3/5] format --- google/generativeai/client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/google/generativeai/client.py b/google/generativeai/client.py index 5072f5ea6..80e3c3b76 100644 --- a/google/generativeai/client.py +++ b/google/generativeai/client.py @@ -27,7 +27,8 @@ __version__ = "0.0.0" USER_AGENT = "genai-py" -GENAI_API_DISCOVERY_URL="https://generativelanguage.googleapis.com/$discovery/rest" +GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest" + class FileServiceClient(glm.FileServiceClient): def __init__(self, *args, **kwargs): From e80512e1f224d3254ecdfbe80589da608878d2a0 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 28 Mar 2024 10:07:04 -0700 Subject: [PATCH 4/5] Increase default page_size --- google/generativeai/files.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 9d892ddcb..ffd9e6ed8 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -55,7 +55,7 @@ def upload_file( return file_types.File(response) -def list_files(page_size=50) -> Iterable[file_types.File]: +def list_files(page_size=100) -> Iterable[file_types.File]: client = get_default_file_client() response = client.list_files(glm.ListFilesRequest(page_size=page_size)) From 1ca35d34b411418f09a5afb7576c912b4f8567c1 Mon Sep 17 00:00:00 2001 From: Michael Chen Date: Mon, 8 Apr 2024 10:57:27 -0400 Subject: [PATCH 5/5] add default timeout --- google/generativeai/generative_models.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py index ceb08ef07..d3091931b 100644 --- a/google/generativeai/generative_models.py +++ b/google/generativeai/generative_models.py @@ -85,6 +85,9 @@ def __init__( self._client = None self._async_client = None + self._default_request_options = { + "timeout": 600 + } @property def model_name(self): @@ -220,7 +223,7 @@ def generate_content( self._client = client.get_default_generative_client() if request_options is None: - request_options = {} + request_options = self._default_request_options try: if stream: @@ -265,7 +268,7 @@ async def generate_content_async( self._async_client = client.get_default_generative_async_client() if request_options is None: - request_options = {} + request_options = self._default_request_options try: if stream: