Add ability to configure timeouts (#204)

* Add ability to configure timeouts

* Add timeout to docstrings

* Replace timeout default w/ None

* Add missing timeout arg to mocked methed

* Make async/sync generate_content have same formatting

* Add missing timeout arg to embedding mocks

* Add request_options param

* Test generate answer called w/ req opts

* Test embed_content called w/ req opts

* Test count_tokens called w/ req opts

* Test async generate_content called w/ req opts

* Test embed/generate text called w/ req opts

* Fixup mocked methods w/ missing req opts

* Test batch_embed_text called w/ req opts

* Test count_text_tokens called w/ req opts

* Test create_corpus called w/ req opts

* Test corpus called w/ req opts

* Unpack req opts, hande None case

* Handle exc instead of assert raises

* Handle failing tests

* Test batch_embed_contents called w/ req opts

* Test get_model called w/ req opts

* Test get_tuned_model called w/ req opts

* Test list_models called w/ req opts

* Test list_tuned_models called w/ req opts

* Test update_tuned_model called w/ req opts

* Test delete_tuned_model called w/ req opts

* Test create_tuned_model called w/ req opts

* Add missing kwarg to mocked client method

* Pass req opts to update tuned model client method

* Test update_corpus called w/ req opts

* Test list_corpora called w/ req opts

* Test query_corpus called w/ req opts

* Test delete_corpus called w/ req opts

* Rm unneeded try catch on delete corpus

* Test create_document called w/ req opts

* Import missing Any

* Test embed_content called w/ req opts

* Test sync count_tokens called w/ req opts

* Test create_corpus called w/ req opts

* Test async embed_content called w/ req opts

* Test generate_message called w/ req opts

* Test get_corpus called w/ req opts

* Test update_corpus called w/ req opts

* Test get_document called w/ req opts

* Test update_document called w/ req opts

* Test list_documents called w/ req opts

* Test delete_document called w/ req opts

* Test query_document called w/ req opts

* Test create_chunk called w/ req opts

* Test batch_create_chunks called w/ req opts

* Test get_chunk called w/ req opts

* Test list_chunks called w/ req opts

* Test update_chunk called w/ req opts

* Test batch_update_chunks called w/ req opts

* Test delete_chunk called w/ req opts

* Test batch_delete_chunks called w/ req opts

* Add req opts to docstrings

* Add couple more docstrs 4 req opts

* Remove stray timeout.

---------

Co-authored-by: Mark Daoust <[email protected]>

Author: aidoskanapyanov

google/generativeai/answer.py

@@ -165,6 +165,7 @@ def generate_answer(
     safety_settings: safety_types.SafetySettingOptions | None = None,
     temperature: float | None = None,
     client: glm.GenerativeServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ):
     """
     Calls the API and returns a `types.Answer` containing the answer.
@@ -177,10 +178,14 @@ def generate_answer(
         answer_style: Style in which the grounded answer should be returned.
         safety_settings: Safety settings for generated output. Defaults to None.
         client: If you're not relying on a default client, you pass a `glm.TextServiceClient` instead.
+        request_options: Options for the request.
 
     Returns:
         A `types.Answer` containing the model's text answer response.
     """
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_generative_client()
 
@@ -193,6 +198,6 @@ def generate_answer(
         answer_style=answer_style,
     )
 
-    response = client.generate_answer(request)
+    response = client.generate_answer(request, **request_options)
 
     return response

google/generativeai/discuss.py

@@ -18,7 +18,7 @@
 import sys
 import textwrap
 
-from typing import Iterable, List, Optional, Union
+from typing import Any, Iterable, List, Optional, Union
 
 import google.ai.generativelanguage as glm
 
@@ -316,6 +316,7 @@ def chat(
     top_k: float | None = None,
     prompt: discuss_types.MessagePromptOptions | None = None,
     client: glm.DiscussServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> discuss_types.ChatResponse:
     """Calls the API and returns a `types.ChatResponse` containing the response.
 
@@ -382,6 +383,7 @@ def chat(
             setting `context`/`examples`/`messages`, but not both.
         client: If you're not relying on the default client, you pass a
             `glm.DiscussServiceClient` instead.
+        request_options: Options for the request.
 
     Returns:
         A `types.ChatResponse` containing the model's reply.
@@ -398,7 +400,7 @@ def chat(
         prompt=prompt,
     )
 
-    return _generate_response(client=client, request=request)
+    return _generate_response(client=client, request=request, request_options=request_options)
 
 
 @string_utils.set_doc(chat.__doc__)
@@ -414,6 +416,7 @@ async def chat_async(
     top_k: float | None = None,
     prompt: discuss_types.MessagePromptOptions | None = None,
     client: glm.DiscussServiceAsyncClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> discuss_types.ChatResponse:
     request = _make_generate_message_request(
         model=model,
@@ -427,7 +430,9 @@ async def chat_async(
         prompt=prompt,
     )
 
-    return await _generate_response_async(client=client, request=request)
+    return await _generate_response_async(
+        client=client, request=request, request_options=request_options
+    )
 
 
 if (sys.version_info.major, sys.version_info.minor) >= (3, 10):
@@ -461,7 +466,11 @@ def last(self, message: discuss_types.MessageOptions):
         self.messages[-1] = message
 
     @string_utils.set_doc(discuss_types.ChatResponse.reply.__doc__)
-    def reply(self, message: discuss_types.MessageOptions) -> discuss_types.ChatResponse:
+    def reply(
+        self,
+        message: discuss_types.MessageOptions,
+        request_options: dict[str, Any] | None = None,
+    ) -> discuss_types.ChatResponse:
         if isinstance(self._client, glm.DiscussServiceAsyncClient):
             raise TypeError(f"reply can't be called on an async client, use reply_async instead.")
         if self.last is None:
@@ -477,7 +486,9 @@ def reply(self, message: discuss_types.MessageOptions) -> discuss_types.ChatResp
         request["messages"] = list(request["messages"])
         request["messages"].append(_make_message(message))
         request = _make_generate_message_request(**request)
-        return _generate_response(request=request, client=self._client)
+        return _generate_response(
+            request=request, client=self._client, request_options=request_options
+        )
 
     @string_utils.set_doc(discuss_types.ChatResponse.reply.__doc__)
     async def reply_async(
@@ -526,23 +537,31 @@ def _build_chat_response(
 def _generate_response(
     request: glm.GenerateMessageRequest,
     client: glm.DiscussServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> ChatResponse:
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_discuss_client()
 
-    response = client.generate_message(request)
+    response = client.generate_message(request, **request_options)
 
     return _build_chat_response(request, response, client)
 
 
 async def _generate_response_async(
     request: glm.GenerateMessageRequest,
     client: glm.DiscussServiceAsyncClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> ChatResponse:
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_discuss_async_client()
 
-    response = await client.generate_message(request)
+    response = await client.generate_message(request, **request_options)
 
     return _build_chat_response(request, response, client)
 
@@ -555,13 +574,17 @@ def count_message_tokens(
     messages: discuss_types.MessagesOptions | None = None,
     model: model_types.AnyModelNameOptions = DEFAULT_DISCUSS_MODEL,
     client: glm.DiscussServiceAsyncClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> discuss_types.TokenCount:
     model = model_types.make_model_name(model)
     prompt = _make_message_prompt(prompt, context=context, examples=examples, messages=messages)
 
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_discuss_client()
 
-    result = client.count_message_tokens(model=model, prompt=prompt)
+    result = client.count_message_tokens(model=model, prompt=prompt, **request_options)
 
     return type(result).to_dict(result)

google/generativeai/embedding.py

@@ -17,7 +17,7 @@
 import dataclasses
 from collections.abc import Iterable, Sequence, Mapping
 import itertools
-from typing import Iterable, overload, TypeVar, Union, Mapping
+from typing import Any, Iterable, overload, TypeVar, Union, Mapping
 
 import google.ai.generativelanguage as glm
 
@@ -95,6 +95,7 @@ def embed_content(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.EmbeddingDict: ...
 
 
@@ -105,6 +106,7 @@ def embed_content(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.BatchEmbeddingDict: ...
 
 
@@ -114,6 +116,7 @@ def embed_content(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceClient = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.EmbeddingDict | text_types.BatchEmbeddingDict:
     """Calls the API to create embeddings for content passed in.
 
@@ -132,13 +135,18 @@ def embed_content(
         title:
             An optional title for the text. Only applicable when task_type is
             `RETRIEVAL_DOCUMENT`.
+        request_options:
+            Options for the request.
 
     Return:
         Dictionary containing the embedding (list of float values) for the
         input content.
     """
     model = model_types.make_model_name(model)
 
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_generative_client()
 
@@ -160,15 +168,21 @@ def embed_content(
         )
         for batch in _batched(requests, EMBEDDING_MAX_BATCH_SIZE):
             embedding_request = glm.BatchEmbedContentsRequest(model=model, requests=batch)
-            embedding_response = client.batch_embed_contents(embedding_request)
+            embedding_response = client.batch_embed_contents(
+                embedding_request,
+                **request_options,
+            )
             embedding_dict = type(embedding_response).to_dict(embedding_response)
             result["embedding"].extend(e["values"] for e in embedding_dict["embeddings"])
         return result
     else:
         embedding_request = glm.EmbedContentRequest(
             model=model, content=content_types.to_content(content), task_type=task_type, title=title
         )
-        embedding_response = client.embed_content(embedding_request)
+        embedding_response = client.embed_content(
+            embedding_request,
+            **request_options,
+        )
         embedding_dict = type(embedding_response).to_dict(embedding_response)
         embedding_dict["embedding"] = embedding_dict["embedding"]["values"]
         return embedding_dict
@@ -181,6 +195,7 @@ async def embed_content_async(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceAsyncClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.EmbeddingDict: ...
 
 
@@ -191,6 +206,7 @@ async def embed_content_async(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceAsyncClient | None = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.BatchEmbeddingDict: ...
 
 
@@ -200,10 +216,14 @@ async def embed_content_async(
     task_type: EmbeddingTaskTypeOptions | None = None,
     title: str | None = None,
     client: glm.GenerativeServiceAsyncClient = None,
+    request_options: dict[str, Any] | None = None,
 ) -> text_types.EmbeddingDict | text_types.BatchEmbeddingDict:
     """The async version of `genai.embed_content`."""
     model = model_types.make_model_name(model)
 
+    if request_options is None:
+        request_options = {}
+
     if client is None:
         client = get_default_generative_async_client()
 
@@ -225,15 +245,21 @@ async def embed_content_async(
         )
         for batch in _batched(requests, EMBEDDING_MAX_BATCH_SIZE):
             embedding_request = glm.BatchEmbedContentsRequest(model=model, requests=batch)
-            embedding_response = await client.batch_embed_contents(embedding_request)
+            embedding_response = await client.batch_embed_contents(
+                embedding_request,
+                **request_options,
+            )
             embedding_dict = type(embedding_response).to_dict(embedding_response)
             result["embedding"].extend(e["values"] for e in embedding_dict["embeddings"])
         return result
     else:
         embedding_request = glm.EmbedContentRequest(
             model=model, content=content_types.to_content(content), task_type=task_type, title=title
         )
-        embedding_response = await client.embed_content(embedding_request)
+        embedding_response = await client.embed_content(
+            embedding_request,
+            **request_options,
+        )
         embedding_dict = type(embedding_response).to_dict(embedding_response)
         embedding_dict["embedding"] = embedding_dict["embedding"]["values"]
         return embedding_dict