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..80e3c3b76 100644 --- 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 index b3a8a3e42..ffd9e6ed8 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(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) @@ -63,3 +66,11 @@ 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): + 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 index 135a34ce1..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: @@ -239,8 +242,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 @@ -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: @@ -284,8 +287,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)