feat: edit inline with AI

frontend/src/components/shortcuts/renderShortcut.tsx

@@ -108,6 +108,15 @@ function prettyPrintHotkey(key: string): [label: string, symbol?: string] {
   return [lowerKey];
 }
 
+export function getSymbol(key: string) {
+  const platform = isPlatformMac() ? "mac" : "default";
+  const keyData = KEY_MAPPINGS[key.toLowerCase()];
+  if (keyData) {
+    return keyData.symbols[platform] || keyData.symbols.default;
+  }
+  return key;
+}
+
 interface KeyData {
   symbols: {
     mac?: string;
@@ -190,6 +199,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/cm.ts

@@ -56,7 +56,9 @@ import { historyCompartment } from "./editing/extensions";
 import { goToDefinitionBundle } from "./go-to-definition/extension";
 import type { HotkeyProvider } from "../hotkeys/hotkeys";
 import { lightTheme } from "./theme/light";
-
+import { promptPlugin } from "./prompt/prompt";
+import { requestEditCompletion } from "./prompt/request";
+import { getCurrentLanguageAdapter } from "./language/commands";
 export interface CodeMirrorSetupOpts {
   cellId: CellId;
   showPlaceholder: boolean;
@@ -79,6 +81,7 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     cellCodeCallbacks,
     keymapConfig,
     hotkeys,
+    enableAI,
   } = opts;
 
   return [
@@ -91,6 +94,20 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     basicBundle(opts),
     // Underline cmd+clickable placeholder
     goToDefinitionBundle(),
+    // AI prompt edit
+    enableAI
+      ? promptPlugin({
+          complete: (req) => {
+            return requestEditCompletion({
+              prompt: req.prompt,
+              selection: req.selection,
+              codeBefore: req.codeBefore,
+              code: req.editorView.state.doc.toString(),
+              language: getCurrentLanguageAdapter(req.editorView),
+            });
+          },
+        })
+      : [],
   ];
 };
 

frontend/src/core/codemirror/prompt/complete.ts

@@ -0,0 +1,51 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { API } from "@/core/network/api";
+import { asURL } from "@/utils/url";
+import type { LanguageAdapterType } from "../language/types";
+import { getCodes } from "../copilot/getCodes";
+
+/**
+ * Request to edit code with AI
+ */
+export async function requestEditCompletion(opts: {
+  prompt: string;
+  selection: string;
+  code: string;
+  codeBefore: string;
+  language: LanguageAdapterType;
+}): Promise<string> {
+  const currentCode = opts.code;
+
+  const otherCodes = getCodes(currentCode);
+  // Other code to include is the codeBefore and the other codes
+  const includeOtherCode = `${opts.codeBefore}\n${otherCodes}`;
+
+  const response = await fetch(asURL("api/ai/completion").toString(), {
+    method: "POST",
+    headers: API.headers(),
+    body: JSON.stringify({
+      prompt: opts.prompt,
+      code: opts.selection,
+      includeOtherCode: includeOtherCode,
+      language: opts.language,
+    }),
+  });
+
+  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);
+  }
+
+  return result;
+}