Final changes

Author: leawn

agent_apis/src/functions/weather.py

@@ -1,5 +1,5 @@
 import aiohttp
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function, log
 
 HTTP_OK = 200
 
@@ -21,6 +21,6 @@ async def weather() -> str:
                 return str(data)
             error_message = f"Error: {response.status}"
             raise_exception(error_message)
-    except Exception:
-        log.error("Error: {e}")
-        raise
+    except Exception as e:
+        error_message = f"Error: {e}"
+        raise NonRetryableError(error_message) from e

agent_apis/src/workflows/multistep.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel, Field
-from restack_ai.workflow import import_functions, log, workflow, NonRetryableError
+from restack_ai.workflow import NonRetryableError, import_functions, log, workflow
 
 with import_functions():
     from src.functions.llm import FunctionInputParams, llm

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, NonRetryableError
+from restack_ai.agent import NonRetryableError, agent, import_functions, log
 
 with import_functions():
     from src.functions.llm_chat import LlmChatInput, Message, llm_chat

agent_rag/src/functions/lookup_sales.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function, log
 
 
 class SalesItem(BaseModel):
@@ -85,5 +85,5 @@ async def lookup_sales() -> str:
 
         return str(items)
     except Exception as e:
-        log.error("lookup_sales function failed", error=e)
-        raise
+        error_message = f"lookup_sales function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_stream/src/functions/llm_chat.py

@@ -3,7 +3,7 @@
 
 from openai import OpenAI
 from pydantic import BaseModel, Field
-from restack_ai.function import function, log, stream_to_websocket
+from restack_ai.function import NonRetryableError, function, stream_to_websocket
 
 from src.client import api_address
 
@@ -49,5 +49,5 @@ async def llm_chat(function_input: LlmChatInput) -> str:
         return await stream_to_websocket(api_address=api_address, data=response)
 
     except Exception as e:
-        log.error("llm_chat function failed", error=str(e))
-        raise
+        error_message = f"llm_chat function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_telephony/twilio/agent_twilio/src/functions/livekit_dispatch.py

@@ -3,7 +3,7 @@
 
 from livekit import api
 from livekit.protocol.agent_dispatch import AgentDispatch
-from restack_ai.function import function, function_info, log
+from restack_ai.function import NonRetryableError, function, function_info
 
 
 @dataclass
@@ -37,8 +37,8 @@ async def livekit_dispatch(function_input: LivekitDispatchInput) -> AgentDispatc
         await lkapi.aclose()
 
     except Exception as e:
-        log.error("livekit_dispatch function failed", error=str(e))
-        raise
+        error_message = f"livekit_dispatch function failed: {e}"
+        raise NonRetryableError(error_message) from e
 
     else:
         return dispatch

agent_telephony/twilio/agent_twilio/src/functions/livekit_room.py

@@ -2,7 +2,7 @@
 
 from livekit import api
 from livekit.api import CreateRoomRequest, Room
-from restack_ai.function import function, function_info, log
+from restack_ai.function import NonRetryableError, function, function_info
 
 
 @function.defn()
@@ -27,8 +27,8 @@ async def livekit_room() -> Room:
         await lkapi.aclose()
 
     except Exception as e:
-        log.error("livekit_dispatch function failed", error=str(e))
-        raise
+        error_message = f"livekit_room function failed: {e}"
+        raise NonRetryableError(error_message) from e
 
     else:
         return room

agent_telephony/twilio/agent_twilio/src/functions/llm_chat.py

@@ -3,7 +3,7 @@
 
 from openai import OpenAI
 from pydantic import BaseModel, Field
-from restack_ai.function import function, log, stream_to_websocket
+from restack_ai.function import NonRetryableError, function, stream_to_websocket
 
 from src.client import api_address
 
@@ -48,5 +48,5 @@ async def llm_chat(function_input: LlmChatInput) -> str:
         return await stream_to_websocket(api_address=api_address, data=response)
 
     except Exception as e:
-        log.error("llm_chat function failed", error=str(e))
-        raise
+        error_message = f"llm_chat function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_telephony/vapi/agent_vapi/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
+from restack_ai.agent import NonRetryableError, agent, import_functions, log
 
 with import_functions():
     from src.functions.llm_chat import LlmChatInput, Message, llm_chat
@@ -32,27 +32,38 @@ 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"llm_chat function failed: {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 call(self, call_input: CallInput) -> None:
         log.info("Call", call_input=call_input)
         assistant_id = call_input.assistant_id
         phone_number = call_input.phone_number
 
-        return await agent.step(
-            function=vapi_call,
-            function_input=VapiCallInput(
-                assistant_id=assistant_id,
-                phone_number=phone_number,
-            ),
-        )
+        try:
+            result = agent.step(
+                function=vapi_call,
+                function_input=VapiCallInput(
+                    assistant_id=assistant_id,
+                    phone_number=phone_number,
+                ),
+            )
+        except Exception as e:
+            error_message = f"vapi_call function failed: {e}"
+            raise NonRetryableError(error_message) from e
+        else:
+            return result
 
     @agent.event
     async def end(self, end: EndEvent) -> EndEvent:

agent_telephony/vapi/agent_vapi/src/functions/llm_chat.py

@@ -3,7 +3,7 @@
 
 from openai import OpenAI
 from pydantic import BaseModel, Field
