handle errors better

Author: adrianlyjak

src/{{ project_name_snake }}/qa_workflows.py

@@ -172,6 +172,12 @@ class QueryConversationHistoryEvent(HumanResponseEvent):
     pass
 
 
+class ErrorEvent(Event):
+    """Event emitted when an error occurs"""
+
+    error: str
+
+
 class ChatWorkflowState(BaseModel):
     index_name: str | None = None
     conversation_history: list[ConversationMessage] = Field(default_factory=list)
@@ -219,7 +225,7 @@ class ChatWorkflow(Workflow):
     @step
     async def initialize_chat(
         self, ev: ChatEvent, ctx: Context[ChatWorkflowState]
-    ) -> InputRequiredEvent | StopEvent:
+    ) -> InputRequiredEvent | StopEvent | None:
         """Initialize the chat session and request first input"""
         try:
             logger.info(f"Initializing chat {ev.index_name}")
@@ -236,17 +242,14 @@ async def initialize_chat(
             if ev.conversation_history:
                 async with ctx.store.edit_state() as state:
                     state.conversation_history.extend(ev.conversation_history)
-            # Request first user input
-            return InputRequiredEvent()
 
         except Exception as e:
             logger.error(f"Error initializing chat: {str(e)}", exc_info=True)
-            return StopEvent(
-                result={
-                    "success": False,
-                    "error": f"Failed to initialize chat: {str(e)}",
-                }
+            ctx.write_event_to_stream(
+                ErrorEvent(error=f"Failed to initialize chat: {str(e)}")
             )
+        return InputRequiredEvent()
+
 
     @step
     async def get_conversation_history(
@@ -260,7 +263,7 @@ async def get_conversation_history(
     @step
     async def process_user_response(
         self, ev: HumanResponseEvent, ctx: Context[ChatWorkflowState]
-    ) -> InputRequiredEvent | HumanResponseEvent | StopEvent | None:
+    ) -> InputRequiredEvent | HumanResponseEvent | None:
         """Process user input and generate response"""
         try:
             logger.info(f"Processing user response {ev.response}")
@@ -326,13 +329,11 @@ async def process_user_response(
                         assistant_response,
                     ]
                 )
-            return InputRequiredEvent()
 
         except Exception as e:
             logger.error(f"Error processing query: {str(e)}", exc_info=True)
-            return StopEvent(
-                result={"success": False, "error": f"Error processing query: {str(e)}"}
-            )
+            ctx.write_event_to_stream(ErrorEvent(error=str(e)))
+        return InputRequiredEvent()
 
 
 upload = DocumentUploadWorkflow(timeout=None)

test-proj/src/test_proj/qa_workflows.py

@@ -172,6 +172,12 @@ class QueryConversationHistoryEvent(HumanResponseEvent):
     pass
 
 
+class ErrorEvent(Event):
+    """Event emitted when an error occurs"""
+
+    error: str
+
+
 class ChatWorkflowState(BaseModel):
     index_name: str | None = None
     conversation_history: list[ConversationMessage] = Field(default_factory=list)
@@ -219,7 +225,7 @@ class ChatWorkflow(Workflow):
     @step
     async def initialize_chat(
         self, ev: ChatEvent, ctx: Context[ChatWorkflowState]
-    ) -> InputRequiredEvent | StopEvent:
+    ) -> InputRequiredEvent | StopEvent | None:
         """Initialize the chat session and request first input"""
         try:
             logger.info(f"Initializing chat {ev.index_name}")
@@ -236,17 +242,14 @@ async def initialize_chat(
             if ev.conversation_history:
                 async with ctx.store.edit_state() as state:
                     state.conversation_history.extend(ev.conversation_history)
-            # Request first user input
-            return InputRequiredEvent()
 
         except Exception as e:
             logger.error(f"Error initializing chat: {str(e)}", exc_info=True)
-            return StopEvent(
-                result={
-                    "success": False,
-                    "error": f"Failed to initialize chat: {str(e)}",
-                }
+            ctx.write_event_to_stream(
+                ErrorEvent(error=f"Failed to initialize chat: {str(e)}")
             )
+        return InputRequiredEvent()
+
 
     @step
     async def get_conversation_history(
@@ -260,7 +263,7 @@ async def get_conversation_history(
     @step
     async def process_user_response(
         self, ev: HumanResponseEvent, ctx: Context[ChatWorkflowState]
-    ) -> InputRequiredEvent | HumanResponseEvent | StopEvent | None:
+    ) -> InputRequiredEvent | HumanResponseEvent | None:
         """Process user input and generate response"""
         try:
             logger.info(f"Processing user response {ev.response}")
@@ -326,13 +329,11 @@ async def process_user_response(
                         assistant_response,
                     ]
                 )
-            return InputRequiredEvent()
 
         except Exception as e:
             logger.error(f"Error processing query: {str(e)}", exc_info=True)
-            return StopEvent(
-                result={"success": False, "error": f"Error processing query: {str(e)}"}
-            )
+            ctx.write_event_to_stream(ErrorEvent(error=str(e)))
+        return InputRequiredEvent()
 
 
 upload = DocumentUploadWorkflow(timeout=None)

test-proj/ui/src/libs/useChatbot.ts

@@ -166,6 +166,9 @@ export function useChatbot({
 interface AppendChatMessageData {
   message: ChatMessage;
 }
+interface ErrorEventData {
+  error: string;
+}
 interface ChatMessage {
   role: "user" | "assistant";
   text: string;
@@ -239,6 +242,22 @@ function toMessages(events: WorkflowEvent[]): Message[] {
         timestamp: new Date(content.message.timestamp),
         isPartial: false,
       });
+    } else if (type.endsWith(".ErrorEvent")) {
+      if (
+        lastMessage &&
+        lastMessage.isPartial &&
+        lastMessage.role === "assistant"
+      ) {
+        messages.pop();
+      }
+      const content = ev.data as unknown as ErrorEventData;
+      messages.push({
+        role: "assistant",
+        content: content.error,
+        timestamp: new Date(),
+        isPartial: false,
+        error: true,
+      });
     }
   }
   return messages;

ui/src/libs/useChatbot.ts

@@ -166,6 +166,9 @@ export function useChatbot({
 interface AppendChatMessageData {
   message: ChatMessage;
 }
+interface ErrorEventData {
+  error: string;
+}
 interface ChatMessage {
   role: "user" | "assistant";
   text: string;
@@ -239,6 +242,22 @@ function toMessages(events: WorkflowEvent[]): Message[] {
         timestamp: new Date(content.message.timestamp),
         isPartial: false,
       });
+    } else if (type.endsWith(".ErrorEvent")) {
+      if (
+        lastMessage &&
+        lastMessage.isPartial &&
+        lastMessage.role === "assistant"
+      ) {
+        messages.pop();
+      }
+      const content = ev.data as unknown as ErrorEventData;
+      messages.push({
+        role: "assistant",
+        content: content.error,
+        timestamp: new Date(),
+        isPartial: false,
+        error: true,
+      });
     }
   }
   return messages;