Update agent_chat example

Author: leawn

agent_chat/pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "watchfiles>=1.0.4",
     "python-dotenv==1.0.1",
     "openai>=1.61.0",
-    "restack-ai>=0.0.77",
+    "restack-ai>=0.0.78",
 ]
 
 [project.scripts]

agent_chat/src/agents/agent.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import agent, import_functions, log, AgentError
+from restack_ai.agent import NonRetryableError, agent, import_functions, log
 
 with import_functions():
     from src.functions.llm_chat import LlmChatInput, Message, llm_chat
@@ -23,22 +23,22 @@ def __init__(self) -> None:
 
     @agent.event
     async def messages(self, messages_event: MessagesEvent) -> list[Message]:
+        log.info(f"Received messages: {messages_event.messages}")
+        self.messages.extend(messages_event.messages)
+
+        log.info(f"Calling llm_chat with messages: {self.messages}")
         try:
-            log.info(f"Received messages: {messages_event.messages}")
-            self.messages.extend(messages_event.messages)
-            
-            log.info(f"Calling llm_chat with messages: {self.messages}")
             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(assistant_message)
             return self.messages
-        except Exception as e:
-                log.error(f"Error in messages: {e}")
-                raise AgentError(f"Error in messages: {e}")
 
     @agent.event
     async def end(self, end: EndEvent) -> EndEvent: