Add additional functionalities with files SDK (google-gemini#260)

TYMichaelChen · MarkDaoust · web-flow · commit a187374fda7f · 2024-03-28T10:27:17.000-07:00
* Add additional functionalities with files SDK

* Apply suggestions from code review

* format

* Increase default page_size

---------

Co-authored-by: Mark Daoust &lt;markdaoust@google.com&gt;
diff --git 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
@@ -27,6 +27,7 @@
     __version__ = "0.0.0"
 
 USER_AGENT = "genai-py"
+GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest"
 
 
 class FileServiceClient(glm.FileServiceClient):
@@ -37,14 +38,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 +83,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
@@ -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,14 +55,22 @@ 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(page_size=page_size)
+    response = client.list_files(glm.ListFilesRequest(page_size=page_size))
     for proto in response:
         yield file_types.File(proto)
 
 
 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):
+    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)
diff --git 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
@@ -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)
