dataset pagination

This commit is contained in:
Xi Yan 2025-03-15 13:53:32 -07:00
parent c7d741d89e
commit c45f083a84

View file

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