Skip to content

Commit 74d154c

Browse files
Light2DarkLight2Dark
authored
add floating banner for AI completions & support deletes (#6825)
## 📝 Summary <!-- Provide a concise summary of what this pull request is addressing. If this PR fixes any issues, list them here by number (e.g., Fixes #123). --> https://github.com/user-attachments/assets/05af6efe-65c2-4312-9571-f0cb88c378a6 ## 🔍 Description of Changes <!-- Detail the specific changes made in this pull request. Explain the problem addressed and how it was resolved. If applicable, provide before and after comparisons, screenshots, or any relevant details to help reviewers understand the changes easily. --> ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [x] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] I have added tests for the changes made. - [x] I have run the code and verified that it works as expected.
1 parent a827331 commit 74d154c

File tree

11 files changed

+307
-60
lines changed

11 files changed

+307
-60
lines changed

‎frontend/src/components/editor/ai/completion-handlers.tsx‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ export const CompletionActionsCellFooter: React.FC<{
114114
export const AcceptCompletionButton: React.FC<{
115115
isLoading: boolean;
116116
onAccept: () => void;
117+
text?: string;
117118
size?: "xs" | "sm";
118119
buttonStyles?: string;
119120
playButtonStyles?: string;
@@ -122,6 +123,7 @@ export const AcceptCompletionButton: React.FC<{
122123
}> = ({
123124
isLoading,
124125
onAccept,
126+
text = "Accept",
125127
size = "sm",
126128
buttonStyles,
127129
acceptShortcut,
@@ -150,7 +152,7 @@ export const AcceptCompletionButton: React.FC<{
150152
onClick={onAccept}
151153
className={`${baseClasses} rounded-r-none ${buttonStyles}`}
152154
>
153-
Accept
155+
{text}
154156
{acceptShortcut && (
155157
<MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
156158
)}
@@ -178,7 +180,7 @@ export const AcceptCompletionButton: React.FC<{
178180
onClick={onAccept}
179181
className={`${baseClasses} rounded px-3 ${buttonStyles}`}
180182
>
181-
Accept
183+
{text}
182184
{acceptShortcut && (
183185
<MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
184186
)}
@@ -188,10 +190,17 @@ export const AcceptCompletionButton: React.FC<{
188190

189191
export const RejectCompletionButton: React.FC<{
190192
onDecline: () => void;
193+
text?: string;
191194
size?: "xs" | "sm";
192195
className?: string;
193196
declineShortcut?: string;
194-
}> = ({ onDecline, size = "sm", className, declineShortcut }) => {
197+
}> = ({
198+
onDecline,
199+
text = "Reject",
200+
size = "sm",
201+
className,
202+
declineShortcut,
203+
}) => {
195204
return (
196205
<Button
197206
variant="text"
@@ -202,7 +211,7 @@ export const RejectCompletionButton: React.FC<{
202211
active:bg-(--red-5) dark:active:bg-(--red-4)
203212
border-(--red-6) border hover:shadow-xs ${className}`}
204213
>
205-
Reject
214+
{text}
206215
{declineShortcut && (
207216
<MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
208217
)}

‎frontend/src/components/editor/cell/StagedAICell.tsx‎

Lines changed: 72 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
/* Copyright 2024 Marimo. All rights reserved. */
22

33
import { useAtomValue, useStore } from "jotai";
4-
import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
4+
import {
5+
type Edit,
6+
stagedAICellsAtom,
7+
useStagedCells,
8+
} from "@/core/ai/staged-cells";
59
import { getCellEditorView } from "@/core/cells/cells";
610
import type { CellId } from "@/core/cells/ids";
711
import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
@@ -14,12 +18,18 @@ export const StagedAICellBackground: React.FC<{
1418
className?: string;
1519
}> = ({ cellId, className }) => {
1620
const stagedAICells = useAtomValue(stagedAICellsAtom);
21+
const stagedCell = stagedAICells.get(cellId);
1722

18-
if (!stagedAICells.has(cellId)) {
23+
if (!stagedCell) {
1924
return null;
2025
}
2126

22-
return <div className={cn("mo-ai-generated-cell", className)} />;
27+
const cellClass =
28+
stagedCell.type === "delete_cell"
29+
? "mo-ai-deleted-cell"
30+
: "mo-ai-generated-cell";
31+
32+
return <div className={cn(cellClass, className)} />;
2333
};
2434

2535
export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
@@ -35,43 +45,73 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
3545
return null;
3646
}
3747

38-
const handleAcceptCompletion = () => {
39-
removeStagedCell(cellId);
40-
};
41-
42-
const handleDeclineCompletion = () => {
43-
switch (stagedAiCell.type) {
44-
case "update_cell": {
45-
// Revert cell code
46-
const editorView = getCellEditorView(cellId);
47-
if (!editorView) {
48-
Logger.error("Editor for this cell not found", { cellId });
49-
break;
50-
}
51-
52-
updateEditorCodeFromPython(editorView, stagedAiCell.previousCode);
53-
removeStagedCell(cellId);
54-
break;
55-
}
56-
case "add_cell":
57-
// Delete the cell since it's newly created
58-
deleteStagedCell(cellId);
59-
break;
60-
case "delete_cell":
61-
// TODO: Revert delete
62-
removeStagedCell(cellId);
63-
break;
64-
}
48+
const handleCompletion = (type: "accept" | "reject") => {
49+
const completionFunc =
50+
type === "accept" ? acceptStagedCell : rejectStagedCell;
51+
completionFunc(cellId, stagedAiCell, removeStagedCell, deleteStagedCell);
6552
};
6653

6754
return (
6855
<div className="flex items-center justify-end gap-1.5 w-full pb-1 pt-2">
6956
<CompletionActionsCellFooter
7057
isLoading={false}
71-
onAccept={handleAcceptCompletion}
72-
onDecline={handleDeclineCompletion}
58+
onAccept={() => handleCompletion("accept")}
59+
onDecline={() => handleCompletion("reject")}
7360
size="xs"
7461
/>
7562
</div>
7663
);
7764
};
65+
66+
/**
67+
* Accept a staged cell and apply the changes.
68+
*/
69+
export function acceptStagedCell(
70+
cellId: CellId,
71+
edit: Edit,
72+
removeStagedCell: (cellId: CellId) => void,
73+
deleteStagedCell: (cellId: CellId) => void,
74+
): void {
75+
switch (edit.type) {
76+
case "delete_cell":
77+
// For delete cells, the cell is deleted when the completion is accepted
78+
deleteStagedCell(cellId);
79+
break;
80+
default:
81+
removeStagedCell(cellId);
82+
break;
83+
}
84+
}
85+
86+
/**
87+
* Reject a staged cell and revert the changes.
88+
*/
89+
export function rejectStagedCell(
90+
cellId: CellId,
91+
edit: Edit,
92+
removeStagedCell: (cellId: CellId) => void,
93+
deleteStagedCell: (cellId: CellId) => void,
94+
): void {
95+
switch (edit.type) {
96+
case "update_cell": {
97+
// Revert cell code
98+
const editorView = getCellEditorView(cellId);
99+
if (!editorView) {
100+
Logger.error("Editor for this cell not found", { cellId });
101+
break;
102+
}
103+
104+
updateEditorCodeFromPython(editorView, edit.previousCode);
105+
removeStagedCell(cellId);
106+
break;
107+
}
108+
case "add_cell":
109+
// Delete the cell since it's newly created
110+
deleteStagedCell(cellId);
111+
break;
112+
case "delete_cell":
113+
// Just remove the deletion marker - cell stays in notebook
114+
removeStagedCell(cellId);
115+
break;
116+
}
117+
}

‎frontend/src/components/editor/chrome/wrapper/app-chrome.tsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import { ContextAwarePanel } from "../panels/context-aware-panel/context-aware-p
2525
import { useChromeActions, useChromeState } from "../state";
2626
import { Minimap } from "./minimap";
2727
import { PanelsWrapper } from "./panels";
28+
import { PendingAICells } from "./pending-ai-cells";
2829
import { createStorage } from "./storage";
2930
import { handleDragging } from "./utils";
3031

@@ -272,6 +273,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
272273
<ContextAwarePanel />
273274
</PanelGroup>
274275
<Minimap />
276+
<PendingAICells />
275277
<ErrorBoundary>
276278
<TooltipProvider>
277279
<Footer />

‎frontend/src/components/editor/chrome/wrapper/minimap.tsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ const MinimapCell: React.FC<MinimapCellProps> = (props) => {
6464

6565
return (
6666
<button
67+
type="button"
6768
data-node-id={cell.id}
6869
className={cn(
6970
"group bg-transparent text-left w-full flex relative justify-between items-center",
@@ -107,6 +108,7 @@ const MinimapCell: React.FC<MinimapCellProps> = (props) => {
107108
width="1"
108109
height="1"
109110
>
111+
<title>Cell dependency connections for cell {cell.id}</title>
110112
{isSelected ? (
111113
<SelectedCell
112114
cell={cell}

‎frontend/src/components/editor/chrome/wrapper/pending-ai-cells.tsx‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
/* Copyright 2024 Marimo. All rights reserved. */
2+
3+
import { useAtomValue, useStore } from "jotai";
4+
import { ChevronDown, ChevronUp, SparklesIcon } from "lucide-react";
5+
import { useState } from "react";
6+
import { Button } from "@/components/ui/button";
7+
import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
8+
import type { CellId } from "@/core/cells/ids";
9+
import { getNextIndex } from "@/utils/arrays";
10+
import { cn } from "@/utils/cn";
11+
import {
12+
AcceptCompletionButton,
13+
RejectCompletionButton,
14+
} from "../../ai/completion-handlers";
15+
import { acceptStagedCell, rejectStagedCell } from "../../cell/StagedAICell";
16+
import { scrollAndHighlightCell } from "../../links/cell-link";
17+
18+
export const PendingAICells: React.FC = () => {
19+
const [currentIndex, setCurrentIndex] = useState<number | null>(null);
20+
21+
const stagedAiCells = useAtomValue(stagedAICellsAtom);
22+
const listStagedCells = [...stagedAiCells.keys()];
23+
const store = useStore();
24+
const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
25+
26+
if (stagedAiCells.size === 0) {
27+
return null;
28+
}
29+
30+
const scrollToCell = (cellId: CellId) => {
31+
scrollAndHighlightCell(cellId, "focus");
32+
};
33+
34+
const clickNext = (direction: "up" | "down") => {
35+
const newIndex = getNextIndex(
36+
currentIndex,
37+
listStagedCells.length,
38+
direction,
39+
);
40+
setCurrentIndex(newIndex);
41+
scrollToCell(listStagedCells[newIndex]);
42+
};
43+
44+
const acceptAllCompletions = () => {
45+
for (const [cellId, edit] of stagedAiCells) {
46+
acceptStagedCell(cellId, edit, removeStagedCell, deleteStagedCell);
47+
}
48+
};
49+
50+
const rejectAllCompletions = () => {
51+
for (const [cellId, edit] of stagedAiCells) {
52+
rejectStagedCell(cellId, edit, removeStagedCell, deleteStagedCell);
53+
}
54+
};
55+
56+
const cyanShadow = "shadow-[0_0_6px_0_#00A2C733]";
57+
58+
return (
59+
<div
60+
className={cn(
61+
"fixed bottom-16 left-1/2 transform -translate-x-1/2 z-50 bg-background/95 backdrop-blur-sm supports-backdrop-filter:bg-background/80 border border-border rounded-lg px-3 py-2 flex items-center justify-between gap-2.5 w-100",
62+
cyanShadow,
63+
)}
64+
>
65+
<SparklesIcon className="h-4 w-4 text-primary" />
66+
67+
<div className="flex items-center">
68+
<Button variant="ghost" size="icon" onClick={() => clickNext("up")}>
69+
<ChevronUp className="h-3.5 w-3.5" />
70+
</Button>
71+
<span className="text-xs font-mono min-w-[3.5rem] text-center">
72+
{currentIndex === null
73+
? `${listStagedCells.length} pending`
74+
: `${currentIndex + 1} / ${listStagedCells.length}`}
75+
</span>
76+
<Button variant="ghost" size="icon" onClick={() => clickNext("down")}>
77+
<ChevronDown className="h-3.5 w-3.5" />
78+
</Button>
79+
</div>
80+
81+
<div className="h-5 w-px bg-border" />
82+
83+
<div className="flex items-center gap-1.5">
84+
<AcceptCompletionButton
85+
text="Accept all"
86+
onAccept={acceptAllCompletions}
87+
isLoading={false}
88+
size="xs"
89+
buttonStyles="h-6.5"
90+
/>
91+
<RejectCompletionButton
92+
text="Reject all"
93+
onDecline={rejectAllCompletions}
94+
size="xs"
95+
className="h-6.5"
96+
/>
97+
</div>
98+
</div>
99+
);
100+
};

‎frontend/src/core/ai/staged-cells.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ import type { EditType } from "./tools/edit-notebook-tool";
2828
* And we only track one set of staged cells at a time.
2929
*/
3030

31-
type Edit =
31+
export type Edit =
3232
| { type: Extract<EditType, "update_cell">; previousCode: string }
3333
| { type: Extract<EditType, "add_cell"> }
3434
| { type: Extract<EditType, "delete_cell">; previousCode: string };

‎frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -333,7 +333,7 @@ describe("EditNotebookTool", () => {
333333
});
334334

335335
describe("delete_cell operation", () => {
336-
it.fails("should delete a cell", async () => {
336+
it("should delete a cell", async () => {
337337
const cellCode = "x = 1";
338338
const editorView = createMockEditorView(cellCode);
339339
const notebook = MockNotebook.notebookState({
@@ -363,7 +363,7 @@ describe("EditNotebookTool", () => {
363363
});
364364
});
365365

366-
it.fails("should throw error when cell ID doesn't exist", async () => {
366+
it("should throw error when cell ID doesn't exist", async () => {
367367
const notebook = MockNotebook.notebookState({
368368
cellData: {
369369
[cellId1]: { code: "x = 1" },
@@ -381,7 +381,7 @@ describe("EditNotebookTool", () => {
381381
).rejects.toThrow("Cell not found");
382382
});
383383

384-
it.fails("should throw error when cell editor not found", async () => {
384+
it("should throw error when cell editor not found", async () => {
385385
const notebook = MockNotebook.notebookState({
386386
cellData: {
387387
[cellId1]: { code: "x = 1" },

0 commit comments

Comments
 (0)