feat: (experimental) inline edits

frontend/package.json

@@ -38,6 +38,7 @@
     "@lezer/highlight": "^1.2.1",
     "@lezer/lr": "^1.4.2",
     "@lezer/python": "^1.1.15",
+    "@marimo-team/codemirror-ai": "^0.1.3",
     "@marimo-team/marimo-api": "file:../openapi",
     "@marimo-team/react-slotz": "^0.1.8",
     "@open-rpc/client-js": "^1.8.1",

frontend/src/components/app-config/user-config-form.tsx

@@ -1087,6 +1087,30 @@ export const UserConfigForm: React.FC = () => {
                 </div>
               )}
             />
+            <FormField
+              control={form.control}
+              name="experimental.inline_ai_tooltip"
+              render={({ field }) => (
+                <div className="flex flex-col gap-y-1">
+                  <FormItem className={formItemClasses}>
+                    <FormLabel className="font-normal">
+                      AI Edit Tooltip
+                    </FormLabel>
+                    <FormControl>
+                      <Checkbox
+                        data-testid="inline-ai-checkbox"
+                        checked={field.value === true}
+                        onCheckedChange={field.onChange}
+                      />
+                    </FormControl>
+                  </FormItem>
+                  <FormDescription>
+                    Enable experimental "Edit with AI" tooltip when selecting
+                    code.
+                  </FormDescription>
+                </div>
+              )}
+            />
             {!isWasm() && (
               <FormField
                 control={form.control}

frontend/src/components/shortcuts/renderShortcut.tsx

@@ -196,6 +196,10 @@ const KEY_MAPPINGS: Record<string, KeyData> = {
     symbols: { mac: "↘", default: "End" },
     label: "End",
   },
+  mod: {
+    symbols: { mac: "⌘", windows: "⊞ Win", default: "Ctrl" },
+    label: "Control",
+  },
 };
 
 function capitalize(str: string) {

frontend/src/core/codemirror/ai/request.ts

@@ -0,0 +1,76 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { API } from "@/core/network/api";
+import { asURL } from "@/utils/url";
+import type { LanguageAdapterType } from "../language/types";
+
+/**
+ * Request to edit code with AI
+ */
+export async function requestEditCompletion(opts: {
+  prompt: string;
+  selection: string;
+  codeBefore: string;
+  codeAfter: string;
+  language: LanguageAdapterType;
+}): Promise<string> {
+  // TODO: maybe include other code
+  // const otherCodes = getCodes(currentCode);
+
+  const finalPrompt = `
+Given the following code context, ${opts.prompt}
+
+SELECTED CODE:
+${opts.selection}
+
+CODE BEFORE SELECTION:
+${opts.codeBefore}
+
+CODE AFTER SELECTION:
+${opts.codeAfter}
+
+Instructions:
+1. Modify ONLY the selected code
+2. Keep the same indentation selected code
+3. Maintain consistent style with surrounding code
+4. Ensure the edit is complete and can be inserted directly
+5. Return ONLY the replacement code, no explanations, no code fences.
+
+Your task: ${opts.prompt}`;
+
+  const response = await fetch(asURL("api/ai/completion").toString(), {
+    method: "POST",
+    headers: API.headers(),
+    body: JSON.stringify({
+      prompt: finalPrompt,
+      code: "",
+      includeOtherCode: "",
+      language: opts.language,
+    }),
+  });
+
+  const firstLineIndent = opts.selection.match(/^\s*/)?.[0] || "";
+
+  const reader = response.body?.getReader();
+  if (!reader) {
+    throw new Error("Failed to get response reader");
+  }
+
+  let result = "";
+  // eslint-disable-next-line no-constant-condition
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) {
+      break;
+    }
+    result += new TextDecoder().decode(value);
+  }
+
+  // Add back the indent if it was stripped, which can happen with
+  // LLM responses
+  if (!result.startsWith(firstLineIndent)) {
+    result = firstLineIndent + result;
+  }
+
+  return result;
+}

frontend/src/core/codemirror/cm.ts

@@ -61,6 +61,11 @@ import { dndBundle } from "./misc/dnd";
 import { jupyterHelpExtension } from "./compat/jupyter";
 import { pasteBundle } from "./misc/paste";
 
+import { requestEditCompletion } from "./ai/request";
+import { getCurrentLanguageAdapter } from "./language/commands";
+import { aiExtension } from "@marimo-team/codemirror-ai";
+import { getFeatureFlag } from "../config/feature-flag";
+
 export interface CodeMirrorSetupOpts {
   cellId: CellId;
   showPlaceholder: boolean;
@@ -83,6 +88,7 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     cellCodeCallbacks,
     keymapConfig,
     hotkeys,
+    enableAI,
   } = opts;
 
   return [
@@ -98,6 +104,20 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     basicBundle(opts),
     // Underline cmd+clickable placeholder
     goToDefinitionBundle(),
+    // AI edit inline
+    enableAI && getFeatureFlag("inline_ai_tooltip")
+      ? aiExtension({
+          prompt: (req) => {
+            return requestEditCompletion({
+              prompt: req.prompt,
+              selection: req.selection,
+              codeBefore: req.codeBefore,
+              codeAfter: req.codeAfter,
+              language: getCurrentLanguageAdapter(req.editorView),
+            });
+          },
+        })
+      : [],
   ];
 };
 

frontend/src/core/config/feature-flag.tsx

@@ -8,6 +8,7 @@ import { getResolvedMarimoConfig } from "./config";
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface ExperimentalFeatures {
   markdown: boolean;
+  inline_ai_tooltip: boolean;
   wasm_layouts: boolean;
   scratchpad: boolean;
   chat_sidebar: boolean;
@@ -17,6 +18,7 @@ export interface ExperimentalFeatures {
 
 const defaultValues: ExperimentalFeatures = {
   markdown: true,
+  inline_ai_tooltip: import.meta.env.DEV,
   wasm_layouts: false,
   scratchpad: true,
   chat_sidebar: import.meta.env.DEV,

marimo/_server/errors.py

@@ -44,7 +44,7 @@ async def handle_error(request: Request, response: Any) -> Any:
     if isinstance(response, MarimoHTTPException):
         # Log server errors
         if not is_client_error(response.status_code):
-            LOGGER.exception(response)
+            LOGGER.error(response)
         return JSONResponse(
             {"detail": response.detail},
             status_code=response.status_code,

tests/_server/ai/test_prompts.py

@@ -102,6 +102,28 @@ def test_user_prompts():
     snapshot("user_prompts.txt", result)
 
 
+def test_edit_inline_prompts():
+    # Nothing fancy here, just making sure
+    # that if we already give a fully formatted prompt without
+    # any code, we don't add any more instructions
+    prompt = """
+    Given the following code context, {opts.prompt}
+
+    SELECTED CODE:
+    {opts.selection}
+
+    Instructions:
+    1. Modify ONLY the selected code`;
+    """
+
+    result = Prompter(code="").get_prompt(
+        user_prompt=prompt,
+        include_other_code="",
+    )
+
+    assert result == prompt
+
+
 def test_chat_system_prompts():
     result: str = ""
     result += _header("no custom rules")