google-genai 0.0.1__py3-none-any.whl

google/genai/pagers.py

@@ -0,0 +1,245 @@
+# Copyright 2024 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.
+#
+
+"""Pagers for the GenAI List APIs."""
+
+# pylint: disable=protected-access
+
+import copy
+from typing import Any, AsyncIterator, Awaitable, Callable, Generic, Iterator, Literal, TypeVar
+
+T = TypeVar('T')
+
+PagedItem = Literal[
+    'batch_jobs', 'models', 'tuning_jobs', 'files', 'cached_contents'
+]
+
+
+class _BasePager(Generic[T]):
+  """Base pager class for iterating through paginated results."""
+
+  def __init__(
+      self,
+      name: PagedItem,
+      request: Callable[Any, Any],
+      response: Any,
+      config: Any,
+  ):
+    self._name = name
+    self._request = request
+
+    self._page = getattr(response, self._name) or []
+    self._idx = 0
+
+    if not config:
+      request_config = {}
+    elif isinstance(config, dict):
+      request_config = copy.deepcopy(config)
+    else:
+      request_config = dict(config)
+    request_config['page_token'] = getattr(response, 'next_page_token')
+    self._config = request_config
+
+    self._page_size = request_config.get('page_size', len(self._page))
+
+  @property
+  def page(self) -> list[T]:
+    """Returns the current page, which is a list of items.
+
+    The returned list of items is a subset of the entire list.
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = client.batches.list(config={'page_size': 5})
+      print(f"first page: {batch_jobs_pager.page}")
+      # first page: [BatchJob(name='projects/./locations/./batchPredictionJobs/1
+    """
+
+    return self._page
+
+  @property
+  def name(self) -> str:
+    """Returns the type of paged item (for example, ``batch_jobs``).
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = client.batches.list(config={'page_size': 5})
+      print(f"name: {batch_jobs_pager.name}")
+      # name: batch_jobs
+    """
+
+    return self._name
+
+  @property
+  def page_size(self) -> int:
+    """Returns the length of the page fetched each time by this pager.
+
+    The number of items in the page is less than or equal to the page length.
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = client.batches.list(config={'page_size': 5})
+      print(f"page_size: {batch_jobs_pager.page_size}")
+      # page_size: 5
+    """
+
+    return self._page_size
+
+  @property
+  def config(self) -> dict[str, Any]:
+    """Returns the configuration when making the API request for the next page.
+
+    A configuration is a set of optional parameters and arguments that can be
+    used to customize the API request. For example, the ``page_token`` parameter
+    contains the token to request the next page.
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = client.batches.list(config={'page_size': 5})
+      print(f"config: {batch_jobs_pager.config}")
+      # config: {'page_size': 5, 'page_token': 'AMEw9yO5jnsGnZJLHSKDFHJJu'}
+    """
+
+    return self._config
+
+  def __len__(self) -> int:
+    """Returns the total number of items in the current page."""
+    return len(self.page)
+
+  def __getitem__(self, index: int) -> T:
+    """Returns the item at the given index."""
+    return self.page[index]
+
+  def _init_next_page(self, response: Any) -> None:
+    """Initializes the next page from the response.
+
+    This is an internal method that should be called by subclasses after
+    fetching the next page.
+
+    Args:
+      response: The response object from the API request.
+    """
+    self.__init__(self.name, self._request, response, self.config)
+
+
+class Pager(_BasePager[T]):
+  """Pager class for iterating through paginated results."""
+
+  def __next__(self) -> T:
+    """Returns the next item."""
+    if self._idx >= len(self):
+      try:
+        self.next_page()
+      except IndexError:
+        raise StopIteration
+
+    item = self.page[self._idx]
+    self._idx += 1
+    return item
+
+  def __iter__(self) -> Iterator[T]:
+    """Returns an iterator over the items."""
+    self._idx = 0
+    return self
+
+  def next_page(self) -> list[T]:
+    """Fetches the next page of items. This makes a new API request.
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = client.batches.list(config={'page_size': 5})
+      print(f"current page: {batch_jobs_pager.page}")
+      batch_jobs_pager.next_page()
+      print(f"next page: {batch_jobs_pager.page}")
+      # current page: [BatchJob(name='projects/.../batchPredictionJobs/1
+      # next page: [BatchJob(name='projects/.../batchPredictionJobs/6
+    """
+
+    if not self.config.get('page_token'):
+      raise IndexError('No more pages to fetch.')
+
+    response = self._request(config=self.config)
+    self._init_next_page(response)
+    return self.page
+
+
+class AsyncPager(_BasePager[T]):
+  """AsyncPager class for iterating through paginated results."""
+
+  def __init__(
+      self,
+      name: PagedItem,
+      request: Callable[Any, Awaitable[Any]],
+      response: Any,
+      config: Any,
+  ):
+    super().__init__(name, request, response, config)
+
+  def __aiter__(self) -> AsyncIterator[T]:
+    """Returns an async iterator over the items."""
+    self._idx = 0
+    return self
+
+  async def __anext__(self) -> Awaitable[T]:
+    """Returns the next item asynchronously."""
+    if self._idx >= len(self):
+      try:
+        await self.next_page()
+      except IndexError:
+        raise StopAsyncIteration
+
+    item = self.page[self._idx]
+    self._idx += 1
+    return item
+
+  async def next_page(self) -> list[T]:
+    """Fetches the next page of items asynchronously.
+
+    This makes a new API request.
+
+    Returns:
+      The next page of items.
+
+    Raises:
+      IndexError: No more pages to fetch.
+
+    Usage:
+
+    .. code-block:: python
+
+      batch_jobs_pager = await client.aio.batches.list(config={'page_size': 5})
+      print(f"current page: {batch_jobs_pager.page}")
+      await batch_jobs_pager.next_page()
+      print(f"next page: {batch_jobs_pager.page}")
+      # current page: [BatchJob(name='projects/.../batchPredictionJobs/1
+      # next page: [BatchJob(name='projects/.../batchPredictionJobs/6
+    """
+
+    if not self.config.get('page_token'):
+      raise IndexError('No more pages to fetch.')
+
+    response = await self._request(config=self.config)
+    self._init_next_page(response)
+    return self.page