chore(ui): better assistant error handling (#1158)

src/leapfrogai_ui/src/lib/components/Message.svelte

@@ -59,7 +59,10 @@
   });
 
   let assistantImage = isRunAssistantMessage(message)
-    ? getAssistantImage($assistantsStore.assistants || [], message.assistant_id!)
+    ? getAssistantImage(
+        $assistantsStore.assistants || [],
+        message.assistant_id || message.metadata?.assistant_id
+      )
     : null;
 
   let messageIsHovered = false;
@@ -167,7 +170,9 @@
         >
           <div class="flex flex-col gap-2">
             <div class="font-bold">
-              {message.role === 'user' ? 'You' : getAssistantName(message.assistant_id)}
+              {message.role === 'user'
+                ? 'You'
+                : getAssistantName(message.assistant_id || message.metadata?.assistant_id)}
             </div>
             {#if fileMetadata}
               <div id="uploaded-files" class={'flex max-w-full  gap-2 overflow-x-auto bg-gray-900'}>

src/leapfrogai_ui/src/lib/constants/errors.ts

@@ -1,2 +1,4 @@
 export const FILE_CONTEXT_TOO_LARGE_ERROR_MSG = 'Error: Upload fewer or smaller files';
 export const ERROR_UPLOADING_FILE_MSG = 'Error uploading file';
+export const ASSISTANT_ERROR_MSG =
+  "I'm sorry but I've experienced an error. Please try again, or contact support.";

src/leapfrogai_ui/src/lib/constants/toastMessages.ts

@@ -19,7 +19,7 @@ export const ERROR_GETTING_ASSISTANT_MSG_TOAST = (
 ): ToastData => ({
   kind: 'error',
   title: 'Error',
-  subtitle: 'Error getting Assistant Response',
+  subtitle: 'Error getting assistant response',
   ...override
 });
 

src/leapfrogai_ui/src/lib/helpers/chatHelpers.ts

@@ -263,3 +263,11 @@ export const getCitations = (message: OpenAIMessage, files: FileObject[]) => {
   }
   return [];
 };
+
+export const refetchThread = async (threadId: string) => {
+  const res = await fetch(`/api/threads/${threadId}`);
+  if (res.ok) {
+    const thread = await res.json();
+    threadsStore.updateThread(thread);
+  }
+};

src/leapfrogai_ui/src/routes/api/chat/assistants/+server.ts

@@ -46,7 +46,6 @@ export const POST: RequestHandler = async ({ request, locals: { session } }) =>
             throw new Error('assistant_id is not set');
           })()
       });
-
       // forward run status would stream message deltas
       let runResult = await forwardStream(runStream);
 

src/leapfrogai_ui/src/routes/chat/(dashboard)/[[thread_id]]/+page.svelte

@@ -13,6 +13,7 @@
   import { twMerge } from 'tailwind-merge';
   import {
     isRunAssistantMessage,
+    refetchThread,
     resetMessages,
     saveMessage,
     stopThenSave
@@ -29,8 +30,12 @@
   import ChatFileUploadForm from '$components/ChatFileUpload.svelte';
   import FileChatActions from '$components/FileChatActions.svelte';
   import LFCarousel from '$components/LFCarousel.svelte';
+  import { ASSISTANT_ERROR_MSG } from '$constants/errors';
+  import { delay } from 'msw';
   import type { LFThread } from '$lib/types/threads';
 
+  export let data;
+
   /** LOCAL VARS **/
   let lengthInvalid: boolean; // bound to child LFTextArea
   let uploadingFiles = false;
@@ -69,6 +74,8 @@
     resetFiles(); // attachment of files w/assistants disabled
   }
 
