google-gemini · markmcd · May 17, 2024 · Apr 22, 2024 · Apr 26, 2024 · Apr 26, 2024
diff --git a/google/generativeai/answer.py b/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
 
@@ -247,7 +245,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.
@@ -320,7 +318,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.

diff --git a/google/generativeai/discuss.py b/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 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)

diff --git a/google/generativeai/embedding.py b/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)

diff --git a/google/generativeai/generative_models.py b/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
 
 
@@ -181,7 +181,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.
 
@@ -281,7 +281,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(
@@ -328,7 +328,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 = {}
@@ -355,7 +355,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 = {}

diff --git a/google/generativeai/models.py b/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` or `types.TunedModel` object.
 
@@ -62,7 +63,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.
 
@@ -99,7 +100,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.
 
@@ -162,7 +163,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.
 
@@ -196,7 +197,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.
 
@@ -244,7 +245,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.
 
@@ -344,6 +345,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
     )
@@ -357,7 +359,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
 
@@ -368,7 +370,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
 
@@ -378,7 +380,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:
@@ -395,6 +397,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)
@@ -436,7 +439,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 = {}