Use correct error for functions

Author: leawn

agent_apis/src/functions/llm.py

@@ -3,7 +3,7 @@
 
 from dotenv import load_dotenv
 from openai import OpenAI
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 load_dotenv()
 
@@ -47,4 +47,4 @@ async def llm(function_input: FunctionInputParams) -> str:
         return response.choices[0].message.content
     except Exception as e:
         error_message = "llm function failed"
-        raise FunctionFailure(error_message, non_retryable=True) from e
+        raise NonRetryableError(error_message) from e

agent_chat/src/functions/llm_chat.py

@@ -4,7 +4,7 @@
 from dotenv import load_dotenv
 from openai import OpenAI
 from pydantic import BaseModel
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 load_dotenv()
 
@@ -22,7 +22,7 @@ class LlmChatInput(BaseModel):
 
 def raise_exception(message: str) -> None:
     log.error(message)
-    raise FunctionFailure(message, non_retryable=True)
+    raise NonRetryableError(message)
 
 
 @function.defn()
@@ -48,8 +48,8 @@ async def llm_chat(agent_input: LlmChatInput) -> dict[str, str]:
             messages=agent_input.messages,
         )
     except Exception as e:
-        log.error("llm_chat function failed", error=e)
-        raise
+        error_message = f"LLM chat failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         log.info(
             "llm_chat function completed", assistant_raw_response=assistant_raw_response

agent_rag/src/functions/llm_chat.py

@@ -5,7 +5,7 @@
 from openai import OpenAI
 from openai.types.chat.chat_completion import ChatCompletion
 from pydantic import BaseModel
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 load_dotenv()
 
@@ -23,7 +23,7 @@ class LlmChatInput(BaseModel):
 
 def raise_exception(message: str) -> None:
     log.error(message)
-    raise FunctionFailure(message, non_retryable=True)
+    raise NonRetryableError(message)
 
 
 @function.defn()
@@ -49,8 +49,8 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
             messages=function_input.messages,
         )
     except Exception as e:
-        log.error("llm_chat function failed", error=e)
-        raise
+        error_message = f"LLM chat failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         log.info("llm_chat function completed", response=response)
         return response

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

@@ -2,7 +2,7 @@
 
 from livekit import api
 from livekit.protocol.sip import CreateSIPParticipantRequest, SIPParticipantInfo
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 
 @dataclass
@@ -37,8 +37,7 @@ async def livekit_call(function_input: LivekitCallInput) -> SIPParticipantInfo:
 
         log.info("livekit_call SIPParticipantInfo:", participant=participant)
     except Exception as e:
-        log.error("livekit_call function failed", error=str(e))
-        failure_message = "livekit_call function failed: " + str(e)
-        raise FunctionFailure(failure_message, non_retryable=True) from e
+        error_message = f"livekit_call function failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         return participant

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

@@ -6,7 +6,7 @@
     ListSIPOutboundTrunkRequest,
     SIPOutboundTrunkInfo,
 )
-from restack_ai.function import FunctionFailure, function, function_info, log
+from restack_ai.function import NonRetryableError, function, function_info, log
 
 
 @function.defn()
@@ -45,9 +45,8 @@ async def livekit_outbound_trunk() -> str:
         await livekit_api.aclose()
 
     except Exception as e:  # Consider catching a more specific exception if possible
-        log.error("livekit_outbound_trunk function failed", error=str(e))
-        failure_message = "livekit_outbound_trunk function failed: " + str(e)
-        raise FunctionFailure(failure_message, non_retryable=True) from e
+        error_message = f"livekit_outbound_trunk function failed: {e}"
+        raise NonRetryableError(error_message) from e
 
     else:
         return trunk.sip_trunk_id

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

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass
 
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 from vapi import AsyncVapi, Call
 
 
@@ -26,8 +26,7 @@ async def vapi_call(function_input: VapiCallInput) -> Call:
         log.info("vapi_call: ", call=call)
 
     except Exception as e:
-        log.error("vapi_call function failed", error=str(e))
-        failure_message = "vapi_call function failed: " + str(e)
-        raise FunctionFailure(failure_message, non_retryable=True) from e
+        error_message = f"vapi_call function failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         return call

agent_todo/src/functions/llm_chat.py

@@ -9,7 +9,7 @@
 )
 from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
 from pydantic import BaseModel
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 load_dotenv()
 
@@ -30,7 +30,7 @@ class LlmChatInput(BaseModel):
 
 def raise_exception(message: str) -> None:
     log.error("llm_chat function failed", error=message)
-    raise FunctionFailure(message, non_retryable=True)
+    raise NonRetryableError(message)
 
 
 @function.defn()
@@ -58,8 +58,8 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
             tools=function_input.tools,
         )
     except Exception as e:
-        log.error("llm_chat function failed", error=e)
-        raise
+        error_message = f"LLM chat failed: {e}"
+        raise NonRetryableError(error_message) from e
     else:
         log.info("llm_chat function completed", response=response)
         return response

agent_tool/src/functions/llm_chat.py

@@ -9,7 +9,7 @@
 )
 from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
 from pydantic import BaseModel
-from restack_ai.function import FunctionFailure, function, log
+from restack_ai.function import NonRetryableError, function, log
 
 load_dotenv()
 
@@ -30,7 +30,7 @@ class LlmChatInput(BaseModel):
 
 def raise_exception(message: str) -> None:
     log.error(message)
-    raise FunctionFailure(message, non_retryable=True)
+    raise NonRetryableError(message)
 
 
 @function.defn()
@@ -58,5 +58,5 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
             tools=function_input.tools,
         )
     except Exception as e:
-        log.error("llm_chat function failed", error=e)
-        raise
+        error_message = f"LLM chat failed: {e}"
+        raise NonRetryableError(error_message) from e

pdf_ocr/src/functions/torch_ocr.py

@@ -7,7 +7,7 @@
 from numpy.typing import NDArray
 from PIL import Image
 from pydantic import BaseModel, Field
-from restack_ai.function import function, log, FunctionFailure
+from restack_ai.function import function, log, NonRetryableError
 
 from doctr.models import ocr_predictor
 import requests
@@ -42,14 +42,14 @@ async def torch_ocr(input: OcrInput) -> str:
             processed_img: NDArray[np.uint8] = service._preprocess_image(image)
             doc = DocumentFile.from_images(processed_img)
         else:
-            raise FunctionFailure("Unsupported file type", non_retryable=True)
+            raise NonRetryableError("Unsupported file type")
 
         result = service.predictor(doc)
         json_output = OCRPrediction.model_validate(result.export())
         return service._process_predictions(json_output)
     except Exception as e:
-        log.error(f"Failed to process file: {str(e)}")
-        raise FunctionFailure(f"Failed to process file: {str(e)}", non_retryable=True)
+        error_message = f"Failed to process file: {str(e)}"
+        raise NonRetryableError(error_message) from e
 
 class DocumentExtractionService:
     def __init__(self) -> None: