Improve request_options (#297)

* Working on request_options

* Add helper_types

Change-Id: Idc3e813616413f4ce085c05b771c0127e4dfc886

* format

Change-Id: I186e015de97ceece56ee5a97f6edef47ef223d18

* UpdateRequestOptions

Change-Id: I9f92466967fb1aa605d442cb143699da4308409b

* Add docs

Change-Id: I209b2b2ad8d783001b1828cbcac84ca301c11bec

* work

Change-Id: I00a2e2edb1e9bf3d4f51c0a868a34e044be3c6ff

* Fix Py3.9

Change-Id: I8cf0ccac90ba3c4548e7549fec7d0b9b58925e7e

* use RequestOptions in tests

Change-Id: I92b68bc86330ad874c3765f428a2e64ba220750f

* annotations

Change-Id: Idbc428075729255d66d2ba8b3bcce0a1d6e8f048

* Update tests/test_discuss.py

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

* tests

Change-Id: Ife30e2cc47bd4c52d2dddafdd85a51df0e42e160

---------

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

Author: MarkDaoust

google/generativeai/answer.py

@@ -26,12 +26,10 @@
     get_default_generative_client,
     get_default_generative_async_client,
 )
-from google.generativeai import string_utils
 from google.generativeai.types import model_types
-from google.generativeai import models
+from google.generativeai.types import helper_types
 from google.generativeai.types import safety_types
 from google.generativeai.types import content_types
-from google.generativeai.types import answer_types
 from google.generativeai.types import retriever_types
 from google.generativeai.types.retriever_types import MetadataFilter
 
@@ -245,7 +243,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,
+    request_options: helper_types.RequestOptionsType | None = None,
 ):
     """
     Calls the GenerateAnswer API and returns a `types.Answer` containing the response.
@@ -318,7 +316,7 @@ async def generate_answer_async(
     safety_settings: safety_types.SafetySettingOptions | None = None,
     temperature: float | None = None,
     client: glm.GenerativeServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ):
     """
     Calls the API and returns a `types.Answer` containing the answer.

google/generativeai/discuss.py

@@ -26,6 +26,7 @@
 from google.generativeai.client import get_default_discuss_async_client
 from google.generativeai import string_utils
 from google.generativeai.types import discuss_types
+from google.generativeai.types import helper_types
 from google.generativeai.types import model_types
 from google.generativeai.types import palm_safety_types
 
@@ -316,7 +317,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,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> discuss_types.ChatResponse:
     """Calls the API and returns a `types.ChatResponse` containing the response.
 
@@ -416,7 +417,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,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> discuss_types.ChatResponse:
     request = _make_generate_message_request(
         model=model,
@@ -469,7 +470,7 @@ def last(self, message: discuss_types.MessageOptions):
     def reply(
         self,
         message: discuss_types.MessageOptions,
-        request_options: dict[str, Any] | None = None,
+        request_options: helper_types.RequestOptionsType | 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.")
@@ -537,7 +538,7 @@ def _build_chat_response(
 def _generate_response(
     request: glm.GenerateMessageRequest,
     client: glm.DiscussServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> ChatResponse:
     if request_options is None:
         request_options = {}
@@ -553,7 +554,7 @@ def _generate_response(
 async def _generate_response_async(
     request: glm.GenerateMessageRequest,
     client: glm.DiscussServiceAsyncClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> ChatResponse:
     if request_options is None:
         request_options = {}
@@ -574,7 +575,7 @@ 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,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> discuss_types.TokenCount:
     model = model_types.make_model_name(model)
     prompt = _make_message_prompt(prompt, context=context, examples=examples, messages=messages)

google/generativeai/embedding.py

@@ -14,8 +14,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import dataclasses
-from collections.abc import Iterable, Sequence, Mapping
 import itertools
 from typing import Any, Iterable, overload, TypeVar, Union, Mapping
 
@@ -24,7 +22,7 @@
 from google.generativeai.client import get_default_generative_client
 from google.generativeai.client import get_default_generative_async_client
 
-from google.generativeai import string_utils
+from google.generativeai.types import helper_types
 from google.generativeai.types import text_types
 from google.generativeai.types import model_types
 from google.generativeai.types import content_types
@@ -104,7 +102,7 @@ def embed_content(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.EmbeddingDict: ...
 
 
@@ -116,7 +114,7 @@ def embed_content(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.BatchEmbeddingDict: ...
 
 
@@ -127,7 +125,7 @@ def embed_content(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceClient = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.EmbeddingDict | text_types.BatchEmbeddingDict:
     """Calls the API to create embeddings for content passed in.
 
