Skip to content

Commit 2f53a3f

Browse files
feat: convert Datasets API to use FastAPI router
1 parent 32c453f commit 2f53a3f

File tree

8 files changed

+110
-24
lines changed

8 files changed

+110
-24
lines changed

.stats.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
configured_endpoints: 103
2-
openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/llamastack%2Fllama-stack-client-0ee7c2f0b6a30d3cd938f7508fa72003432f69bf8881f03817486510b8e2c7ab.yml
3-
openapi_spec_hash: a84bb7c34c365485bdf243ce1d960817
2+
openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/llamastack%2Fllama-stack-client-00b875a79e05a6b0af1115082977fcb29da1efbc553e3a2aa7b34d0c2bce0dee.yml
3+
openapi_spec_hash: 4e7110460ea24b382d51dc1911431df3
44
config_hash: 39578cfdeb4a10121f2cb3fa3e4d5e20

src/llama_stack_client/resources/beta/datasets.py

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ def retrieve(
7070
Get a dataset by its ID.
7171
7272
Args:
73+
dataset_id: The ID of the dataset to get.
74+
7375
extra_headers: Send extra headers
7476
7577
extra_query: Add additional query parameters to the request
@@ -217,14 +219,17 @@ def register(
217219
extra_body: Body | None = None,
218220
timeout: float | httpx.Timeout | None | NotGiven = not_given,
219221
) -> DatasetRegisterResponse:
220-
"""Register a new dataset.
222+
"""
223+
Register a new dataset.
221224
222225
Args:
223-
purpose: Purpose of the dataset.
226+
purpose: The purpose of the dataset.
224227
225-
Each purpose has a required input data schema.
228+
source: The data source of the dataset.
226229
227-
source: A dataset that can be obtained from a URI.
230+
dataset_id: The ID of the dataset. If not provided, an ID will be generated.
231+
232+
metadata: The metadata for the dataset.
228233
229234
extra_headers: Send extra headers
230235
@@ -267,6 +272,8 @@ def unregister(
267272
Unregister a dataset by its ID.
268273
269274
Args:
275+
dataset_id: The ID of the dataset to unregister.
276+
270277
extra_headers: Send extra headers
271278
272279
extra_query: Add additional query parameters to the request
@@ -322,6 +329,8 @@ async def retrieve(
322329
Get a dataset by its ID.
323330
324331
Args:
332+
dataset_id: The ID of the dataset to get.
333+
325334
extra_headers: Send extra headers
326335
327336
extra_query: Add additional query parameters to the request
@@ -469,14 +478,17 @@ async def register(
469478
extra_body: Body | None = None,
470479
timeout: float | httpx.Timeout | None | NotGiven = not_given,
471480
) -> DatasetRegisterResponse:
472-
"""Register a new dataset.
481+
"""
482+
Register a new dataset.
473483
474484
Args:
475-
purpose: Purpose of the dataset.
485+
purpose: The purpose of the dataset.
476486
477-
Each purpose has a required input data schema.
487+
source: The data source of the dataset.
478488
479-
source: A dataset that can be obtained from a URI.
489+
dataset_id: The ID of the dataset. If not provided, an ID will be generated.
490+
491+
metadata: The metadata for the dataset.
480492
481493
extra_headers: Send extra headers
482494
@@ -519,6 +531,8 @@ async def unregister(
519531
Unregister a dataset by its ID.
520532
521533
Args:
534+
dataset_id: The ID of the dataset to unregister.
535+
522536
extra_headers: Send extra headers
523537
524538
extra_query: Add additional query parameters to the request

src/llama_stack_client/types/beta/dataset_list_response.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,28 @@ class DatasetListResponseItemSourceUriDataSource(BaseModel):
2525
"""A dataset that can be obtained from a URI."""
2626

2727
uri: str
28+
"""The dataset can be obtained from a URI.
29+
30+
E.g. "https://mywebsite.com/mydata.jsonl", "lsfs://mydata.jsonl",
31+
"data:csv;base64,{base64_content}"
32+
"""
2833

2934
type: Optional[Literal["uri"]] = None
35+
"""The type of data source."""
3036

3137

3238
class DatasetListResponseItemSourceRowsDataSource(BaseModel):
3339
"""A dataset stored in rows."""
3440

3541
rows: List[Dict[str, object]]
42+
"""The dataset is stored in rows.
43+
44+
E.g. [{"messages": [{"role": "user", "content": "Hello, world!"}, {"role":
45+
"assistant", "content": "Hello, world!"}]}]
46+
"""
3647

3748
type: Optional[Literal["rows"]] = None
49+
"""The type of data source."""
3850

3951

4052
DatasetListResponseItemSource: TypeAlias = Annotated[
@@ -53,10 +65,10 @@ class DatasetListResponseItem(BaseModel):
5365
"""ID of the provider that owns this resource"""
5466

5567
purpose: Literal["post-training/messages", "eval/question-answer", "eval/messages-answer"]
56-
"""Purpose of the dataset. Each purpose has a required input data schema."""
68+
"""Purpose of the dataset indicating its intended use"""
5769

5870
source: DatasetListResponseItemSource
59-
"""A dataset that can be obtained from a URI."""
71+
"""Data source configuration for the dataset"""
6072

6173
metadata: Optional[Dict[str, object]] = None
6274
"""Any additional metadata for this dataset"""
@@ -65,6 +77,7 @@ class DatasetListResponseItem(BaseModel):
6577
"""Unique identifier for this resource in the provider"""
6678

6779
type: Optional[Literal["dataset"]] = None
80+
"""Type of resource, always 'dataset' for datasets"""
6881

6982

7083
DatasetListResponse: TypeAlias = List[DatasetListResponseItem]

src/llama_stack_client/types/beta/dataset_register_params.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,30 +16,44 @@
1616

1717
class DatasetRegisterParams(TypedDict, total=False):
1818
purpose: Required[Literal["post-training/messages", "eval/question-answer", "eval/messages-answer"]]
19-
"""Purpose of the dataset. Each purpose has a required input data schema."""
19+
"""The purpose of the dataset."""
2020

2121
source: Required[Source]
22-
"""A dataset that can be obtained from a URI."""
22+
"""The data source of the dataset."""
2323

2424
dataset_id: Optional[str]
25+
"""The ID of the dataset. If not provided, an ID will be generated."""
2526

2627
metadata: Optional[Dict[str, object]]
28+
"""The metadata for the dataset."""
2729

2830

2931
class SourceUriDataSource(TypedDict, total=False):
3032
"""A dataset that can be obtained from a URI."""
3133

3234
uri: Required[str]
35+
"""The dataset can be obtained from a URI.
36+
37+
E.g. "https://mywebsite.com/mydata.jsonl", "lsfs://mydata.jsonl",
38+
"data:csv;base64,{base64_content}"
39+
"""
3340

3441
type: Literal["uri"]
42+
"""The type of data source."""
3543

3644

3745
class SourceRowsDataSource(TypedDict, total=False):
3846
"""A dataset stored in rows."""
3947

4048
rows: Required[Iterable[Dict[str, object]]]
49+
"""The dataset is stored in rows.
50+
51+
E.g. [{"messages": [{"role": "user", "content": "Hello, world!"}, {"role":
52+
"assistant", "content": "Hello, world!"}]}]
53+
"""
4154

4255
type: Literal["rows"]
56+
"""The type of data source."""
4357

4458

4559
Source: TypeAlias = Union[SourceUriDataSource, SourceRowsDataSource]

src/llama_stack_client/types/beta/dataset_register_response.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,28 @@ class SourceUriDataSource(BaseModel):
1919
"""A dataset that can be obtained from a URI."""
2020

2121
uri: str
22+
"""The dataset can be obtained from a URI.
23+
24+
E.g. "https://mywebsite.com/mydata.jsonl", "lsfs://mydata.jsonl",
25+
"data:csv;base64,{base64_content}"
26+
"""
2227

2328
type: Optional[Literal["uri"]] = None
29+
"""The type of data source."""
2430

2531

2632
class SourceRowsDataSource(BaseModel):
2733
"""A dataset stored in rows."""
2834

2935
rows: List[Dict[str, object]]
36+
"""The dataset is stored in rows.
37+
38+
E.g. [{"messages": [{"role": "user", "content": "Hello, world!"}, {"role":
39+
"assistant", "content": "Hello, world!"}]}]
40+
"""
3041

3142
type: Optional[Literal["rows"]] = None
43+
"""The type of data source."""
3244

3345

3446
Source: TypeAlias = Annotated[Union[SourceUriDataSource, SourceRowsDataSource], PropertyInfo(discriminator="type")]
@@ -44,10 +56,10 @@ class DatasetRegisterResponse(BaseModel):
4456
"""ID of the provider that owns this resource"""
4557

4658
purpose: Literal["post-training/messages", "eval/question-answer", "eval/messages-answer"]
47-
"""Purpose of the dataset. Each purpose has a required input data schema."""
59+
"""Purpose of the dataset indicating its intended use"""
4860

4961
source: Source
50-
"""A dataset that can be obtained from a URI."""
62+
"""Data source configuration for the dataset"""
5163

5264
metadata: Optional[Dict[str, object]] = None
5365
"""Any additional metadata for this dataset"""
@@ -56,3 +68,4 @@ class DatasetRegisterResponse(BaseModel):
5668
"""Unique identifier for this resource in the provider"""
5769

5870
type: Optional[Literal["dataset"]] = None
71+
"""Type of resource, always 'dataset' for datasets"""

src/llama_stack_client/types/beta/dataset_retrieve_response.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,28 @@ class SourceUriDataSource(BaseModel):
1919
"""A dataset that can be obtained from a URI."""
2020

2121
uri: str
22+
"""The dataset can be obtained from a URI.
23+
24+
E.g. "https://mywebsite.com/mydata.jsonl", "lsfs://mydata.jsonl",
25+
"data:csv;base64,{base64_content}"
26+
"""
2227

2328
type: Optional[Literal["uri"]] = None
29+
"""The type of data source."""
2430

2531

2632
class SourceRowsDataSource(BaseModel):
2733
"""A dataset stored in rows."""
2834

2935
rows: List[Dict[str, object]]
36+
"""The dataset is stored in rows.
37+
38+
E.g. [{"messages": [{"role": "user", "content": "Hello, world!"}, {"role":
39+
"assistant", "content": "Hello, world!"}]}]
40+
"""
3041

3142
type: Optional[Literal["rows"]] = None
43+
"""The type of data source."""
3244

3345

3446
Source: TypeAlias = Annotated[Union[SourceUriDataSource, SourceRowsDataSource], PropertyInfo(discriminator="type")]
@@ -44,10 +56,10 @@ class DatasetRetrieveResponse(BaseModel):
4456
"""ID of the provider that owns this resource"""
4557

4658
purpose: Literal["post-training/messages", "eval/question-answer", "eval/messages-answer"]
47-
"""Purpose of the dataset. Each purpose has a required input data schema."""
59+
"""Purpose of the dataset indicating its intended use"""
4860

4961
source: Source
50-
"""A dataset that can be obtained from a URI."""
62+
"""Data source configuration for the dataset"""
5163

5264
metadata: Optional[Dict[str, object]] = None
5365
"""Any additional metadata for this dataset"""
@@ -56,3 +68,4 @@ class DatasetRetrieveResponse(BaseModel):
5668
"""Unique identifier for this resource in the provider"""
5769

5870
type: Optional[Literal["dataset"]] = None
71+
"""Type of resource, always 'dataset' for datasets"""

src/llama_stack_client/types/beta/list_datasets_response.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,4 @@ class ListDatasetsResponse(BaseModel):
1616
"""Response from listing datasets."""
1717

1818
data: DatasetListResponse
19+
"""List of datasets"""

tests/api_resources/beta/test_datasets.py

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,10 @@ def test_method_register(self, client: LlamaStackClient) -> None:
187187
with pytest.warns(DeprecationWarning):
188188
dataset = client.beta.datasets.register(
189189
purpose="post-training/messages",
190-
source={"uri": "uri"},
190+
source={
191+
"uri": "uri",
192+
"type": "uri",
193+
},
191194
)
192195

193196
assert_matches_type(DatasetRegisterResponse, dataset, path=["response"])
@@ -212,7 +215,10 @@ def test_raw_response_register(self, client: LlamaStackClient) -> None:
212215
with pytest.warns(DeprecationWarning):
213216
response = client.beta.datasets.with_raw_response.register(
214217
purpose="post-training/messages",
215-
source={"uri": "uri"},
218+
source={
219+
"uri": "uri",
220+
"type": "uri",
221+
},
216222
)
217223

218224
assert response.is_closed is True
@@ -225,7 +231,10 @@ def test_streaming_response_register(self, client: LlamaStackClient) -> None:
225231
with pytest.warns(DeprecationWarning):
226232
with client.beta.datasets.with_streaming_response.register(
227233
purpose="post-training/messages",
228-
source={"uri": "uri"},
234+
source={
235+
"uri": "uri",
236+
"type": "uri",
237+
},
229238
) as response:
230239
assert not response.is_closed
231240
assert response.http_request.headers.get("X-Stainless-Lang") == "python"
@@ -441,7 +450,10 @@ async def test_method_register(self, async_client: AsyncLlamaStackClient) -> Non
441450
with pytest.warns(DeprecationWarning):
442451
dataset = await async_client.beta.datasets.register(
443452
purpose="post-training/messages",
444-
source={"uri": "uri"},
453+
source={
454+
"uri": "uri",
455+
"type": "uri",
456+
},
445457
)
446458

447459
assert_matches_type(DatasetRegisterResponse, dataset, path=["response"])
@@ -466,7 +478,10 @@ async def test_raw_response_register(self, async_client: AsyncLlamaStackClient)
466478
with pytest.warns(DeprecationWarning):
467479
response = await async_client.beta.datasets.with_raw_response.register(
468480
purpose="post-training/messages",
469-
source={"uri": "uri"},
481+
source={
482+
"uri": "uri",
483+
"type": "uri",
484+
},
470485
)
471486

472487
assert response.is_closed is True
@@ -479,7 +494,10 @@ async def test_streaming_response_register(self, async_client: AsyncLlamaStackCl
479494
with pytest.warns(DeprecationWarning):
480495
async with async_client.beta.datasets.with_streaming_response.register(
481496
purpose="post-training/messages",
482-
source={"uri": "uri"},
497+
source={
498+
"uri": "uri",
499+
"type": "uri",
500+
},
483501
) as response:
484502
assert not response.is_closed
485503
assert response.http_request.headers.get("X-Stainless-Lang") == "python"

0 commit comments

Comments
 (0)