Implement error handling for extract values (#382)

* Remove chat streaming

* Improve error handling and logging in extract values and chat title steps

Co-authored-by: dallin <dallin@buster.so>

---------

Co-authored-by: Nate Kelley <nate@buster.so>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>

packages/ai/src/steps/extract-values-search-step.ts

@@ -35,7 +35,7 @@ export const extractValuesSearchOutputSchema = z.object({
         name: z.string(),
         versionNumber: z.number(),
         metricIds: z.array(z.string()),
-      })
+      }),
     )
     .optional(),
 });
@@ -93,7 +93,7 @@ Focus only on extracting meaningful, specific values that could be searched for
  * Organizes search results by schema.table structure with columns and their values
  */
 function organizeResultsBySchemaTable(
-  results: StoredValueResult[]
+  results: StoredValueResult[],
 ): Record<string, Record<string, string[]>> {
   const organized: Record<string, Record<string, string[]>> = {};
 
@@ -149,7 +149,7 @@ function formatSearchResults(results: StoredValueResult[]): string {
  */
 async function searchStoredValues(
   values: string[],
-  dataSourceId: string
+  dataSourceId: string,
 ): Promise<{
   searchResults: string;
   foundValues: Record<string, Record<string, string[]>>;
@@ -173,7 +173,7 @@ async function searchStoredValues(
       } catch (error) {
         console.error(
           `[StoredValues] Failed to generate embedding for "${value}":`,
-          error instanceof Error ? error.message : 'Unknown error'
+          error instanceof Error ? error.message : 'Unknown error',
         );
         return null;
       }
@@ -181,7 +181,7 @@ async function searchStoredValues(
 
     const embeddingResults = await Promise.all(embeddingPromises);
     const validEmbeddings = embeddingResults.filter(
-      (result): result is { value: string; embedding: number[] } => result !== null
+      (result): result is { value: string; embedding: number[] } => result !== null,
     );
 
     if (validEmbeddings.length === 0) {
@@ -201,7 +201,7 @@ async function searchStoredValues(
       } catch (error) {
         console.error(
           `[StoredValues] Failed to search stored values for "${value}":`,
-          error instanceof Error ? error.message : 'Unknown error'
+          error instanceof Error ? error.message : 'Unknown error',
         );
         return [];
       }
@@ -264,21 +264,34 @@ const extractValuesSearchStepExecution = async ({
       messages = standardizeMessages(prompt);
     }
 
-    const tracedValuesExtraction = wrapTraced(
-      async () => {
-        const response = await valuesAgent.generate(messages, {
-          maxSteps: 0,
-          output: extractValuesAgentOutputSchema, // Use the agent-specific schema
-        });
+    let extractedValues: { values: string[] } = { values: [] };
 
-        return response.object;
-      },
-      {
-        name: 'Extract Values',
-      }
-    );
+    try {
+      const tracedValuesExtraction = wrapTraced(
+        async () => {
+          const response = await valuesAgent.generate(messages, {
+            maxSteps: 0,
+            output: extractValuesSearchOutputSchema,
+          });
 
-    const extractedValues = await tracedValuesExtraction();
+          return response.object;
+        },
+        {
+          name: 'Extract Values',
+        },
+      );
+
+      extractedValues = await tracedValuesExtraction();
+    } catch (llmError) {
+      // Handle LLM generation errors specifically
+      console.warn('[ExtractValues] LLM failed to generate valid response:', {
+        error: llmError instanceof Error ? llmError.message : 'Unknown error',
+        errorType: llmError instanceof Error ? llmError.name : 'Unknown',
+      });
+
+      // Continue with empty values instead of failing
+      extractedValues = { values: [] };
+    }
 
     // Get dataSourceId from runtime context for stored values search
     const dataSourceId = runtimeContext.get('dataSourceId') as string | undefined;
@@ -296,6 +309,7 @@ const extractValuesSearchStepExecution = async ({
   } catch (error) {
     // Handle AbortError gracefully
     if (error instanceof Error && error.name === 'AbortError') {
+      console.info('[ExtractValues] Step was aborted');
       // Return empty values when aborted
       return {
         values: [],
@@ -306,13 +320,14 @@ const extractValuesSearchStepExecution = async ({
       };
     }
 
-    console.error('Failed to extract values:', error);
+    console.error('[ExtractValues] Unexpected error in step execution:', error);
     // Return empty values array instead of crashing
     return {
       values: [],
       searchResults: '',
       foundValues: {},
       searchPerformed: false,
+      dashboardFiles: inputData.dashboardFiles, // Pass through dashboard context
     };
   }
 };

packages/ai/src/steps/generate-chat-title-step.ts

@@ -14,12 +14,16 @@ const inputSchema = thinkAndPrepWorkflowInputSchema;
 export const generateChatTitleOutputSchema = z.object({
   title: z.string().describe('The title for the chat.'),
   // Pass through dashboard context
-  dashboardFiles: z.array(z.object({
-    id: z.string(),
-    name: z.string(),
-    versionNumber: z.number(),
-    metricIds: z.array(z.string()),
-  })).optional(),
+  dashboardFiles: z
+    .array(
+      z.object({
+        id: z.string(),
+        name: z.string(),
+        versionNumber: z.number(),
+        metricIds: z.array(z.string()),
+      }),
+    )
+    .optional(),
 });
 
 const generateChatTitleInstructions = `
@@ -54,21 +58,34 @@ const generateChatTitleExecution = async ({
       messages = standardizeMessages(prompt);
     }
 
-    const tracedChatTitle = wrapTraced(
-      async () => {
-        const response = await todosAgent.generate(messages, {
-          maxSteps: 0,
-          output: generateChatTitleOutputSchema,
-        });
+    let title: { title: string };
 
-        return response.object;
-      },
-      {
-        name: 'Generate Chat Title',
-      }
-    );
+    try {
+      const tracedChatTitle = wrapTraced(
+        async () => {
+          const response = await todosAgent.generate(messages, {
+            maxSteps: 0,
+            output: generateChatTitleOutputSchema,
+          });
 
-    const title = await tracedChatTitle();
+          return response.object;
+        },
+        {
+          name: 'Generate Chat Title',
+        },
+      );
+
+      title = await tracedChatTitle();
+    } catch (llmError) {
+      // Handle LLM generation errors specifically
+      console.warn('[GenerateChatTitle] LLM failed to generate valid response:', {
+        error: llmError instanceof Error ? llmError.message : 'Unknown error',
+        errorType: llmError instanceof Error ? llmError.name : 'Unknown',
+      });
+
+      // Continue with fallback title instead of failing
+      title = { title: 'New Analysis' };
+    }
 
     const chatId = runtimeContext.get('chatId');
     const messageId = runtimeContext.get('messageId');
@@ -80,7 +97,7 @@ const generateChatTitleExecution = async ({
       updatePromises.push(
         updateChat(chatId, {
           title: title.title,
-        })
+        }),
       );
     }
 
@@ -88,7 +105,7 @@ const generateChatTitleExecution = async ({
       updatePromises.push(
         updateMessage(messageId, {
           title: title.title,
-        })
+        }),
       );
     }
 
@@ -108,10 +125,11 @@ const generateChatTitleExecution = async ({
       };
     }
 
-    console.error('Failed to generate chat title:', error);
+    console.error('[GenerateChatTitle] Failed to generate chat title:', error);
     // Return a fallback title instead of crashing
     return {
       title: 'New Analysis',
+      dashboardFiles: inputData.dashboardFiles, // Pass through dashboard context
     };
   }
 };