+  $: if ($assistantError) handleAssistantResponseError();
+
   /** END REACTIVE STATE **/
 
   const handleThreadChange = () => {
@@ -134,14 +141,60 @@
 
   const handleCompletedAssistantResponse = async () => {
     if (componentHasMounted && $status === 'awaiting_message') {
-      const assistantResponseId = $assistantMessages[$assistantMessages.length - 1].id;
+      if ($assistantError) return;
+      if (latestAssistantMessage.role === 'user') {
+        await handleAssistantResponseError();
+        return;
+      }
+
+      const assistantResponseId = latestAssistantMessage.id;
       const messageRes = await fetch(
         `/api/messages?thread_id=${$page.params.thread_id}&message_id=${assistantResponseId}`
       );
+      if (!messageRes.ok) {
+        //useAssistants onError hook will handle this
+        return;
+      }
+
       const message = await messageRes.json();
-      await threadsStore.addMessageToStore(message);
-      threadsStore.setStreamingMessage(null);
+      if (message && !getMessageText(message)) {
+        // error with response(empty response)/timeout
+        await handleAssistantResponseError();
+      } else {
+        await threadsStore.addMessageToStore(message);
+        threadsStore.setStreamingMessage(null);
+      }
+    }
+  };
+
+  const createAssistantErrorResponse = async () => {
+    await delay(1000); // ensure error response timestamp is after user's msg
+    const newMessage = await saveMessage({
+      thread_id: data.thread.id,
+      content: ASSISTANT_ERROR_MSG,
+      role: 'assistant',
+      metadata: {
+        assistant_id: latestAssistantMessage.assistant_id || $threadsStore.selectedAssistantId
+      }
+    });
+
+    await threadsStore.addMessageToStore(newMessage);
+  };
+
+  const handleAssistantResponseError = async () => {
+    await refetchThread($page.params.thread_id); // if there was an error in the stream, we need to re-fetch to get the user's msg from the db
+    toastStore.addToast({
+      ...ERROR_GETTING_ASSISTANT_MSG_TOAST()
+    });
+    if (latestAssistantMessage.role === 'assistant') {
+      await threadsStore.deleteMessage($page.params.thread_id, latestAssistantMessage.id);
+      threadsStore.removeMessageFromStore($page.params.thread_id, latestAssistantMessage.id);
+      $assistantMessages = [...$assistantMessages.splice(-1)];
     }
+    await createAssistantErrorResponse();
+
+    threadsStore.setStreamingMessage(null);
+    await threadsStore.setSendingBlocked(false);
   };
 
   /** useChat - streams messages with the /api/chat route**/
@@ -193,19 +246,11 @@
     submitMessage: submitAssistantMessage,
     stop: assistantStop,
     setMessages: setAssistantMessages,
-    append: assistantAppend
+    append: assistantAppend,
+    error: assistantError
   } = useAssistant({
     api: '/api/chat/assistants',
-    threadId: activeThread?.id,
-    onError: async (e) => {
-      // ignore this error b/c it is expected on cancel
-      if (e.message !== 'BodyStreamBuffer was aborted') {
-        toastStore.addToast({
-          ...ERROR_GETTING_ASSISTANT_MSG_TOAST()
-        });
-      }
-      await threadsStore.setSendingBlocked(false);
-    }
+    threadId: activeThread?.id
   });
 
   const sendAssistantMessage = async (e: SubmitEvent | KeyboardEvent) => {

src/leapfrogai_ui/tests/assistant-progress.test.ts

@@ -9,6 +9,7 @@ import {
   uploadFileWithApi
 } from './helpers/fileHelpers';
 import { loadNewAssistantPage } from './helpers/navigationHelpers';
+import type { FileObject } from 'openai/resources/files';
 
 // Note - fully testing the assistant progress toast has proven difficult with Playwright. Sometimes the websocket
 // connection for the Supabase realtime listeners works, and sometimes it does not. Here we test that the
@@ -18,12 +19,17 @@ test('when creating an assistant with files, an assistant progress toast is disp
   openAIClient
 }) => {
   const assistantInput = getFakeAssistantInput();
-  const filename1 = `${faker.word.noun()}-test.pdf`;
-  const filename2 = `${faker.word.noun()}-test.pdf`;
-  await createPDF({ filename: filename1 });
-  await createPDF({ filename: filename2 });
-  const uploadedFile1 = await uploadFileWithApi(filename1, 'application/pdf', openAIClient);
-  const uploadedFile2 = await uploadFileWithApi(filename2, 'application/pdf', openAIClient);
+  const numFiles = 2;
+  const filenames: string[] = [];
+  const uploadedFiles: FileObject[] = [];
+
+  for (let i = 0; i < numFiles; i++) {
+    const filename = `${faker.word.noun()}-test.pdf`;
+    filenames.push(filename);
+    await createPDF({ filename });
+    const uploadedFile = await uploadFileWithApi(filename, 'application/pdf', openAIClient);
+    uploadedFiles.push(uploadedFile);
+  }
 
   await loadNewAssistantPage(page);
 
@@ -33,19 +39,36 @@ test('when creating an assistant with files, an assistant progress toast is disp
 
   await page.getByTestId('file-select-dropdown-btn').click();
   const fileSelectContainer = page.getByTestId('file-select-container');
-  await fileSelectContainer.getByTestId(`${uploadedFile1.id}-checkbox`).check();
-  await fileSelectContainer.getByTestId(`${uploadedFile2.id}-checkbox`).check();
+  for (const file of uploadedFiles) {
+    await fileSelectContainer.getByTestId(`${file.id}-checkbox`).check();
+  }
 
   await page.getByRole('button', { name: 'Save' }).click();
-  await page.waitForURL('/chat/assistants-management');
 
-  await expect(page.getByTestId(`file-${uploadedFile1.id}-vector-in-progress`)).toBeVisible();
-  await expect(page.getByTestId(`file-${uploadedFile2.id}-vector-pending`)).toBeVisible();
+  const inProgressSelector = `file-${uploadedFiles[0].id}-vector-in-progress`;
+  const completedSelector = `file-${uploadedFiles[0].id}-vector-completed`;
+
+  // Second file is pending
+  await expect(page.getByTestId(`file-${uploadedFiles[1].id}-vector-pending`)).toBeVisible();
+
+  // Check for either "in-progress" or "completed" state for the first file, it can happen really fast so this prevents
+  // a flaky test
+  const progressToast = await page.waitForSelector(
+    `[data-testid="${inProgressSelector}"], [data-testid="${completedSelector}"]`,
+    {
+      timeout: 30000
+    }
+  );
+  expect(progressToast).toBeTruthy();
+
+  await page.waitForURL('/chat/assistants-management');
 
   // cleanup
-  deleteFixtureFile(filename1);
-  deleteFixtureFile(filename2);
+  for (const filename of filenames) {
+    deleteFixtureFile(filename);
+  }
+  for (const file of uploadedFiles) {
+    await deleteFileWithApi(file.id, openAIClient);
+  }
   await deleteAssistantCard(assistantInput.name, page);
-  await deleteFileWithApi(uploadedFile1.id, openAIClient);
-  await deleteFileWithApi(uploadedFile2.id, openAIClient);
 });

