fix: handle case when no csv file selected

gregnr · gregnr · commit 08163cd67164 · 2024-08-10T15:19:32.000-05:00
diff --git a/apps/postgres-new/app/api/chat/route.ts b/apps/postgres-new/app/api/chat/route.ts
@@ -41,6 +41,9 @@ export async function POST(req: Request) {
       When querying data, limit to 5 by default.
 
       When performing FTS, always use 'simple' (languages aren't available).
+
+      When importing CSVs try to solve the problem yourself (eg. use a generic text column, then refine)
+      vs. asking the user to change the CSV.
       
       You also know math. All math equations and expressions must be written in KaTex and must be wrapped in double dollar \`$$\`:
         - Inline: $$\\sqrt{26}$$
diff --git a/apps/postgres-new/components/chat.tsx b/apps/postgres-new/components/chat.tsx
@@ -73,6 +73,28 @@ export default function Chat() {
 
   const sendCsv = useCallback(
     async (file: File) => {
+      if (file.type !== 'text/csv') {
+        // Add an artificial tool call requesting the CSV
+        // with an error indicating the file wasn't a CSV
+        appendMessage({
+          role: 'assistant',
+          content: '',
+          toolInvocations: [
+            {
+              state: 'result',
+              toolCallId: generateId(),
+              toolName: 'requestCsv',
+              args: {},
+              result: {
+                success: false,
+                error: `The file has type '${file.type}'. Let the user know that only CSV imports are currently supported.`,
+              },
+            },
+          ],
+        })
+        return
+      }
+
       const fileId = generateId()
 
       await saveFile(fileId, file)
@@ -120,7 +142,7 @@ export default function Chat() {
 
       const [file] = files
 
-      if (file && file.type === 'text/csv') {
+      if (file) {
         await sendCsv(file)
       }
     },
@@ -402,7 +424,7 @@ export default function Chat() {
                 const changeEvent = event as unknown as ChangeEvent<HTMLInputElement>
                 const [file] = Array.from(changeEvent.target?.files ?? [])
 
-                if (file && file.type === 'text/csv') {
+                if (file) {
                   await sendCsv(file)
                 }
 
diff --git a/apps/postgres-new/components/tools/csv-request.tsx b/apps/postgres-new/components/tools/csv-request.tsx
@@ -23,11 +23,20 @@ export default function CsvRequest({ toolInvocation }: CsvRequestProps) {
     const { result } = toolInvocation
 
     if (!result.success) {
-      return null
+      return (
+        <m.div
+          layout="position"
+          layoutId={toolInvocation.toolCallId}
+          className="self-end px-5 py-2.5 text-base rounded-full bg-destructive flex gap-2 items-center text-lighter italic"
+        >
+          No CSV file selected
+        </m.div>
+      )
     }
 
     return (
       <m.div
+        layout="position"
         layoutId={toolInvocation.toolCallId}
         className="self-end px-5 py-2.5 text-base rounded-full bg-border flex gap-2 items-center text-lighter italic"
         style={{
@@ -52,7 +61,7 @@ export default function CsvRequest({ toolInvocation }: CsvRequestProps) {
   }
 
   return (
-    <m.div layoutId={toolInvocation.toolCallId}>
+    <m.div layout="position" layoutId={toolInvocation.toolCallId}>
       <input
         type="file"
         onChange={async (e) => {
diff --git a/apps/postgres-new/components/workspace.tsx b/apps/postgres-new/components/workspace.tsx
@@ -7,7 +7,7 @@ import { useMessagesQuery } from '~/data/messages/messages-query'
 import { useTablesQuery } from '~/data/tables/tables-query'
 import { useOnToolCall } from '~/lib/hooks'
 import { useBreakpoint } from '~/lib/use-breakpoint'
-import { ensureMessageId } from '~/lib/util'
+import { ensureMessageId, ensureToolResult } from '~/lib/util'
 import { useApp } from './app-provider'
 import Chat, { getInitialMessages } from './chat'
 import IDE from './ide'
@@ -48,6 +48,7 @@ export default function Workspace({ databaseId, visibility, onStart }: Workspace
 
   const {
     messages,
+    setMessages,
     append,
     stop: stopReply,
   } = useChat({
@@ -84,11 +85,15 @@ export default function Workspace({ databaseId, visibility, onStart }: Workspace
 
   const appendMessage = useCallback(
     async (message: Message | CreateMessage) => {
+      setMessages((messages) => {
+        const isModified = ensureToolResult(messages)
+        return isModified ? [...messages] : messages
+      })
       ensureMessageId(message)
       append(message)
       saveMessage({ message })
     },
-    [saveMessage, append]
+    [setMessages, saveMessage, append]
   )
 
   const isConversationStarted =
diff --git a/apps/postgres-new/lib/util.ts b/apps/postgres-new/lib/util.ts
@@ -23,6 +23,34 @@ export function ensureMessageId(message: Message | CreateMessage): asserts messa
   }
 }
 
+/**
+ * Ensures that all tool invocations have a result before submitting,
+ * otherwise the LLM provider will return an error.
+ */
+export function ensureToolResult(messages: Message[]) {
+  let modified = false
+
+  for (const message of messages) {
+    if (!message.toolInvocations) {
+      continue
+    }
+
+    for (const toolInvocation of message.toolInvocations) {
+      if (!('result' in toolInvocation)) {
+        Object.assign(toolInvocation, {
+          result: {
+            success: false,
+            error: 'Failed to complete',
+          },
+        })
+        modified = true
+      }
+    }
+  }
+
+  return modified
+}
+
 /**
  * Checks if the message is a user message sent by the
  * application instead of the user.
