Skip to content

Commit 747a054

Browse files
leawnleawn
committed
Update agent-todo example
1 parent 8d9d29c commit 747a054

File tree

3 files changed

+155
-121
lines changed

3 files changed

+155
-121
lines changed

‎agent_todo/pyproject.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ dependencies = [
1010
"watchfiles>=1.0.4",
1111
"python-dotenv==1.0.1",
1212
"openai>=1.61.0",
13-
"restack-ai>=0.0.77",
13+
"restack-ai>=0.0.78",
1414
]
1515

1616
[project.scripts]
@@ -24,6 +24,9 @@ include = ["src"]
2424
[tool.hatch.build.targets.wheel]
2525
include = ["src"]
2626

27+
[tool.uv.sources]
28+
restack-ai = { path = "../../../sdk/engine/libraries/python/dist/restack_ai-0.0.78-py3-none-any.whl" }
29+
2730
[build-system]
2831
requires = ["hatchling"]
2932
build-backend = "hatchling.build"

‎agent_todo/src/agents/agent_todo.py‎

Lines changed: 115 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from datetime import timedelta
22

33
from pydantic import BaseModel
4-
from restack_ai.agent import agent, import_functions, log
4+
from restack_ai.agent import NonRetryableError, agent, import_functions, log
55

66
from src.workflows.todo_execute import TodoExecute, TodoExecuteParams
77

