From 13a65217b3167b9b09a7a98f2344e289e53cc613 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Thu, 6 Feb 2025 16:44:22 -0500 Subject: [PATCH 1/2] feat: (experimental) inline edits --- frontend/package.json | 1 + frontend/pnpm-lock.yaml | 14 ++++ .../app-config/user-config-form.tsx | 24 ++++++ .../components/shortcuts/renderShortcut.tsx | 4 + frontend/src/core/codemirror/ai/request.ts | 76 +++++++++++++++++++ frontend/src/core/codemirror/cm.ts | 20 +++++ frontend/src/core/config/feature-flag.tsx | 2 + tests/_server/ai/test_prompts.py | 22 ++++++ 8 files changed, 163 insertions(+) create mode 100644 frontend/src/core/codemirror/ai/request.ts diff --git a/frontend/package.json b/frontend/package.json index 2264781b40d..9cce977a214 100644 --- a/frontend/package.json +++ b/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", diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml index 208e959a746..376700a170b 100644 --- a/frontend/pnpm-lock.yaml +++ b/frontend/pnpm-lock.yaml @@ -90,6 +90,9 @@ importers: '@lezer/python': specifier: ^1.1.15 version: 1.1.15 + '@marimo-team/codemirror-ai': + specifier: ^0.1.3 + version: 0.1.3(@codemirror/state@6.5.2)(@codemirror/view@6.36.2) '@marimo-team/marimo-api': specifier: file:../openapi version: file:../openapi @@ -1678,6 +1681,12 @@ packages: '@marijn/find-cluster-break@1.0.2': resolution: {integrity: sha512-l0h88YhZFyKdXIFNfSWpyjStDjGHwZ/U7iobcK1cQQD8sejsONdQtTVU+1wVN1PBw40PiiHB1vA5S7VTfQiP9g==} + '@marimo-team/codemirror-ai@0.1.3': + resolution: {integrity: sha512-ZRtw9F8OGXVSpvGhGG38dAXcgumN/++gvfWPTig043kERbkDns99EbBHiPR5kmWWuTYIJLJtg9jFJKhnDSW9BQ==} + peerDependencies: + '@codemirror/state': ^6 + '@codemirror/view': ^6 + '@marimo-team/marimo-api@file:../openapi': resolution: {directory: ../openapi, type: directory} @@ -10391,6 +10400,11 @@ snapshots: '@marijn/find-cluster-break@1.0.2': {} + '@marimo-team/codemirror-ai@0.1.3(@codemirror/state@6.5.2)(@codemirror/view@6.36.2)': + dependencies: + '@codemirror/state': 6.5.2 + '@codemirror/view': 6.36.2 + '@marimo-team/marimo-api@file:../openapi': dependencies: openapi-fetch: 0.9.7 diff --git a/frontend/src/components/app-config/user-config-form.tsx b/frontend/src/components/app-config/user-config-form.tsx index 173684bd437..67239edca76 100644 --- a/frontend/src/components/app-config/user-config-form.tsx +++ b/frontend/src/components/app-config/user-config-form.tsx @@ -1087,6 +1087,30 @@ export const UserConfigForm: React.FC = () => { )} /> + ( +
+ + + AI Edit Tooltip + + + + + + + Enable experimental "Edit with AI" tooltip when selecting + code. + +
+ )} + /> {!isWasm() && ( = { symbols: { mac: "↘", default: "End" }, label: "End", }, + mod: { + symbols: { mac: "⌘", windows: "⊞ Win", default: "Ctrl" }, + label: "Control", + }, }; function capitalize(str: string) { diff --git a/frontend/src/core/codemirror/ai/request.ts b/frontend/src/core/codemirror/ai/request.ts new file mode 100644 index 00000000000..44fc7d5fe97 --- /dev/null +++ b/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 { + // 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; +} diff --git a/frontend/src/core/codemirror/cm.ts b/frontend/src/core/codemirror/cm.ts index fd59b3465b5..a072bf88922 100644 --- a/frontend/src/core/codemirror/cm.ts +++ b/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), + }); + }, + }) + : [], ]; }; diff --git a/frontend/src/core/config/feature-flag.tsx b/frontend/src/core/config/feature-flag.tsx index c58f0a80150..2d84356db42 100644 --- a/frontend/src/core/config/feature-flag.tsx +++ b/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, diff --git a/tests/_server/ai/test_prompts.py b/tests/_server/ai/test_prompts.py index 9eb841fcb58..d7eab90beca 100644 --- a/tests/_server/ai/test_prompts.py +++ b/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") From f57018af59a43d8a614f31025868ec70fc53e72a Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Thu, 6 Feb 2025 16:46:54 -0500 Subject: [PATCH 2/2] fix lint --- marimo/_server/errors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/marimo/_server/errors.py b/marimo/_server/errors.py index 99a88ee29c7..228f2ee3ae3 100644 --- a/marimo/_server/errors.py +++ b/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,