Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 55 additions & 59 deletions google/genai/batches.py
Original file line number Diff line number Diff line change
Expand Up @@ -1902,6 +1902,34 @@ def delete(
self._api_client._verify_response(return_value)
return return_value

def list(
self, *, config: Optional[types.ListBatchJobsConfigOrDict] = None
) -> Pager[types.BatchJob]:
"""Lists batch jobs.

Args:
config (ListBatchJobsConfig): Optional configuration for the list request.

Returns:
A Pager object that contains one page of batch jobs. When iterating over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python
config = {'page_size': 10}
for batch_job in client.batches.list(config):
print(batch_job.name)
"""

list_request = self._list
return Pager(
'batch_jobs',
list_request,
self._list(config=config),
config,
)

def create(
self,
*,
Expand Down Expand Up @@ -1997,35 +2025,6 @@ def create_embeddings(
else:
return self._create_embeddings(model=model, src=src, config=config)

def list(
self, *, config: Optional[types.ListBatchJobsConfigOrDict] = None
) -> Pager[types.BatchJob]:
"""Lists batch jobs.

Args:
config (ListBatchJobsConfig): Optional configuration for the list request.

Returns:
A Pager object that contains one page of batch jobs. When iterating over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python

batch_jobs = client.batches.list(config={"page_size": 10})
for batch_job in batch_jobs:
print(f"Batch job: {batch_job.name}, state {batch_job.state}")
"""
if config is None:
config = types.ListBatchJobsConfig()
return Pager(
'batch_jobs',
self._list,
self._list(config=config),
config,
)


class AsyncBatches(_api_module.BaseModule):

Expand Down Expand Up @@ -2452,6 +2451,33 @@ async def delete(
self._api_client._verify_response(return_value)
return return_value

async def list(
self, *, config: Optional[types.ListBatchJobsConfigOrDict] = None
) -> AsyncPager[types.BatchJob]:
"""Lists batch jobs asynchronously.

Args:
config (ListBatchJobsConfig): Optional configuration for the list request.

Returns:
A Pager object that contains one page of batch jobs. When iterating over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python
async for batch_job in await client.aio.batches.list():
print(batch_job.name)
"""

list_request = self._list
return AsyncPager(
'batch_jobs',
list_request,
await self._list(config=config),
config,
)

async def create(
self,
*,
Expand Down Expand Up @@ -2552,33 +2578,3 @@ async def create_embeddings(
raise ValueError('Vertex AI does not support batches.create_embeddings.')
else:
return await self._create_embeddings(model=model, src=src, config=config)

async def list(
self, *, config: Optional[types.ListBatchJobsConfigOrDict] = None
) -> AsyncPager[types.BatchJob]:
"""Lists batch jobs asynchronously.

Args:
config (ListBatchJobsConfig): Optional configuration for the list request.

Returns:
A Pager object that contains one page of batch jobs. When iterating over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python

batch_jobs = await client.aio.batches.list(config={'page_size': 5})
print(f"current page: {batch_jobs.page}")
await batch_jobs_pager.next_page()
print(f"next page: {batch_jobs_pager.page}")
"""
if config is None:
config = types.ListBatchJobsConfig()
return AsyncPager(
'batch_jobs',
self._list,
await self._list(config=config),
config,
)
60 changes: 40 additions & 20 deletions google/genai/caches.py
Original file line number Diff line number Diff line change
Expand Up @@ -1129,15 +1129,6 @@ def update(
def _list(
self, *, config: Optional[types.ListCachedContentsConfigOrDict] = None
) -> types.ListCachedContentsResponse:
"""Lists cached content configurations.

.. code-block:: python

cached_contents = client.caches.list(config={'page_size': 2})
for cached_content in cached_contents:
print(cached_content)
"""

parameter_model = types._ListCachedContentsParameters(
config=config,
)
Expand Down Expand Up @@ -1196,9 +1187,28 @@ def _list(
def list(
self, *, config: Optional[types.ListCachedContentsConfigOrDict] = None
) -> Pager[types.CachedContent]:
"""Lists cached contents.

Args:
config (ListCachedContentsConfig): Optional configuration for the list
request.

Returns:
A Pager object that contains one page of cached contents. When iterating
over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python
for cached_content in client.caches.list():
print(cached_content.name)
"""

list_request = self._list
return Pager(
'cached_contents',
self._list,
list_request,
self._list(config=config),
config,
)
Expand Down Expand Up @@ -1505,15 +1515,6 @@ async def update(
async def _list(
self, *, config: Optional[types.ListCachedContentsConfigOrDict] = None
) -> types.ListCachedContentsResponse:
"""Lists cached content configurations.

.. code-block:: python

cached_contents = await client.aio.caches.list(config={'page_size': 2})
async for cached_content in cached_contents:
print(cached_content)
"""

parameter_model = types._ListCachedContentsParameters(
config=config,
)
Expand Down Expand Up @@ -1574,9 +1575,28 @@ async def _list(
async def list(
self, *, config: Optional[types.ListCachedContentsConfigOrDict] = None
) -> AsyncPager[types.CachedContent]:
"""Lists cached contents asynchronously.

Args:
config (ListCachedContentsConfig): Optional configuration for the list
request.

Returns:
A Pager object that contains one page of cached contents. When iterating
over
the pager, it automatically fetches the next page if there are more.

Usage:

.. code-block:: python
async for cached_content in await client.aio.caches.list():
print(cached_content.name)
"""

list_request = self._list
return AsyncPager(
'cached_contents',
self._list,
list_request,
await self._list(config=config),
config,
)
26 changes: 3 additions & 23 deletions google/genai/documents.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,17 +249,6 @@ def _list(
parent: str,
config: Optional[types.ListDocumentsConfigOrDict] = None,
) -> types.ListDocumentsResponse:
"""Lists all Documents in a FileSearchStore.

Args:
parent (str): The name of the FileSearchStore containing the Documents.
config (ListDocumentsConfig | None): Optional parameters for the request,
such as page_size.

Returns:
ListDocumentsResponse: A paginated list of Documents.
"""

parameter_model = types._ListDocumentsParameters(
parent=parent,
config=config,
Expand Down Expand Up @@ -328,6 +317,7 @@ def list(
for document in client.documents.list(parent='rag_store_name'):
print(f"document: {document.name} - {document.display_name}")
"""

list_request = partial(self._list, parent=parent)
return Pager(
'documents',
Expand Down Expand Up @@ -461,17 +451,6 @@ async def _list(
parent: str,
config: Optional[types.ListDocumentsConfigOrDict] = None,
) -> types.ListDocumentsResponse:
"""Lists all Documents in a FileSearchStore.

Args:
parent (str): The name of the FileSearchStore containing the Documents.
config (ListDocumentsConfig | None): Optional parameters for the request,
such as page_size.

Returns:
ListDocumentsResponse: A paginated list of Documents.
"""

parameter_model = types._ListDocumentsParameters(
parent=parent,
config=config,
Expand Down Expand Up @@ -540,9 +519,10 @@ async def list(
Usage:
.. code-block:: python
async for document in await
client.documents.list(parent='rag_store_name'):
client.aio.documents.list(parent='rag_store_name'):
print(f"document: {document.name} - {document.display_name}")
"""

list_request = partial(self._list, parent=parent)
return AsyncPager(
'documents',
Expand Down
Loading