@@ -47,107 +47,128 @@ async def messages(self, messages_event: MessagesEvent) -> list[Message]:
4747
description="Execute a todo, needs to be created first and need confirmation from user before executing.",
4848
),
4949
]
50-
51-
completion = await agent.step(
52-
function=llm_chat,
53-
function_input=LlmChatInput(messages=self.messages, tools=tools),
54-
start_to_close_timeout=timedelta(seconds=120),
55-
)
56-
57-
log.info(f"completion: {completion}")
58-
59-
tool_calls = completion.choices[0].message.tool_calls
60-
self.messages.append(
61-
Message(
62-
role="assistant",
63-
content=completion.choices[0].message.content or "",
64-
tool_calls=tool_calls,
50+
try:
51+
completion = await agent.step(
52+
function=llm_chat,
53+
function_input=LlmChatInput(messages=self.messages, tools=tools),
54+
start_to_close_timeout=timedelta(seconds=120),
6555
)
66-
)
67-
68-
log.info(f"tool_calls: {tool_calls}")
69-
70-
if tool_calls:
71-
for tool_call in tool_calls:
72-
log.info(f"tool_call: {tool_call}")
73-
74-
name = tool_call.function.name
75-
76-
match name:
77-
case todo_create.__name__:
78-
args = TodoCreateParams.model_validate_json(
79-
tool_call.function.arguments
80-
)
81-
82-
result = await agent.step(
83-
function=todo_create,
84-
function_input=args,
85-
)
86-
self.messages.append(
87-
Message(
88-
role="tool",
89-
tool_call_id=tool_call.id,
90-
content=str(result),
91-
)
92-
)
93-
94-
completion_with_tool_call = await agent.step(
95-
function=llm_chat,
96-
function_input=LlmChatInput(
97-
messages=self.messages, tools=tools
98-
),
99-
start_to_close_timeout=timedelta(seconds=120),
100-
)
101-
self.messages.append(
102-
Message(
103-
role="assistant",
104-
content=completion_with_tool_call.choices[
105-
0
106-
].message.content
107-
or "",
108-
)
109-
)
110-
case TodoExecute.__name__:
111-
args = TodoExecuteParams.model_validate_json(
112-
tool_call.function.arguments
113-
)
114-
115-
result = await agent.child_execute(
116-
workflow=TodoExecute,
117-
workflow_id=tool_call.id,
118-
workflow_input=args,
119-
)
120-
self.messages.append(
121-
Message(
122-
role="tool",
123-
tool_call_id=tool_call.id,
124-
content=str(result),
125-
)
126-
)
127-
128-
completion_with_tool_call = await agent.step(
129-
function=llm_chat,
130-
function_input=LlmChatInput(
131-
messages=self.messages, tools=tools
132-
),
133-
start_to_close_timeout=timedelta(seconds=120),
134-
)
135-
self.messages.append(
136-
Message(
137-
role="assistant",
138-
content=completion_with_tool_call.choices[
139-
0
140-
].message.content
141-
or "",
142-
)
143-
)
56+
except Exception as e:
57+
error_message = f"Error during llm_chat: {e}"
58+
raise NonRetryableError(error_message) from e
14459
else:
60+
log.info(f"completion: {completion}")
61+
tool_calls = completion.choices[0].message.tool_calls
14562
self.messages.append(
14663
Message(
14764
role="assistant",
14865
content=completion.choices[0].message.content or "",
66+
tool_calls=tool_calls,
14967
)
15068
)
69+
70+
log.info(f"tool_calls: {tool_calls}")
71+
72+
if tool_calls:
73+
for tool_call in tool_calls:
74+
log.info(f"tool_call: {tool_call}")
75+
76+
name = tool_call.function.name
77+
78+
match name:
79+
case todo_create.__name__:
80+
args = TodoCreateParams.model_validate_json(
81+
tool_call.function.arguments
82+
)
83+
84+
try:
85+
result = await agent.step(
86+
function=todo_create,
87+
function_input=args,
88+
)
89+
except Exception as e:
90+
error_message = f"Error during todo_create: {e}"
91+
raise NonRetryableError(error_message) from e
92+
else:
93+
self.messages.append(
94+
Message(
95+
role="tool",
96+
tool_call_id=tool_call.id,
97+
content=str(result),
98+
)
99+
)
100+
try:
101+
completion_with_tool_call = await agent.step(
102+
function=llm_chat,
103+
function_input=LlmChatInput(
104+
messages=self.messages, tools=tools
105+
),
106+
start_to_close_timeout=timedelta(seconds=120),
107+
)
108+
except Exception as e:
109+
error_message = f"Error during llm_chat: {e}"
110+
raise NonRetryableError(error_message) from e
111+
else:
112+
self.messages.append(
113+
Message(
114+
role="assistant",
115+
content=completion_with_tool_call.choices[
116+
0
117+
].message.content
118+
or "",
119+
)
120+
)
121+
case TodoExecute.__name__:
122+
args = TodoExecuteParams.model_validate_json(
123+
tool_call.function.arguments
124+
)
125+
126+
try:
127+
result = await agent.child_execute(
128+
workflow=TodoExecute,
129+
workflow_id=tool_call.id,
130+
workflow_input=args,
131+
)
132+
except Exception as e:
133+
error_message = f"Error during TodoExecute: {e}"
134+
raise NonRetryableError(error_message) from e
135+
else:
136+
self.messages.append(
137+
Message(
138+
role="tool",
139+
tool_call_id=tool_call.id,
140+
content=str(result),
141+
)
142+
)
143+
144+
try:
145+
completion_with_tool_call = await agent.step(
146+
function=llm_chat,
147+
function_input=LlmChatInput(
148+
messages=self.messages, tools=tools
149+
),
150+
start_to_close_timeout=timedelta(seconds=120),
151+
)
152+
except Exception as e:
153+
error_message = f"Error during llm_chat: {e}"
154+
raise NonRetryableError(error_message) from e
155+
else:
156+
self.messages.append(
157+
Message(
158+
role="assistant",
159+
content=completion_with_tool_call.choices[
160+
0
161+
].message.content
162+
or "",
163+
)
164+
)
165+
else:
166+
self.messages.append(
167+
Message(
168+
role="assistant",
169+
content=completion.choices[0].message.content or "",
170+
)
171+
)
151172
except Exception as e:
152173
log.error(f"Error during message event: {e}")
153174
raise

‎agent_todo/src/workflows/todo_execute.py‎

Lines changed: 36 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from datetime import timedelta
22

33
from pydantic import BaseModel
4-
from restack_ai.workflow import import_functions, log, workflow
4+
from restack_ai.workflow import NonRetryableError, import_functions, log, workflow
55

66
with import_functions():
77
from src.functions.get_random import RandomParams, get_random
@@ -25,28 +25,38 @@ class TodoExecute:
2525
@workflow.run
2626
async def run(self, workflow_input: TodoExecuteParams) -> TodoExecuteResponse:
2727
log.info("TodoExecuteWorkflow started")
28-
random = await workflow.step(
29-
function=get_random,
30-
function_input=RandomParams(todo_title=workflow_input.todo_title),
31-
start_to_close_timeout=timedelta(seconds=120),
32-
)
33-
34-
await workflow.sleep(2)
35-
36-
result = await workflow.step(
37-
function=get_result,
38-
function_input=ResultParams(
39-
todo_title=workflow_input.todo_title,
40-
todo_id=workflow_input.todo_id,
41-
),
42-
start_to_close_timeout=timedelta(seconds=120),
43-
)
44-
45-
todo_details = TodoExecuteResponse(
46-
todo_id=workflow_input.todo_id,
47-
todo_title=workflow_input.todo_title,
48-
details=random,
49-
status=result.status,
50-
)
51-
log.info("TodoExecuteWorkflow done", result=todo_details)
52-
return todo_details
28+
try:
29+
random = await workflow.step(
30+
function=get_random,
31+
function_input=RandomParams(todo_title=workflow_input.todo_title),
32+
start_to_close_timeout=timedelta(seconds=120),
33+
)
34+
except Exception as e:
35+
error_message = f"Error during get_random: {e}"
36+
raise NonRetryableError(error_message) from e
37+
else:
38+
log.info("get_random done", random=random)
39+
40+
await workflow.sleep(2)
41+
42+
try:
43+
result = await workflow.step(
44+
function=get_result,
45+
function_input=ResultParams(
46+
todo_title=workflow_input.todo_title,
47+
todo_id=workflow_input.todo_id,
48+
),
49+
start_to_close_timeout=timedelta(seconds=120),
50+
)
51+
except Exception as e:
52+
error_message = f"Error during get_result: {e}"
53+
raise NonRetryableError(error_message) from e
54+
else:
55+
todo_details = TodoExecuteResponse(
56+
todo_id=workflow_input.todo_id,
57+
todo_title=workflow_input.todo_title,
58+
details=random,
59+
status=result.status,
60+
)
61+
log.info("TodoExecuteWorkflow done", result=todo_details)
62+
return todo_details

0 commit comments

Comments
 (0)