src/leapfrogai_ui/tests/assistants.test.ts

@@ -16,6 +16,10 @@ import {
   loadChatPage,
   loadNewAssistantPage
 } from './helpers/navigationHelpers';
+import { ERROR_GETTING_ASSISTANT_MSG_TOAST } from '$constants/toastMessages';
+import { ASSISTANT_ERROR_MSG } from '$constants/errors';
+
+const newMessage1 = getSimpleMathQuestion();
 
 test('it navigates to the assistants page', async ({ page }) => {
   await loadChatPage(page);
@@ -342,3 +346,65 @@ test('displays a toast if there is an error deleting an assistant and response i
   //cleanup
   await deleteAssistantWithApi(assistant.id, openAIClient);
 });
+
+// Note - these error cases do not test all edge cases. ex. completed response comes back empty, /chat/assistants
+// partially completes then fails, stream fails, etc...
+test('displays an error toast if /chat/assistants throws while getting a response from an assistant', async ({
+  page,
+  openAIClient
+}) => {
+  const assistant = await createAssistantWithApi({ openAIClient });
+  await loadChatPage(page);
+
+  const assistantDropdown = page.getByTestId('assistants-select-btn');
+  await assistantDropdown.click();
+  await page.getByText(assistant!.name!).click();
+
+  await page.route('*/**/chat/assistants', async (route) => {
+    await route.abort('failed');
+  });
+  await sendMessage(page, newMessage1);
+
+  await expect(page.getByText(ERROR_GETTING_ASSISTANT_MSG_TOAST().title)).toBeVisible();
+  await expect(page.getByText(ASSISTANT_ERROR_MSG)).toBeVisible();
+});
+
+test('displays an error toast if /chat/assistants returns a 500 when getting a response from an assistant', async ({
+  page,
+  openAIClient
+}) => {
+  const assistant = await createAssistantWithApi({ openAIClient });
+  await loadChatPage(page);
+
+  const assistantDropdown = page.getByTestId('assistants-select-btn');
+  await assistantDropdown.click();
+  await page.getByText(assistant!.name!).click();
+
+  await page.route('*/**/chat/assistants', async (route) => {
+    await route.fulfill({ status: 500 });
+  });
+  await sendMessage(page, newMessage1);
+
+  await expect(page.getByText(ERROR_GETTING_ASSISTANT_MSG_TOAST().title)).toBeVisible();
+  await expect(page.getByText(ASSISTANT_ERROR_MSG)).toBeVisible();
+});
+
+test('displays an error toast if /chat/assistants returns a 200 with no body when getting a response from an assistant', async ({
+  page,
+  openAIClient
+}) => {
+  const assistant = await createAssistantWithApi({ openAIClient });
+  await loadChatPage(page);
+
+  const assistantDropdown = page.getByTestId('assistants-select-btn');
+  await assistantDropdown.click();
+  await page.getByText(assistant!.name!).click();
+
+  await page.route('*/**/chat/assistants', async (route) => {
+    await route.fulfill({ status: 200 });
+  });
+  await sendMessage(page, newMessage1);
+
+  await expect(page.getByText(ERROR_GETTING_ASSISTANT_MSG_TOAST().title)).toBeVisible();
+  await expect(page.getByText(ASSISTANT_ERROR_MSG)).toBeVisible();
+});