Skip to content

feat(dataset api): (1.2/n) datasets.iterrorws pagination api updates #1656

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 15, 2025
Merged

feat(dataset api): (1.2/n) datasets.iterrorws pagination api updates #1656

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c45f083
dataset pagination
yanxi0830 Mar 15, 2025
fbe6eef
dataset pagination
yanxi0830 Mar 15, 2025
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
44 changes: 15 additions & 29 deletions docs/_static/llama-stack-spec.html
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@
}
],
"paths": {
"/v1/datasets/{dataset_id}/rows": {
"/v1/datasets/{dataset_id}/append-rows": {
"post": {
"responses": {
"200": {
Expand Down Expand Up @@ -2185,7 +2185,7 @@
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/PaginatedRowsResult"
"$ref": "#/components/schemas/IterrowsResponse"
}
}
}
Expand All @@ -2206,7 +2206,7 @@
"tags": [
"DatasetIO"
],
"description": "Get a paginated list of rows from a dataset.",
"description": "Get a paginated list of rows from a dataset. Uses cursor-based pagination.",
"parameters": [
{
"name": "dataset_id",
Expand All @@ -2218,30 +2218,21 @@
}
},
{
"name": "rows_in_page",
"name": "start_index",
"in": "query",
"description": "The number of rows to get per page.",
"required": true,
"schema": {
"type": "integer"
}
},
{
"name": "page_token",
"in": "query",
"description": "The token to get the next page of rows.",
"description": "Index into dataset for the first row to get. Get all rows if None.",
"required": false,
"schema": {
"type": "string"
"type": "integer"
}
},
{
"name": "filter_condition",
"name": "limit",
"in": "query",
"description": "(Optional) A condition to filter the rows by.",
"description": "The number of rows to get per page.",
"required": false,
"schema": {
"type": "string"
"type": "integer"
}
}
]
Expand Down Expand Up @@ -8137,10 +8128,10 @@
],
"title": "ToolInvocationResult"
},
"PaginatedRowsResult": {
"IterrowsResponse": {
"type": "object",
"properties": {
"rows": {
"data": {
"type": "array",
"items": {
"type": "object",
Expand Down Expand Up @@ -8169,21 +8160,16 @@
},
"description": "The rows in the current page."
},
"total_count": {
"next_index": {
"type": "integer",
"description": "The total number of rows in the dataset."
},
"next_page_token": {
"type": "string",
"description": "The token to get the next page of rows."
"description": "Index into dataset for the first row in the next page. None if there are no more rows."
}
},
"additionalProperties": false,
"required": [
"rows",
"total_count"
"data"
],
"title": "PaginatedRowsResult",
"title": "IterrowsResponse",
"description": "A paginated list of rows from a dataset."
},
"ListAgentSessionsResponse": {
Expand Down
44 changes: 18 additions & 26 deletions docs/_static/llama-stack-spec.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ info:
servers:
- url: http://any-hosted-llama-stack.com
paths:
/v1/datasets/{dataset_id}/rows:
/v1/datasets/{dataset_id}/append-rows:
post:
responses:
'200':
Expand Down Expand Up @@ -1465,7 +1465,7 @@ paths:
content:
application/json:
schema:
$ref: '#/components/schemas/PaginatedRowsResult'
$ref: '#/components/schemas/IterrowsResponse'
'400':
$ref: '#/components/responses/BadRequest400'
'429':
Expand All @@ -1479,7 +1479,7 @@ paths:
tags:
- DatasetIO
description: >-
Get a paginated list of rows from a dataset.
Get a paginated list of rows from a dataset. Uses cursor-based pagination.
parameters:
- name: dataset_id
in: path
Expand All @@ -1488,25 +1488,19 @@ paths:
required: true
schema:
type: string
- name: rows_in_page
- name: start_index
in: query
description: The number of rows to get per page.
required: true
schema:
type: integer
- name: page_token
in: query
description: The token to get the next page of rows.
description: >-
Index into dataset for the first row to get. Get all rows if None.
required: false
schema:
type: string
- name: filter_condition
type: integer
- name: limit
in: query
description: >-
(Optional) A condition to filter the rows by.
description: The number of rows to get per page.
required: false
schema:
type: string
type: integer
/v1/eval/benchmarks/{benchmark_id}/jobs/{job_id}:
get:
responses:
Expand Down Expand Up @@ -5547,10 +5541,10 @@ components:
required:
- content
title: ToolInvocationResult
PaginatedRowsResult:
IterrowsResponse:
type: object
properties:
rows:
data:
type: array
items:
type: object
Expand All @@ -5563,17 +5557,15 @@ components:
- type: array
- type: object
description: The rows in the current page.
total_count:
next_index:
type: integer
description: The total number of rows in the dataset.
next_page_token:
type: string
description: The token to get the next page of rows.
description: >-
Index into dataset for the first row in the next page. None if there are
no more rows.
additionalProperties: false
required:
- rows
- total_count
title: PaginatedRowsResult
- data
title: IterrowsResponse
description: A paginated list of rows from a dataset.
ListAgentSessionsResponse:
type: object
Expand Down
33 changes: 15 additions & 18 deletions llama_stack/apis/datasetio/datasetio.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,16 @@


@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):
Expand All @@ -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: ...
Loading