Make model name handling more consistent (and fix some pytype errors). (#36)

* Make model name handling more consistent.

* fix test

Author: MarkDaoust

google/generativeai/discuss.py

@@ -501,9 +501,10 @@ def count_message_tokens(
     context: Optional[str] = None,
     examples: Optional[discuss_types.ExamplesOptions] = None,
     messages: Optional[discuss_types.MessagesOptions] = None,
-    model: str = DEFAULT_DISCUSS_MODEL,
+    model: model_types.ModelNameOptions = DEFAULT_DISCUSS_MODEL,
     client: Optional[glm.DiscussServiceAsyncClient] = None,
 ):
+    model = model_types.make_model_name(model)
     prompt = _make_message_prompt(
         prompt, context=context, examples=examples, messages=messages
     )

google/generativeai/models.py

@@ -19,18 +19,13 @@
 from google.generativeai.client import get_default_model_client
 from google.generativeai.types import model_types
 
-# A bare model name, with no preceding namespace. e.g. foo-bar-001
-_BARE_MODEL_NAME = re.compile(r"^\w+-\w+-\d+$")
-
 
 def get_model(name: str, *, client=None) -> model_types.Model:
     """Get the `types.Model` for the given model name."""
     if client is None:
         client = get_default_model_client()
 
-    # If only a bare model name is passed, give it the structure we expect.
-    if _BARE_MODEL_NAME.match(name):
-        name = f"models/{name}"
+    name = model_types.make_model_name(name)
 
     result = client.get_model(name=name)
     result = type(result).to_dict(result)

google/generativeai/text.py

@@ -26,6 +26,8 @@
 from google.generativeai.types import model_types
 from google.generativeai.types import safety_types
 
+DEFAULT_TEXT_MODEL = "models/text-bison-001"
+
 
 def _make_text_prompt(prompt: Union[str, dict[str, str]]) -> glm.TextPrompt:
     if isinstance(prompt, str):
@@ -38,15 +40,15 @@ def _make_text_prompt(prompt: Union[str, dict[str, str]]) -> glm.TextPrompt:
 
 def _make_generate_text_request(
     *,
-    model: model_types.ModelNameOptions = "models/chat-lamda-001",
+    model: model_types.ModelNameOptions = DEFAULT_TEXT_MODEL,
     prompt: Optional[str] = None,
     temperature: Optional[float] = None,
     candidate_count: Optional[int] = None,
     max_output_tokens: Optional[int] = None,
     top_p: Optional[int] = None,
     top_k: Optional[int] = None,
     safety_settings: Optional[List[safety_types.SafetySettingDict]] = None,
-    stop_sequences: Union[str, Iterable[str]] = None,
+    stop_sequences: Optional[Union[str, Iterable[str]]] = None,
 ) -> glm.GenerateTextRequest:
     model = model_types.make_model_name(model)
     prompt = _make_text_prompt(prompt=prompt)
@@ -70,15 +72,15 @@ def _make_generate_text_request(
 
 def generate_text(
     *,
-    model: Optional[model_types.ModelNameOptions] = "models/text-bison-001",
+    model: model_types.ModelNameOptions = DEFAULT_TEXT_MODEL,
     prompt: str,
     temperature: Optional[float] = None,
     candidate_count: Optional[int] = None,
     max_output_tokens: Optional[int] = None,
     top_p: Optional[float] = None,
     top_k: Optional[float] = None,
-    safety_settings: Optional[Iterable[safety.SafetySettingDict]] = None,
-    stop_sequences: Union[str, Iterable[str]] = None,
+    safety_settings: Optional[Iterable[safety_types.SafetySettingDict]] = None,
+    stop_sequences: Optional[Union[str, Iterable[str]]] = None,
     client: Optional[glm.TextServiceClient] = None,
 ) -> text_types.Completion:
     """Calls the API and returns a `types.Completion` containing the response.
@@ -170,7 +172,9 @@ def _generate_response(
     return Completion(_client=client, **response)
 
 
-def generate_embeddings(model: str, text: str, client: glm.TextServiceClient = None):
+def generate_embeddings(
+    model: model_types.ModelNameOptions, text: str, client: glm.TextServiceClient = None
+):
     """Calls the API to create an embedding for the text passed in.
 
     Args:

google/generativeai/types/model_types.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """Type definitions for the models service."""
 
+import re
 import abc
 import dataclasses
 from typing import Iterator, List, Optional, Union
@@ -58,10 +59,18 @@ class Model:
 
 ModelNameOptions = Union[str, Model]
 
+# A bare model name, with no preceding namespace. e.g. foo-bar-001
+_BARE_MODEL_NAME = re.compile(r"^\w+-\w+-\d+$")
+
 
 def make_model_name(name: ModelNameOptions):
     if isinstance(name, Model):
         name = name.name
+    elif isinstance(name, str):
+        # If only a bare model name is passed, give it the structure we expect.
+        if _BARE_MODEL_NAME.match(name):
+            name = f"models/{name}"
+
     return name
 
 

tests/test_text.py

@@ -74,8 +74,10 @@ def test_make_prompt(self, prompt):
         ]
     )
     def test_make_generate_text_request(self, prompt):
-        x = text_service._make_generate_text_request(prompt=prompt)
-        self.assertEqual("models/chat-lamda-001", x.model)
+        x = text_service._make_generate_text_request(
+            model="chat-bison-001", prompt=prompt
+        )
+        self.assertEqual("models/chat-bison-001", x.model)
         self.assertIsInstance(x, glm.GenerateTextRequest)
 
     @parameterized.named_parameters(