Skip to content

Commit fb3ef47

Browse files
yanxi0830yanxi0830
authored
Unwrapping for "data" (#86)
# What does this PR do? - add `unwrap_response_fields: [data]` ## Test Plan ```python from llama_stack_client import LlamaStackClient from rich.pretty import pprint client = LlamaStackClient(base_url="http://localhost:5000") pprint(client.models.list()) ``` <img width="499" alt="image" src="https://github.com/user-attachments/assets/632ae764-7ef0-4e44-bae4-370bdca1d8d7" /> ## Sources Please link relevant resources if necessary. ## Before submitting - [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case). - [ ] Ran pre-commit to handle lint / formatting issues. - [ ] Read the [contributor guideline](https://github.com/meta-llama/llama-stack/blob/main/CONTRIBUTING.md), Pull Request section? - [ ] Updated relevant documentation. - [ ] Wrote necessary unit or integration tests.
1 parent c94bf4a commit fb3ef47

22 files changed

+281
-310
lines changed

‎src/llama_stack_client/_client.py‎

Lines changed: 92 additions & 214 deletions
Large diffs are not rendered by default.

‎src/llama_stack_client/_wrappers.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2+
3+
from typing import Generic, TypeVar
4+
5+
from ._models import GenericModel
6+
7+
__all__ = ["DataWrapper"]
8+
9+
_T = TypeVar("_T")
10+
11+
12+
class DataWrapper(GenericModel, Generic[_T]):
13+
data: _T
14+
15+
@staticmethod
16+
def _unwrapper(obj: "DataWrapper[_T]") -> _T:
17+
return obj.data

‎src/llama_stack_client/resources/datasets.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Dict, Union, Iterable, Optional
5+
from typing import Dict, Type, Union, Iterable, Optional, cast
66

77
import httpx
88

@@ -21,6 +21,7 @@
2121
async_to_raw_response_wrapper,
2222
async_to_streamed_response_wrapper,
2323
)
24+
from .._wrappers import DataWrapper
2425
from .._base_client import make_request_options
2526
from ..types.shared_params.url import URL
2627
from ..types.dataset_list_response import DatasetListResponse
@@ -126,9 +127,13 @@ def list(
126127
return self._get(
127128
"/v1/datasets",
128129
options=make_request_options(
129-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
130+
extra_headers=extra_headers,
131+
extra_query=extra_query,
132+
extra_body=extra_body,
133+
timeout=timeout,
134+
post_parser=DataWrapper[DatasetListResponse]._unwrapper,
130135
),
131-
cast_to=DatasetListResponse,
136+
cast_to=cast(Type[DatasetListResponse], DataWrapper[DatasetListResponse]),
132137
)
133138

134139
def register(
@@ -328,9 +333,13 @@ async def list(
328333
return await self._get(
329334
"/v1/datasets",
330335
options=make_request_options(
331-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
336+
extra_headers=extra_headers,
337+
extra_query=extra_query,
338+
extra_body=extra_body,
339+
timeout=timeout,
340+
post_parser=DataWrapper[DatasetListResponse]._unwrapper,
332341
),
333-
cast_to=DatasetListResponse,
342+
cast_to=cast(Type[DatasetListResponse], DataWrapper[DatasetListResponse]),
334343
)
335344

336345
async def register(

‎src/llama_stack_client/resources/eval_tasks.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Dict, List, Union, Iterable, Optional
5+
from typing import Dict, List, Type, Union, Iterable, Optional, cast
66

77
import httpx
88

@@ -21,6 +21,7 @@
2121
async_to_raw_response_wrapper,
2222
async_to_streamed_response_wrapper,
2323
)
24+
from .._wrappers import DataWrapper
2425
from .._base_client import make_request_options
2526
from ..types.eval_task import EvalTask
2627
from ..types.eval_task_list_response import EvalTaskListResponse
@@ -124,9 +125,13 @@ def list(
124125
return self._get(
125126
"/v1/eval-tasks",
126127
options=make_request_options(
127-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
128+
extra_headers=extra_headers,
129+
extra_query=extra_query,
130+
extra_body=extra_body,
131+
timeout=timeout,
132+
post_parser=DataWrapper[EvalTaskListResponse]._unwrapper,
128133
),
129-
cast_to=EvalTaskListResponse,
134+
cast_to=cast(Type[EvalTaskListResponse], DataWrapper[EvalTaskListResponse]),
130135
)
131136

132137
def register(
@@ -283,9 +288,13 @@ async def list(
283288
return await self._get(
284289
"/v1/eval-tasks",
285290
options=make_request_options(
286-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
291+
extra_headers=extra_headers,
292+
extra_query=extra_query,
293+
extra_body=extra_body,
294+
timeout=timeout,
295+
post_parser=DataWrapper[EvalTaskListResponse]._unwrapper,
287296
),
288-
cast_to=EvalTaskListResponse,
297+
cast_to=cast(Type[EvalTaskListResponse], DataWrapper[EvalTaskListResponse]),
289298
)
290299

291300
async def register(

‎src/llama_stack_client/resources/memory_banks.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Any, Optional, cast
5+
from typing import Any, Type, Optional, cast
66

77
import httpx
88

@@ -21,6 +21,7 @@
2121
async_to_raw_response_wrapper,
2222
async_to_streamed_response_wrapper,
2323
)
24+
from .._wrappers import DataWrapper
2425
from .._base_client import make_request_options
2526
from ..types.memory_bank_list_response import MemoryBankListResponse
2627
from ..types.memory_bank_register_response import MemoryBankRegisterResponse
@@ -130,9 +131,13 @@ def list(
130131
return self._get(
131132
"/v1/memory-banks",
132133
options=make_request_options(
133-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
134+
extra_headers=extra_headers,
135+
extra_query=extra_query,
136+
extra_body=extra_body,
137+
timeout=timeout,
138+
post_parser=DataWrapper[MemoryBankListResponse]._unwrapper,
134139
),
135-
cast_to=MemoryBankListResponse,
140+
cast_to=cast(Type[MemoryBankListResponse], DataWrapper[MemoryBankListResponse]),
136141
)
137142

138143
def register(
@@ -337,9 +342,13 @@ async def list(
337342
return await self._get(
338343
"/v1/memory-banks",
339344
options=make_request_options(
340-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
345+
extra_headers=extra_headers,
346+
extra_query=extra_query,
347+
extra_body=extra_body,
348+
timeout=timeout,
349+
post_parser=DataWrapper[MemoryBankListResponse]._unwrapper,
341350
),
342-
cast_to=MemoryBankListResponse,
351+
cast_to=cast(Type[MemoryBankListResponse], DataWrapper[MemoryBankListResponse]),
343352
)
344353

345354
async def register(

‎src/llama_stack_client/resources/models.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Dict, Union, Iterable, Optional
5+
from typing import Dict, Type, Union, Iterable, Optional, cast
66
from typing_extensions import Literal
77

88
import httpx
@@ -22,6 +22,7 @@
2222
async_to_raw_response_wrapper,
2323
async_to_streamed_response_wrapper,
2424
)
25+
from .._wrappers import DataWrapper
2526
from ..types.model import Model
2627
from .._base_client import make_request_options
2728
from ..types.model_list_response import ModelListResponse
@@ -125,9 +126,13 @@ def list(
125126
return self._get(
126127
"/v1/models",
127128
options=make_request_options(
128-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
129+
extra_headers=extra_headers,
130+
extra_query=extra_query,
131+
extra_body=extra_body,
132+
timeout=timeout,
133+
post_parser=DataWrapper[ModelListResponse]._unwrapper,
129134
),
130-
cast_to=ModelListResponse,
135+
cast_to=cast(Type[ModelListResponse], DataWrapper[ModelListResponse]),
131136
)
132137

133138
def register(
@@ -324,9 +329,13 @@ async def list(
324329
return await self._get(
325330
"/v1/models",
326331
options=make_request_options(
327-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
332+
extra_headers=extra_headers,
333+
extra_query=extra_query,
334+
extra_body=extra_body,
335+
timeout=timeout,
336+
post_parser=DataWrapper[ModelListResponse]._unwrapper,
328337
),
329-
cast_to=ModelListResponse,
338+
cast_to=cast(Type[ModelListResponse], DataWrapper[ModelListResponse]),
330339
)
331340

332341
async def register(

‎src/llama_stack_client/resources/post_training/job.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Optional
5+
from typing import Type, Optional, cast
66

77
import httpx
88

@@ -20,6 +20,7 @@
2020
async_to_raw_response_wrapper,
2121
async_to_streamed_response_wrapper,
2222
)
23+
from ..._wrappers import DataWrapper
2324
from ..._base_client import make_request_options
2425
from ...types.post_training import job_cancel_params, job_status_params, job_artifacts_params
2526
from ...types.post_training.job_list_response import JobListResponse
@@ -83,9 +84,13 @@ def list(
8384
return self._get(
8485
"/v1/post-training/jobs",
8586
options=make_request_options(
86-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
87+
extra_headers=extra_headers,
88+
extra_query=extra_query,
89+
extra_body=extra_body,
90+
timeout=timeout,
91+
post_parser=DataWrapper[JobListResponse]._unwrapper,
8792
),
88-
cast_to=JobListResponse,
93+
cast_to=cast(Type[JobListResponse], DataWrapper[JobListResponse]),
8994
)
9095

9196
def artifacts(
@@ -273,9 +278,13 @@ async def list(
273278
return await self._get(
274279
"/v1/post-training/jobs",
275280
options=make_request_options(
276-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
281+
extra_headers=extra_headers,
282+
extra_query=extra_query,
283+
extra_body=extra_body,
284+
timeout=timeout,
285+
post_parser=DataWrapper[JobListResponse]._unwrapper,
277286
),
278-
cast_to=JobListResponse,
287+
cast_to=cast(Type[JobListResponse], DataWrapper[JobListResponse]),
279288
)
280289

281290
async def artifacts(

‎src/llama_stack_client/resources/providers.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
from __future__ import annotations
44

5+
from typing import Type, cast
6+
57
import httpx
68

79
from .._types import NOT_GIVEN, Body, Query, Headers, NotGiven
@@ -14,6 +16,7 @@
1416
async_to_raw_response_wrapper,
1517
async_to_streamed_response_wrapper,
1618
)
19+
from .._wrappers import DataWrapper
1720
from .._base_client import make_request_options
1821
from ..types.provider_list_response import ProviderListResponse
1922

@@ -74,9 +77,13 @@ def list(
7477
return self._get(
7578
"/v1/providers/list",
7679
options=make_request_options(
77-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
80+
extra_headers=extra_headers,
81+
extra_query=extra_query,
82+
extra_body=extra_body,
83+
timeout=timeout,
84+
post_parser=DataWrapper[ProviderListResponse]._unwrapper,
7885
),
79-
cast_to=ProviderListResponse,
86+
cast_to=cast(Type[ProviderListResponse], DataWrapper[ProviderListResponse]),
8087
)
8188

8289

@@ -134,9 +141,13 @@ async def list(
134141
return await self._get(
135142
"/v1/providers/list",
136143
options=make_request_options(
137-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
144+
extra_headers=extra_headers,
145+
extra_query=extra_query,
146+
extra_body=extra_body,
147+
timeout=timeout,
148+
post_parser=DataWrapper[ProviderListResponse]._unwrapper,
138149
),
139-
cast_to=ProviderListResponse,
150+
cast_to=cast(Type[ProviderListResponse], DataWrapper[ProviderListResponse]),
140151
)
141152

142153

‎src/llama_stack_client/resources/scoring_functions.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Optional
5+
from typing import Type, Optional, cast
66

77
import httpx
88

@@ -21,6 +21,7 @@
2121
async_to_raw_response_wrapper,
2222
async_to_streamed_response_wrapper,
2323
)
24+
from .._wrappers import DataWrapper
2425
from .._base_client import make_request_options
2526
from ..types.scoring_fn import ScoringFn
2627
from ..types.shared_params.return_type import ReturnType
@@ -125,9 +126,13 @@ def list(
125126
return self._get(
126127
"/v1/scoring-functions",
127128
options=make_request_options(
128-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
129+
extra_headers=extra_headers,
130+
extra_query=extra_query,
131+
extra_body=extra_body,
132+
timeout=timeout,
133+
post_parser=DataWrapper[ScoringFunctionListResponse]._unwrapper,
129134
),
130-
cast_to=ScoringFunctionListResponse,
135+
cast_to=cast(Type[ScoringFunctionListResponse], DataWrapper[ScoringFunctionListResponse]),
131136
)
132137

133138
def register(
@@ -284,9 +289,13 @@ async def list(
284289
return await self._get(
285290
"/v1/scoring-functions",
286291
options=make_request_options(
287-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
292+
extra_headers=extra_headers,
293+
extra_query=extra_query,
294+
extra_body=extra_body,
295+
timeout=timeout,
296+
post_parser=DataWrapper[ScoringFunctionListResponse]._unwrapper,
288297
),
289-
cast_to=ScoringFunctionListResponse,
298+
cast_to=cast(Type[ScoringFunctionListResponse], DataWrapper[ScoringFunctionListResponse]),
290299
)
291300

292301
async def register(

0 commit comments

Comments
 (0)