From c45f083a84869f846af4e148a0e136acad8d6fa2 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Sat, 15 Mar 2025 13:53:32 -0700
Subject: [PATCH] dataset pagination

---
 llama_stack/apis/datasetio/datasetio.py | 33 +++++++++++--------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/llama_stack/apis/datasetio/datasetio.py b/llama_stack/apis/datasetio/datasetio.py
index 9ce90da01..caa7c51df 100644
--- a/llama_stack/apis/datasetio/datasetio.py
+++ b/llama_stack/apis/datasetio/datasetio.py
@@ -13,19 +13,16 @@ from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type
-class PaginatedRowsResult(BaseModel):
+class IterrowsResponse(BaseModel):
     """
     A paginated list of rows from a dataset.
 
-    :param rows: The rows in the current page.
-    :param total_count: The total number of rows in the dataset.
-    :param next_page_token: The token to get the next page of rows.
+    :param data: The rows in the current page.
+    :param next_index: Index into dataset for the first row in the next page. None if there are no more rows.
     """
 
-    # the rows obey the DatasetSchema for the given dataset
-    rows: List[Dict[str, Any]]
-    total_count: int
-    next_page_token: Optional[str] = None
+    data: List[Dict[str, Any]]
+    next_index: Optional[int] = None
 
 
 class DatasetStore(Protocol):
@@ -41,18 +38,18 @@ class DatasetIO(Protocol):
     async def iterrows(
         self,
         dataset_id: str,
-        rows_in_page: int,
-        page_token: Optional[str] = None,
-        filter_condition: Optional[str] = None,
-    ) -> PaginatedRowsResult:
-        """Get a paginated list of rows from a dataset.
+        start_index: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> IterrowsResponse:
+        """Get a paginated list of rows from a dataset. Uses cursor-based pagination.
 
         :param dataset_id: The ID of the dataset to get the rows from.
-        :param rows_in_page: The number of rows to get per page.
-        :param page_token: The token to get the next page of rows.
-        :param filter_condition: (Optional) A condition to filter the rows by.
+        :param start_index: Index into dataset for the first row to get. Get all rows if None.
+        :param limit: The number of rows to get per page.
         """
         ...
 
-    @webmethod(route="/datasets/{dataset_id}/rows", method="POST")
-    async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: ...
+    @webmethod(route="/datasets/{dataset_id}/append-rows", method="POST")
+    async def append_rows(
+        self, dataset_id: str, rows: List[Dict[str, Any]]
+    ) -> None: ...