deep-diver
diff --git a/‎google/generativeai/retriever.py‎
Lines changed: 11 additions & 11 deletions b/‎google/generativeai/retriever.py‎
Lines changed: 11 additions & 11 deletions
@@ -17,7 +17,7 @@
 import re
 import string
 import dataclasses
-from typing import Optional
+from typing import AsyncIterable, Iterable, Optional
 
 import google.ai.generativelanguage as glm
 
@@ -178,9 +178,8 @@ async def delete_corpus_async(name: str, force: bool, client: glm.RetrieverServi
 def list_corpora(
     *,
     page_size: Optional[int] = None,
-    page_token: Optional[str] = None,
     client: glm.RetrieverServiceClient | None = None,
-) -> list[retriever_types.Corpus]:
+) -> Iterable[retriever_types.Corpus]:
     """
     List the Corpuses you own in the service.
 
@@ -194,21 +193,22 @@ def list_corpora(
     if client is None:
         client = get_default_retriever_client()
 
-    request = glm.ListCorporaRequest(page_size=page_size, page_token=page_token)
-    response = client.list_corpora(request)
-    return response
+    request = glm.ListCorporaRequest(page_size=page_size)
+    for corpus in client.list_corpora(request):
+        corpus = type(corpus).to_dict(corpus)
+        yield retriever_types.Corpus(**corpus)
 
 
 async def list_corpora_async(
     *,
     page_size: Optional[int] = None,
-    page_token: Optional[str] = None,
     client: glm.RetrieverServiceClient | None = None,
-) -> list[retriever_types.Corpus]:
+) -> AsyncIterable[retriever_types.Corpus]:
     """This is the async version of `retriever.list_corpora`."""
     if client is None:
         client = get_default_retriever_async_client()
 
-    request = glm.ListCorporaRequest(page_size=page_size, page_token=page_token)
-    response = await client.list_corpora(request)
-    return response
+    request = glm.ListCorporaRequest(page_size=page_size)
+    async for corpus in await client.list_corpora(request):
+        corpus = type(corpus).to_dict(corpus)
+        yield retriever_types.Corpus(**corpus)