Skip to content

Commit 5b1dd1c

Browse files
Light2DarkLight2Dark
committed
change to fix mode
1 parent a79ffd0 commit 5b1dd1c

File tree

4 files changed

+124
-57
lines changed

4 files changed

+124
-57
lines changed

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

Lines changed: 109 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,23 @@
11
/* Copyright 2024 Marimo. All rights reserved. */
22

3-
import { useAtomValue, useSetAtom } from "jotai";
4-
import { WrenchIcon, ZapIcon, ZapOffIcon } from "lucide-react";
3+
import { useAtomValue, useSetAtom, useStore } from "jotai";
4+
import { ChevronDownIcon, SparklesIcon, WrenchIcon } from "lucide-react";
55
import { Button } from "@/components/ui/button";
6-
import { Switch } from "@/components/ui/switch";
6+
import {
7+
DropdownMenu,
8+
DropdownMenuContent,
9+
DropdownMenuItem,
10+
DropdownMenuTrigger,
11+
} from "@/components/ui/dropdown-menu";
712
import { Tooltip } from "@/components/ui/tooltip";
813
import { aiCompletionCellAtom } from "@/core/ai/state";
914
import { notebookAtom, useCellActions } from "@/core/cells/cells";
1015
import type { CellId } from "@/core/cells/ids";
1116
import { aiEnabledAtom } from "@/core/config/config";
1217
import { getAutoFixes } from "@/core/errors/errors";
1318
import type { MarimoError } from "@/core/kernel/messages";
14-
import { store } from "@/core/state/jotai";
1519
import { cn } from "@/utils/cn";
16-
import { useInstantAIFix } from "./instant-fix";
20+
import { useFixMode } from "./fix-mode";
1721