-from restack_ai.function import function, log, stream_to_websocket
+from restack_ai.function import NonRetryableError, function, stream_to_websocket
 
 from src.client import api_address
 
@@ -48,5 +48,5 @@ async def llm_chat(function_input: LlmChatInput) -> str:
         return await stream_to_websocket(api_address=api_address, data=response)
 
     except Exception as e:
-        log.error("llm_chat function failed", error=str(e))
-        raise
+        error_message = f"llm_chat function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_todo/src/functions/get_random.py

@@ -1,7 +1,7 @@
 import secrets
 
 from pydantic import BaseModel
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function
 
 
 class RandomParams(BaseModel):
@@ -13,7 +13,7 @@ async def get_random(params: RandomParams) -> str:
     try:
         random_number = secrets.randbelow(100)
     except Exception as e:
-        log.error("random function failed", error=e)
-        raise
+        error_message = f"get_random function failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         return f"The random number for {params.todo_title} is {random_number}."

agent_todo/src/functions/get_result.py

@@ -1,7 +1,7 @@
 import secrets
 
 from pydantic import BaseModel
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function
 
 
 class ResultParams(BaseModel):
@@ -20,5 +20,5 @@ async def get_result(params: ResultParams) -> ResultResponse:
         status = secrets.choice(["completed", "failed"])
         return ResultResponse(todo_id=params.todo_id, status=status)
     except Exception as e:
-        log.error("result function failed", error=e)
-        raise
+        error_message = f"get_result function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_todo/src/functions/todo_create.py

@@ -1,7 +1,7 @@
 import secrets
 
 from pydantic import BaseModel
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function, log
 
 
 class TodoCreateParams(BaseModel):
@@ -15,8 +15,8 @@ async def todo_create(params: TodoCreateParams) -> str:
 
         todo_id = f"todo-{secrets.randbelow(9000) + 1000}"
     except Exception as e:
-        log.error("todo_create function failed", error=e)
-        raise
+        error_message = f"todo_create function failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         log.info("todo_create function completed", todo_id=todo_id)
         return f"Created the todo '{params.title}' with id: {todo_id}"

agent_tool/src/functions/lookup_sales.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
 from pydantic import BaseModel
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function, log
 
 
 class SalesItem(BaseModel):
@@ -105,5 +105,5 @@ async def lookup_sales(function_input: LookupSalesInput) -> LookupSalesOutput:
 
         return LookupSalesOutput(sales=filtered_items)
     except Exception as e:
-        log.error("lookup_sales function failed", error=e)
-        raise
+        error_message = f"lookup_sales function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_video/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
+from restack_ai.agent import NonRetryableError, agent, import_functions, log
 
 with import_functions():
     from src.functions.context_docs import context_docs
@@ -26,13 +26,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"llm_chat function failed: {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:
@@ -42,11 +47,16 @@ async def end(self, end: EndEvent) -> EndEvent:
 
     @agent.run
     async def run(self) -> None:
-        docs = await agent.step(function=context_docs)
-        system_prompt=f"""
-        You are an interactive video assistant, your answers will be used in text to speech so try to keep answers short and concise so that interaction is seamless.
-        You can answer questions about the following documentation:
-        {docs}
-        """
-        self.messages.append(Message(role="system", content=system_prompt))
-        await agent.condition(lambda: self.end)
+        try:
+            docs = await agent.step(function=context_docs)
+        except Exception as e:
+            error_message = f"context_docs function failed: {e}"
+            raise NonRetryableError(error_message) from e
+        else:
+            system_prompt=f"""
+            You are an interactive video assistant, your answers will be used in text to speech so try to keep answers short and concise so that interaction is seamless.
+            You can answer questions about the following documentation:
+            {docs}
+            """
+            self.messages.append(Message(role="system", content=system_prompt))
+            await agent.condition(lambda: self.end)

agent_video/src/functions/context_docs.py

@@ -1,14 +1,14 @@
 import aiohttp
-from restack_ai.function import function, log
+from restack_ai.function import NonRetryableError, function, log
 
 
 async def fetch_content_from_url(url: str) -> str:
     async with aiohttp.ClientSession() as session:
         async with session.get(url) as response:
             if response.status == 200:
                 return await response.text()
-            log.error("Failed to fetch content", status=response.status)
-            raise Exception(f"Failed to fetch content: {response.status}")
+            error_message = f"Failed to fetch content: {response.status}"
+            raise NonRetryableError(error_message)
 
 
 @function.defn()
@@ -20,5 +20,5 @@ async def context_docs() -> str:
         return docs_content
 
     except Exception as e:
-        log.error("llm_chat function failed", error=str(e))
-        raise
+        error_message = f"context_docs function failed: {e}"
+        raise NonRetryableError(error_message) from e

agent_video/src/functions/llm_chat.py

@@ -3,7 +3,7 @@
 
 from openai import OpenAI
 from pydantic import BaseModel, Field
-from restack_ai.function import function, log, stream_to_websocket
+from restack_ai.function import NonRetryableError, function, stream_to_websocket
 
 from src.client import api_address
 
@@ -48,5 +48,5 @@ async def llm_chat(function_input: LlmChatInput) -> str:
         return await stream_to_websocket(api_address=api_address, data=response)
 
     except Exception as e:
-        log.error("llm_chat function failed", error=str(e))
-        raise
+        error_message = f"llm_chat function failed: {e}"
+        raise NonRetryableError(error_message) from e