Support bare model names in palm.get_model (#35)

Right now we require model names to be specified as
`"models/type-size-version"` (with the `models/` prefix), but that seems
slightly unintuitive. This change adds the `models/` prefix if it
detects the expected *bare* model string (e.g. `foo-bar-001`).

Adds `test_models.py` too.

Author: markmcd

google/generativeai/models.py

@@ -12,19 +12,26 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import dataclasses
+import re
 from typing import Optional, List
 
 import google.ai.generativelanguage as glm
 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, *, client=None) -> model_types.Model:
+
+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}"
+
     result = client.get_model(name=name)
     result = type(result).to_dict(result)
     return model_types.Model(**result)
@@ -37,7 +44,7 @@ def __init__(
         page_size: int,
         page_token: Optional[str],
         models: List[model_types.Model],
-        client: Optional[glm.ModelServiceClient]
+        client: Optional[glm.ModelServiceClient],
     ):
         self._page_size = page_size
         self._page_token = page_token

tests/test_models.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from unittest import mock
+
+from absl.testing import absltest
+
+import google.ai.generativelanguage as glm
+from google.ai.generativelanguage_v1beta2.types import model
+
+from google.generativeai import models
+
+_FAKE_MODEL = model.Model(
+    name="models/fake-model-001",
+    base_model_id="",
+    version="001",
+    display_name="Fake Model",
+    description="A fake model",
+    input_token_limit=123,
+    output_token_limit=234,
+    supported_generation_methods=[],
+)
+
+
+class UnitTests(absltest.TestCase):
+    def test_model_prefix(self):
+        """Test `models/` prefix applies to get_model calls when necessary."""
+        # The SUT needs a concrete return type from `get_model`, so set up a real-enough client.
+        fake_client = mock.Mock(spec=glm.ModelServiceClient)
+        fake_client.get_model.return_value = _FAKE_MODEL
+
+        # Ensure that we don't mess with correctly structure args.
+        models.get_model(name="models/text-bison-001", client=fake_client)
+        fake_client.get_model.assert_called_with(name="models/text-bison-001")
+
+        # Ensure that we do correct bare models.
+        models.get_model(name="text-bison-001", client=fake_client)
+        fake_client.get_model.assert_called_with(name="models/text-bison-001")
+
+        # And unknown structure is not touched.
+        models.get_model(name="unknown_string", client=fake_client)
+        fake_client.get_model.assert_called_with(name="unknown_string")