Skip to content

Commit 42ec723

Browse files
ark-hand[bot]ark-hand[bot]
authored
feat(sessions): checkpoint restore + agent version query call layer
Hand-written change from the internal SDK repository. The original commit message is intentionally not mirrored; see the source repository for details. Sync-Source-Commit: ff7a1e744d148487ec735ad7dbcfdfd307d0b56c Hand-Written-Reason: No Ark-APIs provenance marker; treated as a hand-written source commit. Release-Version: 0.9.0
1 parent 1ecf5b2 commit 42ec723

6 files changed

Lines changed: 578 additions & 8 deletions

File tree

‎src/arkruntime/_client.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@ def __init__(
135135
self.environments = resources.Environments(self)
136136
self.memory_stores = resources.MemoryStores(self)
137137
self.sessions = resources.Sessions(self)
138+
self.checkpoints = resources.Checkpoints(self)
138139
self.skills = resources.Skills(self)
139140
self.vaults = resources.Vaults(self)
140141

@@ -289,6 +290,7 @@ def __init__(
289290
self.environments = resources.AsyncEnvironments(self)
290291
self.memory_stores = resources.AsyncMemoryStores(self)
291292
self.sessions = resources.AsyncSessions(self)
293+
self.checkpoints = resources.AsyncCheckpoints(self)
292294
self.skills = resources.AsyncSkills(self)
293295
self.vaults = resources.AsyncVaults(self)
294296

‎src/arkruntime/resources/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from .agents import Agents, AsyncAgents
66
from .batch import AsyncBatch, Batch
77
from .chat import AsyncChat, Chat
8+
from .checkpoints import AsyncCheckpoints, Checkpoints
89
from .content_generation import AsyncContentGeneration, ContentGeneration
910
from .embeddings import AsyncEmbeddings, Embeddings
1011
from .environments import AsyncEnvironments, AsyncEnvironmentWork, Environments, EnvironmentWork
@@ -31,7 +32,9 @@
3132
from .vaults import AsyncCredentials, AsyncVaults, Credentials, Vaults
3233

3334
__all__ = [
35+
"Checkpoints",
3436
"Chat",
37+
"AsyncCheckpoints",
3538
"AsyncChat",
3639
"Embeddings",
3740
"AsyncEmbeddings",

‎src/arkruntime/resources/agents/agents.py‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,13 +83,29 @@ def create(
8383
cast_to=Agent,
8484
)
8585

86-
def retrieve(self, agent_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None) -> Agent:
86+
def retrieve(
87+
self,
88+
agent_id: str,
89+
*,
90+
version: Optional[int] | NotGiven = NOT_GIVEN,
91+
extra_headers=None,
92+
extra_query=None,
93+
extra_body=None,
94+
timeout=None,
95+
) -> Agent:
8796
if not agent_id:
8897
raise ValueError("agent_id is required")
98+
query: dict = {}
99+
if version is not NOT_GIVEN and version is not None:
100+
query["version"] = version
89101
return self._get(
90102
f"{_PREFIX}/{agent_id}",
91103
options=make_request_options(
92-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
104+
query=query or None,
105+
extra_headers=extra_headers,
106+
extra_query=extra_query,
107+
extra_body=extra_body,
108+
timeout=timeout,
93109
),
94110
cast_to=Agent,
95111
)
@@ -248,14 +264,28 @@ async def create(
248264
)
249265

250266
async def retrieve(
251-
self, agent_id: str, *, extra_headers=None, extra_query=None, extra_body=None, timeout=None
267+
self,
268+
agent_id: str,
269+
*,
270+
version: Optional[int] | NotGiven = NOT_GIVEN,
271+
extra_headers=None,
272+
extra_query=None,
273+
extra_body=None,
274+
timeout=None,
252275
) -> Agent:
253276
if not agent_id:
254277
raise ValueError("agent_id is required")
278+
query: dict = {}
279+
if version is not NOT_GIVEN and version is not None:
280+
query["version"] = version
255281
return await self._get(
256282
f"{_PREFIX}/{agent_id}",
257283
options=make_request_options(
258-
extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
284+
query=query or None,
285+
extra_headers=extra_headers,
286+
extra_query=extra_query,
287+
extra_body=extra_body,
288+
timeout=timeout,
259289
),
260290
cast_to=Agent,
261291
)

‎src/arkruntime/resources/checkpoints.py‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
from __future__ import annotations
2+
3+
from typing import List, Optional
4+
5+
from .._base_client import make_request_options
6+
from .._managed_agents_serialize import dump_body
7+
from .._resource import AsyncAPIResource, SyncAPIResource
8+
from .._types import NOT_GIVEN, NotGiven
9+
from ..types.session.checkpoint_status import CheckpointStatus
10+
from ..types.session.create_session_checkpoint_response import (
11+
CreateSessionCheckpointResponse,
12+
)
13+
from ..types.session.delete_session_checkpoint_response import (
14+
DeleteSessionCheckpointResponse,
15+
)
16+
from ..types.session.get_session_checkpoint_response import (
17+
GetSessionCheckpointResponse,
18+
)
19+
from ..types.session.list_session_checkpoints_response import (
20+
ListSessionCheckpointsResponse,
21+
)
22+
23+
__all__ = ["Checkpoints", "AsyncCheckpoints"]
24+
25+
_PREFIX = "/checkpoints"
26+
27+
28+
def _list_query(
29+
*,
30+
source_session_id=NOT_GIVEN,
31+
status=NOT_GIVEN,
32+
limit=NOT_GIVEN,
33+
page=NOT_GIVEN,
34+
) -> dict:
35+
q: dict = {}
36+
for k, v in {
37+
"source_session_id": source_session_id,
38+
"status": status,
39+
"limit": limit,
40+
"page": page,
41+
}.items():
42+
if v is not NOT_GIVEN and v is not None:
43+
q[k] = v
44+
return q
45+
46+
47+
class Checkpoints(SyncAPIResource):
48+
def create(
49+
self,
50+
session_id: str,
51+
*,
52+
name: str,
53+
extra_headers=None,
54+
extra_query=None,
55+
extra_body=None,
56+
timeout=None,
57+
) -> CreateSessionCheckpointResponse:
58+
"""Create a checkpoint of an idle Session.
59+
60+
The returned checkpoint is initially ``creating`` and becomes
61+
``ready`` once the asynchronous snapshot completes; only ready
62+
checkpoints can be used as ``checkpoint_id`` in
63+
``sessions.create``.
64+
"""
65+
if not session_id:
66+
raise ValueError("session_id is required")
67+
return self._post(
68+
f"/sessions/{session_id}/checkpoints",
69+
body=dump_body({"name": name}),
70+
options=make_request_options(
71+
extra_headers=extra_headers,
72+
extra_query=extra_query,
73+
extra_body=extra_body,
74+
timeout=timeout,
75+
),
76+
cast_to=CreateSessionCheckpointResponse,
77+
)
78+
79+
def list(
80+
self,
81+
*,
82+
source_session_id: Optional[str] | NotGiven = NOT_GIVEN,
83+
status: Optional[List[CheckpointStatus]] | NotGiven = NOT_GIVEN,
84+
limit: Optional[int] | NotGiven = NOT_GIVEN,
85+
page: Optional[str] | NotGiven = NOT_GIVEN,
86+
extra_headers=None,
87+
extra_query=None,
88+
extra_body=None,
89+
timeout=None,
90+
) -> ListSessionCheckpointsResponse:
91+
return self._get(
92+
_PREFIX,
93+
options=make_request_options(
94+
query=_list_query(
95+
source_session_id=source_session_id,
96+
status=status,
97+
limit=limit,
98+
page=page,
99+
),
100+
extra_headers=extra_headers,
101+
extra_query=extra_query,
102+
extra_body=extra_body,
103+
timeout=timeout,
104+
),
105+
cast_to=ListSessionCheckpointsResponse,
106+
)
107+
108+
def retrieve(
109+
self,
110+
checkpoint_id: str,
111+
*,
112+
extra_headers=None,
113+
extra_query=None,
114+
extra_body=None,
115+
timeout=None,
116+
) -> GetSessionCheckpointResponse:
117+
if not checkpoint_id:
118+
raise ValueError("checkpoint_id is required")
119+
return self._get(
120+
f"{_PREFIX}/{checkpoint_id}",
121+
options=make_request_options(
122+
extra_headers=extra_headers,
123+
extra_query=extra_query,
124+
extra_body=extra_body,
125+
timeout=timeout,
126+
),
127+
cast_to=GetSessionCheckpointResponse,
128+
)
129+
130+
def delete(
131+
self,
132+
checkpoint_id: str,
133+
*,
134+
extra_headers=None,
135+
extra_query=None,
136+
extra_body=None,
137+
timeout=None,
138+
) -> DeleteSessionCheckpointResponse:
139+
if not checkpoint_id:
140+
raise ValueError("checkpoint_id is required")
141+
return self._delete(
142+
f"{_PREFIX}/{checkpoint_id}",
143+
options=make_request_options(
144+
extra_headers=extra_headers,
145+
extra_query=extra_query,
146+
extra_body=extra_body,
147+
timeout=timeout,
148+
),
149+
cast_to=DeleteSessionCheckpointResponse,
150+
)
151+
152+
153+
class AsyncCheckpoints(AsyncAPIResource):
154+
async def create(
155+
self,
156+
session_id: str,
157+
*,
158+
name: str,
159+
extra_headers=None,
160+
extra_query=None,
161+
extra_body=None,
162+
timeout=None,
163+
) -> CreateSessionCheckpointResponse:
164+
"""Async variant of :meth:`Checkpoints.create`."""
165+
if not session_id:
166+
raise ValueError("session_id is required")
167+
return await self._post(
168+
f"/sessions/{session_id}/checkpoints",
169+
body=dump_body({"name": name}),
170+
options=make_request_options(
171+
extra_headers=extra_headers,
172+
extra_query=extra_query,
173+
extra_body=extra_body,
174+
timeout=timeout,
175+
),
176+
cast_to=CreateSessionCheckpointResponse,
177+
)
178+
179+
async def list(
180+
self,
181+
*,
182+
source_session_id: Optional[str] | NotGiven = NOT_GIVEN,
183+
status: Optional[List[CheckpointStatus]] | NotGiven = NOT_GIVEN,
184+
limit: Optional[int] | NotGiven = NOT_GIVEN,
185+
page: Optional[str] | NotGiven = NOT_GIVEN,
186+
extra_headers=None,
187+
extra_query=None,
188+
extra_body=None,
189+
timeout=None,
190+
) -> ListSessionCheckpointsResponse:
191+
return await self._get(
192+
_PREFIX,
193+
options=make_request_options(
194+
query=_list_query(
195+
source_session_id=source_session_id,
196+
status=status,
197+
limit=limit,
198+
page=page,
199+
),
200+
extra_headers=extra_headers,
201+
extra_query=extra_query,
202+
extra_body=extra_body,
203+
timeout=timeout,
204+
),
205+
cast_to=ListSessionCheckpointsResponse,
206+
)
207+
208+
async def retrieve(
209+
self,
210+
checkpoint_id: str,
211+
*,
212+
extra_headers=None,
213+
extra_query=None,
214+
extra_body=None,
215+
timeout=None,
216+
) -> GetSessionCheckpointResponse:
217+
if not checkpoint_id:
218+
raise ValueError("checkpoint_id is required")
219+
return await self._get(
220+
f"{_PREFIX}/{checkpoint_id}",
221+
options=make_request_options(
222+
extra_headers=extra_headers,
223+
extra_query=extra_query,
224+
extra_body=extra_body,
225+
timeout=timeout,
226+
),
227+
cast_to=GetSessionCheckpointResponse,
228+
)
229+
230+
async def delete(
231+
self,
232+
checkpoint_id: str,
233+
*,
234+
extra_headers=None,
235+
extra_query=None,
236+
extra_body=None,
237+
timeout=None,
238+
) -> DeleteSessionCheckpointResponse:
239+
if not checkpoint_id:
240+
raise ValueError("checkpoint_id is required")
241+
return await self._delete(
242+
f"{_PREFIX}/{checkpoint_id}",
243+
options=make_request_options(
244+
extra_headers=extra_headers,
245+
extra_query=extra_query,
246+
extra_body=extra_body,
247+
timeout=timeout,
248+
),
249+
cast_to=DeleteSessionCheckpointResponse,
250+
)

0 commit comments

Comments
 (0)