@@ -224,7 +222,7 @@ async def embed_content_async(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceAsyncClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.EmbeddingDict: ...
 
 
@@ -236,7 +234,7 @@ async def embed_content_async(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceAsyncClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.BatchEmbeddingDict: ...
 
 
@@ -247,7 +245,7 @@ async def embed_content_async(
     title: str | None = None,
     output_dimensionality: int | None = None,
     client: glm.GenerativeServiceAsyncClient = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> text_types.EmbeddingDict | text_types.BatchEmbeddingDict:
     """The async version of `genai.embed_content`."""
     model = model_types.make_model_name(model)

google/generativeai/generative_models.py

@@ -15,9 +15,9 @@
 import google.api_core.exceptions
 from google.ai import generativelanguage as glm
 from google.generativeai import client
-from google.generativeai import string_utils
 from google.generativeai.types import content_types
 from google.generativeai.types import generation_types
+from google.generativeai.types import helper_types
 from google.generativeai.types import safety_types
 
 
@@ -179,7 +179,7 @@ def generate_content(
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
-        request_options: dict[str, Any] | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> generation_types.GenerateContentResponse:
         """A multipurpose function to generate responses from the model.
 
@@ -279,7 +279,7 @@ async def generate_content_async(
         stream: bool = False,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
-        request_options: dict[str, Any] | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> generation_types.AsyncGenerateContentResponse:
         """The async version of `GenerativeModel.generate_content`."""
         request = self._prepare_request(
@@ -326,7 +326,7 @@ def count_tokens(
         safety_settings: safety_types.SafetySettingOptions | None = None,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
-        request_options: dict[str, Any] | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> glm.CountTokensResponse:
         if request_options is None:
             request_options = {}
@@ -353,7 +353,7 @@ async def count_tokens_async(
         safety_settings: safety_types.SafetySettingOptions | None = None,
         tools: content_types.FunctionLibraryType | None = None,
         tool_config: content_types.ToolConfigType | None = None,
-        request_options: dict[str, Any] | None = None,
+        request_options: helper_types.RequestOptionsType | None = None,
     ) -> glm.CountTokensResponse:
         if request_options is None:
             request_options = {}

google/generativeai/models.py

@@ -21,6 +21,7 @@
 from google.generativeai import operations
 from google.generativeai.client import get_default_model_client
 from google.generativeai.types import model_types
+from google.generativeai.types import helper_types
 from google.api_core import operation
 from google.api_core import protobuf_helpers
 from google.protobuf import field_mask_pb2
@@ -31,7 +32,7 @@ def get_model(
     name: model_types.AnyModelNameOptions,
     *,
     client=None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.Model | model_types.TunedModel:
     """Given a model name, fetch the `types.Model`
 
@@ -64,7 +65,7 @@ def get_base_model(
     name: model_types.BaseModelNameOptions,
     *,
     client=None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.Model:
     """Get the `types.Model` for the given base model name.
 
@@ -101,7 +102,7 @@ def get_tuned_model(
     name: model_types.TunedModelNameOptions,
     *,
     client=None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.TunedModel:
     """Get the `types.TunedModel` for the given tuned model name.
 
@@ -164,7 +165,7 @@ def list_models(
     *,
     page_size: int | None = 50,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.ModelsIterable:
     """Lists available models.
 
@@ -198,7 +199,7 @@ def list_tuned_models(
     *,
     page_size: int | None = 50,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.TunedModelsIterable:
     """Lists available models.
 
@@ -246,7 +247,7 @@ def create_tuned_model(
     input_key: str = "text_input",
     output_key: str = "output",
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> operations.CreateTunedModelOperation:
     """Launches a tuning job to create a TunedModel.
 
@@ -346,6 +347,7 @@ def create_tuned_model(
         top_k=top_k,
         tuning_task=tuning_task,
     )
+
     operation = client.create_tuned_model(
         dict(tuned_model_id=id, tuned_model=tuned_model), **request_options
     )
@@ -359,7 +361,7 @@ def update_tuned_model(
     updates: None = None,
     *,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.TunedModel:
     pass
 
@@ -370,7 +372,7 @@ def update_tuned_model(
     updates: dict[str, Any],
     *,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.TunedModel:
     pass
 
@@ -380,7 +382,7 @@ def update_tuned_model(
     updates: dict[str, Any] | None = None,
     *,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> model_types.TunedModel:
     """Push updates to the tuned model. Only certain attributes are updatable."""
     if request_options is None:
@@ -397,6 +399,7 @@ def update_tuned_model(
                 "`updates` must be a `dict`.\n"
                 f"got: {type(updates)}"
             )
+
         tuned_model = client.get_tuned_model(name=name, **request_options)
 
         updates = flatten_update_paths(updates)
@@ -438,7 +441,7 @@ def _apply_update(thing, path, value):
 def delete_tuned_model(
     tuned_model: model_types.TunedModelNameOptions,
     client: glm.ModelServiceClient | None = None,
-    request_options: dict[str, Any] | None = None,
+    request_options: helper_types.RequestOptionsType | None = None,
 ) -> None:
     if request_options is None:
         request_options = {}