Update agent-rag example

Author: leawn

agent_rag/pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     "watchfiles>=1.0.4",
     "requests==2.32.3",
     "python-dotenv==1.0.1",
-    "restack-ai>=0.0.77",
+    "restack-ai>=0.0.78",
 ]
 
 [project.scripts]

agent_rag/src/agents/chat_rag.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import agent, import_functions, log
+from restack_ai.agent import agent, import_functions, log, NonRetryableError
 
 with import_functions():
     from src.functions.llm_chat import LlmChatInput, Message, llm_chat
@@ -26,30 +26,36 @@ def __init__(self) -> None:
     async def messages(self, messages_event: MessagesEvent) -> list[Message]:
         log.info(f"Received messages: {messages_event.messages}")
         self.messages.extend(messages_event.messages)
-
-        sales_info = await agent.step(
-            function=lookup_sales, start_to_close_timeout=timedelta(seconds=120)
-        )
-
-        system_content = f"You are a helpful assistant that can help with sales data. Here is the sales information: {sales_info}"
-
-        completion = await agent.step(
-            function=llm_chat,
-            function_input=LlmChatInput(
-                messages=self.messages, system_content=system_content
-            ),
-            start_to_close_timeout=timedelta(seconds=120),
-        )
-
-        log.info(f"completion: {completion}")
-
-        self.messages.append(
-            Message(
-                role="assistant", content=completion.choices[0].message.content or ""
+        try:
+            sales_info = await agent.step(
+                function=lookup_sales, start_to_close_timeout=timedelta(seconds=120)
             )
-        )
-
-        return self.messages
+        except Exception as e:
+            error_message = f"Error during lookup_sales: {e}"
+            raise NonRetryableError(error_message) from e
+        else:
+            system_content = f"You are a helpful assistant that can help with sales data. Here is the sales information: {sales_info}"
+
+            try:
+                completion = await agent.step(
+                    function=llm_chat,
+                    function_input=LlmChatInput(
+                        messages=self.messages, system_content=system_content
+                    ),
+                    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:
+                log.info(f"completion: {completion}")
+                self.messages.append(
+                    Message(
+                        role="assistant", content=completion.choices[0].message.content or ""
+                    )
+                )
+
+                return self.messages
 
     @agent.event
     async def end(self) -> EndEvent: