Skip to content

add default timeout #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions google/generativeai/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 4 additions & 5 deletions google/generativeai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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()

Expand Down Expand Up @@ -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
Expand Down
17 changes: 14 additions & 3 deletions google/generativeai/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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)
15 changes: 9 additions & 6 deletions google/generativeai/generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,9 @@ def __init__(

self._client = None
self._async_client = None
self._default_request_options = {
"timeout": 600
}

@property
def model_name(self):
Expand Down Expand Up @@ -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:
Expand All @@ -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

Expand All @@ -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:
Expand All @@ -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

Expand Down
4 changes: 4 additions & 0 deletions google/generativeai/types/file_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)