1822
export const AutoFixButton = ({
1923
errors,
@@ -24,7 +28,7 @@ export const AutoFixButton = ({
2428
cellId: CellId;
2529
className?: string;
2630
}) => {
27-
const { instantAIFix, setInstantAIFix } = useInstantAIFix();
31+
const store = useStore();
2832
const { createNewCell } = useCellActions();
2933
const aiEnabled = useAtomValue(aiEnabledAtom);
3034
const autoFixes = errors.flatMap((error) =>
@@ -40,7 +44,7 @@ export const AutoFixButton = ({
4044
// multiple fixes.
4145
const firstFix = autoFixes[0];
4246

43-
const handleFix = () => {
47+
const handleFix = (triggerFix: boolean) => {
4448
const editorView =
4549
store.get(notebookAtom).cellHandles[cellId].current?.editorView;
4650
firstFix.onFix({
@@ -56,49 +60,118 @@ export const AutoFixButton = ({
5660
cellId: cellId,
5761
aiFix: {
5862
setAiCompletionCell,
59-
instantFix: instantAIFix,
63+
triggerFix,
6064
},
6165
});
6266
// Focus the editor
6367
editorView?.focus();
6468
};
6569

6670
return (
67-
<div className={cn("flex gap-2 my-2 items-center", className)}>
68-
<Tooltip content={firstFix.description} align="start">
71+
<div className={cn("my-2", className)}>
72+
{firstFix.fixType === "ai" ? (
73+
<AIFixButton
74+
tooltip={firstFix.description}
75+
openPrompt={() => handleFix(false)}
76+
applyAutofix={() => handleFix(true)}
77+
/>
78+
) : (
79+
<Tooltip content={firstFix.description} align="start">
80+
<Button
81+
size="xs"
82+
variant="outline"
83+
className="font-normal"
84+
onClick={() => handleFix(false)}
85+
>
86+
<WrenchIcon className="h-3 w-3 mr-2" />
87+
{firstFix.title}
88+
</Button>
89+
</Tooltip>
90+
)}
91+
</div>
92+
);
93+
};
94+
95+
const PromptIcon = SparklesIcon;
96+
const AutofixIcon = WrenchIcon;
97+
98+
const PromptTitle = "Suggest a prompt";
99+
const AutofixTitle = "Fix with AI";
100+
101+
export const AIFixButton = ({
102+
tooltip,
103+
openPrompt,
104+
applyAutofix,
105+
}: {
106+
tooltip: string;
107+
openPrompt: () => void;
108+
applyAutofix: () => void;
109+
}) => {
110+
const { fixMode, setFixMode } = useFixMode();
111+
112+
return (
113+
<div className="flex">
114+
<Tooltip content={tooltip} align="start">
69115
<Button
70116
size="xs"
71117
variant="outline"
72-
className="font-normal"
73-
onClick={handleFix}
118+
className="font-normal rounded-r-none border-r-0"
119+
onClick={fixMode === "prompt" ? openPrompt : applyAutofix}
74120
>
75-
<WrenchIcon className="h-3 w-3 mr-2" />
76-
{firstFix.title}
121+
{fixMode === "prompt" ? (
122+
<PromptIcon className="h-3 w-3 mr-2 mb-0.5" />
123+
) : (
124+
<AutofixIcon className="h-3 w-3 mr-2 mb-0.5" />
125+
)}
126+
{fixMode === "prompt" ? PromptTitle : AutofixTitle}
77127
</Button>
78128
</Tooltip>
79-
80-
{firstFix.fixType === "ai" && (
81-
<div className="flex items-center gap-2">
82-
<Switch
83-
checked={instantAIFix}
84-
onCheckedChange={() => setInstantAIFix(!instantAIFix)}
85-
size="sm"
86-
className="h-4 w-8"
87-
title="Toggle instant AI fix mode"
88-
/>
89-
<Tooltip
90-
content={
91-
instantAIFix ? "Instant fix enabled" : "Instant fix disabled"
92-
}
129+
<DropdownMenu>
130+
<DropdownMenuTrigger asChild={true}>
131+
<Button
132+
size="xs"
133+
variant="outline"
134+
className="rounded-l-none px-2"
135+
aria-label="Fix options"
93136
>
94-
{instantAIFix ? (
95-
<ZapIcon className="h-3 w-3 text-amber-500" />
96-
) : (
97-
<ZapOffIcon className="h-3 w-3 text-muted-foreground" />
98-
)}
99-
</Tooltip>
100-
</div>
101-
)}
137+
<ChevronDownIcon className="h-3 w-3" />
138+
</Button>
139+
</DropdownMenuTrigger>
140+
<DropdownMenuContent align="end" className="w-56">
141+
<DropdownMenuItem
142+
className="flex items-center gap-2"
143+
onClick={() => {
144+
setFixMode(fixMode === "prompt" ? "autofix" : "prompt");
145+
}}
146+
>
147+
<AiModeItem mode={fixMode === "prompt" ? "autofix" : "prompt"} />
148+
</DropdownMenuItem>
149+
</DropdownMenuContent>
150+
</DropdownMenu>
151+
</div>
152+
);
153+
};
154+
155+
const AiModeItem = ({ mode }: { mode: "prompt" | "autofix" }) => {
156+
const icon =
157+
mode === "prompt" ? (
158+
<PromptIcon className="h-4 w-4" />
159+
) : (
160+
<AutofixIcon className="h-4 w-4" />
161+
);
162+
const title = mode === "prompt" ? PromptTitle : AutofixTitle;
163+
const description =
164+
mode === "prompt"
165+
? "Edit the prompt before applying"
166+
: "Apply AI fixes automatically";
167+
168+
return (
169+
<div className="flex items-center gap-2">
170+
{icon}
171+
<div className="flex flex-col">
172+
<span className="font-medium">{title}</span>
173+
<span className="text-xs text-muted-foreground">{description}</span>
174+
</div>
102175
</div>
103176
);
104177
};

‎frontend/src/components/editor/errors/fix-mode.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
/* Copyright 2024 Marimo. All rights reserved. */
2+
3+
import { useAtom } from "jotai";
4+
import { atomWithStorage } from "jotai/utils";
5+
6+
const BASE_KEY = "marimo:ai-autofix-mode";
7+
8+
const fixModeAtom = atomWithStorage<"prompt" | "autofix">(BASE_KEY, "autofix");
9+
10+
export function useFixMode() {
11+
const [fixMode, setFixMode] = useAtom(fixModeAtom);
12+
return { fixMode, setFixMode };
13+
}

‎frontend/src/components/editor/errors/instant-fix.ts‎

Lines changed: 0 additions & 19 deletions
This file was deleted.

‎frontend/src/core/errors/errors.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { wrapInFunction } from "./utils";
99

1010
interface AIFix {
1111
setAiCompletionCell: (opts: AiCompletionCell) => void;
12-
instantFix: boolean;
12+
triggerFix: boolean;
1313
}
1414

1515
export interface AutoFix {
@@ -92,7 +92,7 @@ export function getAutoFixes(
9292
ctx.aiFix?.setAiCompletionCell({
9393
cellId: ctx.cellId,
9494
initialPrompt: initialPrompt,
95-
triggerImmediately: ctx.aiFix.instantFix,
95+
triggerImmediately: ctx.aiFix.triggerFix,
9696
});
9797
},
9898
},

0 commit comments

Comments
 (0)