Update agent-stream example

Author: leawn

agent_stream/pyproject.toml

@@ -10,8 +10,8 @@ dependencies = [
     "watchfiles>=1.0.4",
     "python-dotenv==1.0.1",
     "openai>=1.61.0",
-    "restack-ai>=0.0.77",
     "livekit-api>=0.8.2",
+    "restack-ai>=0.0.78",
 ]
 
 [project.scripts]
@@ -26,6 +26,7 @@ include = ["src"]
 [tool.hatch.build.targets.wheel]
 include = ["src"]
 
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

agent_stream/src/agents/agent.py

@@ -2,7 +2,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import agent, import_functions, log
+from restack_ai.agent import NonRetryableError, agent, import_functions, log
 
 with import_functions():
     from src.functions.llm_chat import LlmChatInput, Message, llm_chat
@@ -32,13 +32,18 @@ async def messages(self, messages_event: MessagesEvent) -> list[Message]:
         log.info(f"Received message: {messages_event.messages}")
         self.messages.extend(messages_event.messages)
 
-        assistant_message = await agent.step(
-            function=llm_chat,
-            function_input=LlmChatInput(messages=self.messages),
-            start_to_close_timeout=timedelta(seconds=120),
-        )
-        self.messages.append(Message(role="assistant", content=str(assistant_message)))
-        return self.messages
+        try:
+            assistant_message = await agent.step(
+                function=llm_chat,
+                function_input=LlmChatInput(messages=self.messages),
+                start_to_close_timeout=timedelta(seconds=120),
+            )
+        except Exception as e:
+            error_message = f"Error during llm_chat: {e}"
+            raise NonRetryableError(error_message) from e
+        else:
+            self.messages.append(Message(role="assistant", content=str(assistant_message)))
+            return self.messages
 
     @agent.event
     async def end(self, end: EndEvent) -> EndEvent: