add simple auto fix button for sql cell

frontend/src/components/editor/errors/auto-fix.tsx

@@ -1,24 +1,29 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { useSetAtom } from "jotai";
 import { WrenchIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
+import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
+import { cn } from "@/utils/cn";
 
 export const AutoFixButton = ({
   errors,
   cellId,
+  className,
 }: {
   errors: MarimoError[];
   cellId: CellId;
   className?: string;
 }) => {
   const { createNewCell } = useCellActions();
   const autoFixes = errors.flatMap((error) => getAutoFixes(error));
+  const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
 
   if (autoFixes.length === 0) {
     return null;
@@ -33,7 +38,7 @@ export const AutoFixButton = ({
       <Button
         size="xs"
         variant="outline"
-        className="my-2 font-normal"
+        className={cn("my-2 font-normal", className)}
         onClick={() => {
           const editorView =
             store.get(notebookAtom).cellHandles[cellId].current?.editorView;
@@ -48,6 +53,7 @@ export const AutoFixButton = ({
             },
             editor: editorView,
             cellId: cellId,
+            setAiCompletionCell,
           });
           // Focus the editor
           editorView?.focus();

frontend/src/components/editor/output/MarimoErrorOutput.tsx

@@ -533,7 +533,13 @@ export const MarimoErrorOutput = ({
               </div>
             );
           })}
-          {cellId && <AutoFixButton errors={sqlErrors} cellId={cellId} />}
+          {cellId && (
+            <AutoFixButton
+              errors={sqlErrors}
+              cellId={cellId}
+              className="mt-2.5"
+            />
+          )}
         </div>,
       );
     }

frontend/src/core/errors/__tests__/errors.test.ts

@@ -42,6 +42,18 @@ describe("getAutoFixes", () => {
     expect(fixes[0].title).toBe("Fix: Add 'import numpy as np'");
   });
 
+  it("returns sql fix for sql-error error", () => {
+    const error: MarimoError = {
+      type: "sql-error",
+      msg: "syntax error",
+      sql_statement: "SELECT * FROM table",
+    };
+
+    const fixes = getAutoFixes(error);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("AI Fix: Fix the SQL error");
+  });
+
   it("returns no fixes for NameError with unknown import", () => {
     const error: MarimoError = {
       type: "exception",

frontend/src/core/errors/errors.ts

@@ -12,6 +12,10 @@ export interface AutoFix {
     addCodeBelow: (code: string) => void;
     editor: EditorView | undefined;
     cellId: CellId;
+    setAiCompletionCell?: (cell: {
+      cellId: CellId;
+      initialPrompt?: string;
+    }) => void;
   }) => Promise<void>;
 }
 
@@ -57,6 +61,21 @@ export function getAutoFixes(error: MarimoError): AutoFix[] {
     ];
   }
 
+  if (error.type === "sql-error") {
+    return [
+      {
+        title: "AI Fix: Fix the SQL error",
+        description: "Fix the SQL statement",
+        onFix: async (ctx) => {
+          ctx.setAiCompletionCell?.({
+            cellId: ctx.cellId,
+            initialPrompt: `Fix the SQL statement: ${error.msg}`,
+          });
+        },
+      },
+    ];
+  }
+
   return [];
 }