Simplify list-models iteration - the GAPIC lib handles this. (#86)

* Simplify list-models iteration - the GAPIC lib handles this.

* resolve comments

Author: MarkDaoust

google/generativeai/models.py

@@ -137,56 +137,9 @@ def get_base_model_name(
     return base_model
 
 
-def _list_base_models_next_page(page_size, page_token, client):
-    """Returns the next page of the base model or tuned model list."""
-    result = client.list_models(page_size=page_size, page_token=page_token)
-    result = result._response
-    result = type(result).to_dict(result)
-    result["models"] = [model_types.Model(**mod) for mod in result["models"]]
-    result["page_size"] = page_size
-    result["page_token"] = result.pop("next_page_token")
-    result["client"] = client
-    return result
-
-
-def _list_tuned_models_next_page(page_size, page_token, client):
-    """Returns the next page of the base model or tuned model list."""
-    result = client.list_tuned_models(
-        glm.ListTunedModelsRequest(page_size=page_size, page_token=page_token)
-    )
-    result = result._response
-    result = type(result).to_dict(result)
-    result["models"] = [model_types.decode_tuned_model(mod) for mod in result.pop("tuned_models")]
-    result["page_size"] = page_size
-    result["page_token"] = result.pop("next_page_token")
-    result["client"] = client
-    return result
-
-
-def _list_models_iter_pages(
-    *,
-    page_size: int | None = None,
-    select: Literal["base", "tuned"],
-    client: glm.ModelServiceClient | None = None,
-):
-    if client is None:
-        client = get_default_model_client()
-
-    page_token = None
-    while True:
-        if select == "base":
-            result = _list_base_models_next_page(page_size, page_token=page_token, client=client)
-        elif select == "tuned":
-            result = _list_tuned_models_next_page(page_size, page_token=page_token, client=client)
-        yield from result["models"]
-        page_token = result["page_token"]
-        if page_token == "":
-            break
-
-
 def list_models(
     *,
-    page_size: int | None = None,
+    page_size: int | None = 50,
     client: glm.ModelServiceClient | None = None,
 ) -> model_types.ModelsIterable:
     """Lists available models.
@@ -205,12 +158,17 @@ def list_models(
         `types.Model` objects.
 
     """
-    return _list_models_iter_pages(page_size=page_size, select="base", client=client)
+    if client is None:
+        client = get_default_model_client()
+
+    for model in client.list_models(page_size=page_size):
+        model = type(model).to_dict(model)
+        yield model_types.Model(**model)
 
 
 def list_tuned_models(
     *,
-    page_size: int | None = None,
+    page_size: int | None = 50,
     client: glm.ModelServiceClient | None = None,
 ) -> model_types.TunedModelsIterable:
     """Lists available models.
@@ -228,7 +186,12 @@ def list_tuned_models(
     Yields:
         `types.TunedModel` objects.
     """
-    return _list_models_iter_pages(page_size=page_size, select="tuned", client=client)
+    if client is None:
+        client = get_default_model_client()
+
+    for model in client.list_tuned_models(page_size=page_size):
+        model = type(model).to_dict(model)
+        yield model_types.decode_tuned_model(model)
 
 
 def create_tuned_model(

tests/test_models.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import copy
+from collections.abc import Iterable
 import datetime
 import dataclasses
 import pathlib
@@ -42,6 +43,8 @@ def setUp(self):
 
         client._client_manager.model_client = self.client
 
+        # TODO(markdaoust): Check if typechecking works better if wee define this as a
+        #                   subclass of `glm.ModelServiceClient`, would pyi files for `glm` help?
         def add_client_method(f):
             name = f.__name__
             setattr(self.client, name, f)
@@ -70,10 +73,6 @@ def get_tuned_model(
             response = copy.copy(self.responses["get_tuned_model"])
             return response
 
-        @dataclasses.dataclass
-        class ListWrapper:
-            _response: Any
-
         @add_client_method
         def list_models(
             request: Union[glm.ListModelsRequest, None] = None,
@@ -85,25 +84,25 @@ def list_models(
                 request = glm.ListModelsRequest(page_size=page_size, page_token=page_token)
             self.assertIsInstance(request, glm.ListModelsRequest)
             self.observed_requests.append(request)
-            response = self.responses["list_models"][request.page_token]
-            return ListWrapper(response)
+            response = self.responses["list_models"]
+            return (item for item in response)
 
         @add_client_method
         def list_tuned_models(
             request: glm.ListTunedModelsRequest = None,
             *,
             page_size=None,
             page_token=None,
-        ) -> glm.ListModelsResponse:
+        ) -> Iterable[glm.TunedModel]:
             if request is None:
                 request = glm.ListTunedModelsRequest(page_size=page_size, page_token=page_token)
             self.assertIsInstance(request, glm.ListTunedModelsRequest)
             self.observed_requests.append(request)
-            response = self.responses["list_tuned_models"][request.page_token]
-            return ListWrapper(response)
+            response = self.responses["list_tuned_models"]
+            return (item for item in response)
 
         @add_client_method
-        def update_tuned_model(request: glm.UpdateTunedModelRequest):
+        def update_tuned_model(request: glm.UpdateTunedModelRequest) -> glm.TunedModel:
             self.observed_requests.append(request)
             response = self.responses.get("update_tuned_model", None)
             if response is None:
@@ -156,24 +155,13 @@ def test_fail_with_unscoped_model_name(self, name):
             model = models.get_model(name)
 
     def test_list_models(self):
+        # The low level lib wraps the response in an iterable, so this is a fair test.
         self.responses = {
-            "list_models": {
-                # The first request doesn't pass a page token
-                "": glm.ListModelsResponse(
-                    models=[
-                        glm.Model(name="models/fake-bison-001"),
-                        glm.Model(name="models/fake-bison-002"),
-                    ],
-                    next_page_token="page1",
-                ),
-                "page1": glm.ListModelsResponse(
-                    models=[
-                        glm.Model(name="models/fake-bison-003"),
-                    ],
-                    # The last page returns an empty page token.
-                    next_page_token="",
-                ),
-            }
+            "list_models": [
+                glm.Model(name="models/fake-bison-001"),
+                glm.Model(name="models/fake-bison-002"),
+                glm.Model(name="models/fake-bison-003"),
+            ]
         }
 
         found_models = list(models.list_models())
@@ -183,23 +171,12 @@ def test_list_models(self):
 
     def test_list_tuned_models(self):
         self.responses = {
-            "list_tuned_models": {
-                # The first request doesn't pass a page token
-                "": glm.ListTunedModelsResponse(
-                    tuned_models=[
-                        glm.TunedModel(name="tunedModels/my-pig-001"),
-                        glm.TunedModel(name="tunedModels/my-pig-002"),
-                    ],
-                    next_page_token="page1",
-                ),
-                "page1": glm.ListTunedModelsResponse(
-                    tuned_models=[
-                        glm.TunedModel(name="tunedModels/my-pig-003"),
-                    ],
-                    # The last page returns an empty page token.
-                    next_page_token="",
-                ),
-            }
+            # The low level lib wraps the response in an iterable, so this is a fair test.
+            "list_tuned_models": [
+                glm.TunedModel(name="tunedModels/my-pig-001"),
+                glm.TunedModel(name="tunedModels/my-pig-002"),
+                glm.TunedModel(name="tunedModels/my-pig-003"),
+            ]
         }
         found_models = list(models.list_tuned_models())
         self.assertLen(found_models, 3)