返回 presentation-ai
ai-menu.tsx
根目录 / src / components / plate / ui / ai-menu.tsx
1 "use client";
2
3 import {
4 AIChatPlugin,
5 AIPlugin,
6 useEditorChat,
7 useLastAssistantMessage,
8 } from "@platejs/ai/react";
9 import { BlockSelectionPlugin, useIsSelecting } from "@platejs/selection/react";
10 import { getTransientSuggestionKey } from "@platejs/suggestion";
11 import { Command as CommandPrimitive } from "cmdk";
12 import {
13 Album,
14 BadgeHelp,
15 BookOpenCheck,
16 Check,
17 CornerUpLeft,
18 FeatherIcon,
19 ListEnd,
20 ListMinus,
21 ListPlus,
22 PauseIcon,
23 PenLine,
24 SmileIcon,
25 Wand,
26 X,
27 } from "lucide-react";
28 import {
29 isHotkey,
30 KEYS,
31 NodeApi,
32 type NodeEntry,
33 type SlateEditor,
34 } from "platejs";
35 import {
36 useEditorPlugin,
37 useEditorRef,
38 useHotkeys,
39 usePluginOption,
40 type PlateEditor,
41 } from "platejs/react";
42 import * as React from "react";
43
44 import { Button } from "@/components/plate/ui/button";
45 import {
46 Command,
47 CommandGroup,
48 CommandItem,
49 CommandList,
50 } from "@/components/plate/ui/command";
51 import {
52 Popover,
53 PopoverAnchor,
54 PopoverContent,
55 } from "@/components/plate/ui/popover";
56 import { AILoadingLabel } from "@/components/ui/ai-loading-label";
57 import { cn } from "@/lib/utils";
58 import { AIChatEditor } from "./ai-chat-editor";
59
60 type EditorChatState =
61 | "cursorCommand"
62 | "cursorSuggestion"
63 | "selectionCommand"
64 | "selectionSuggestion";
65
66 type AIMenuItemSelectArgs = {
67 aiEditor: SlateEditor;
68 editor: PlateEditor;
69 input: string;
70 };
71
72 function isRecord(value: unknown): value is Record<string, unknown> {
73 return typeof value === "object" && value !== null;
74 }
75
76 export function AIMenu() {
77 const { api, editor } = useEditorPlugin(AIChatPlugin);
78 const open = usePluginOption(AIChatPlugin, "open");
79 const mode = usePluginOption(AIChatPlugin, "mode");
80 const toolName = usePluginOption(AIChatPlugin, "toolName");
81 const streaming = usePluginOption(AIChatPlugin, "streaming");
82 const isSelecting = useIsSelecting();
83
84 const [value, setValue] = React.useState("");
85 const [input, setInput] = React.useState("");
86
87 const chat = usePluginOption(AIChatPlugin, "chat");
88 const { messages, status } = chat;
89 const [anchorElement, setAnchorElement] = React.useState<HTMLElement | null>(
90 null,
91 );
92
93 const content = useLastAssistantMessage()?.parts.find(
94 (part) => part.type === "text",
95 )?.text;
96
97 React.useEffect(() => {
98 if (!streaming) {
99 return;
100 }
101
102 const anchor = api.aiChat.node({ anchor: true });
103 if (!anchor?.[0]) {
104 setAnchorElement(null);
105 return;
106 }
107
108 const timeoutId = window.setTimeout(() => {
109 try {
110 const anchorDom = editor.api.toDOMNode(anchor[0]);
111 setAnchorElement(anchorDom ?? null);
112 } catch {
113 setAnchorElement(null);
114 }
115 }, 0);
116
117 return () => {
118 window.clearTimeout(timeoutId);
119 };
120 }, [api.aiChat, editor, streaming]);
121
122 const setPopoverOpen = (nextOpen: boolean) => {
123 if (nextOpen) {
124 api.aiChat.show();
125 return;
126 }
127
128 api.aiChat.hide();
129 };
130
131 const show = React.useCallback(
132 (nextAnchorElement: HTMLElement) => {
133 setAnchorElement(nextAnchorElement);
134 setPopoverOpen(true);
135 },
136 [api.aiChat],
137 );
138
139 useEditorChat({
140 onOpenBlockSelection: (blocks: NodeEntry[]) => {
141 const lastBlock = blocks.at(-1);
142 if (!lastBlock?.[0]) {
143 return;
144 }
145
146 const domNode = editor.api.toDOMNode(lastBlock[0]);
147 if (!domNode) {
148 return;
149 }
150
151 show(domNode);
152 },
153 onOpenChange: (nextOpen) => {
154 if (!nextOpen) {
155 setAnchorElement(null);
156 setInput("");
157 }
158 },
159 onOpenCursor: () => {
160 const ancestorEntry = editor.api.block({ highest: true });
161 if (!ancestorEntry) {
162 return;
163 }
164
165 const [ancestor] = ancestorEntry;
166
167 if (!editor.api.isAt({ end: true }) && !editor.api.isEmpty(ancestor)) {
168 editor
169 .getApi(BlockSelectionPlugin)
170 .blockSelection.set(ancestor.id as string);
171 }
172
173 const domNode = editor.api.toDOMNode(ancestor);
174 if (!domNode) {
175 return;
176 }
177
178 show(domNode);
179 },
180 onOpenSelection: () => {
181 const lastBlock = editor.api.blocks().at(-1);
182 if (!lastBlock?.[0]) {
183 return;
184 }
185
186 const domNode = editor.api.toDOMNode(lastBlock[0]);
187 if (!domNode) {
188 return;
189 }
190
191 show(domNode);
192 },
193 });
194
195 useHotkeys("esc", () => {
196 api.aiChat.stop();
197 });
198
199 const isLoading = status === "streaming" || status === "submitted";
200
201 React.useEffect(() => {
202 if (toolName !== "edit" || mode !== "chat" || isLoading) {
203 return;
204 }
205
206 let anchorNode = editor.api.node({
207 at: [],
208 reverse: true,
209 match: (node) =>
210 isRecord(node) &&
211 Boolean(node[KEYS.suggestion]) &&
212 Boolean(node[getTransientSuggestionKey()]),
213 });
214
215 if (!anchorNode) {
216 anchorNode = editor
217 .getApi(BlockSelectionPlugin)
218 .blockSelection.getNodes({ selectionFallback: true, sort: true })
219 .at(-1);
220 }
221
222 if (!anchorNode) {
223 return;
224 }
225
226 const block = editor.api.block({ at: anchorNode[1] });
227 if (!block?.[0]) {
228 return;
229 }
230
231 const domNode = editor.api.toDOMNode(block[0]);
232 if (!domNode) {
233 return;
234 }
235
236 setAnchorElement(domNode);
237 }, [editor, isLoading, mode, toolName]);
238
239 if (isLoading && mode === "insert") {
240 return null;
241 }
242
243 if (toolName === "edit" && mode === "chat" && isLoading) {
244 return null;
245 }
246
247 return (
248 <Popover open={open} onOpenChange={setPopoverOpen} modal={false}>
249 {anchorElement ? (
250 <PopoverAnchor virtualRef={{ current: anchorElement }} />
251 ) : null}
252
253 <PopoverContent
254 className="border-none bg-transparent p-0 shadow-none"
255 style={{
256 width: anchorElement?.offsetWidth,
257 }}
258 onEscapeKeyDown={(event) => {
259 event.preventDefault();
260 api.aiChat.hide();
261 }}
262 align="center"
263 side="bottom"
264 >
265 <Command
266 className="w-full rounded-lg border shadow-md"
267 value={value}
268 onValueChange={setValue}
269 >
270 {mode === "chat" &&
271 isSelecting &&
272 content &&
273 toolName === "generate" && <AIChatEditor content={content} />}
274
275 {isLoading ? (
276 <div className="flex grow p-2 select-none">
277 <AILoadingLabel
278 label={messages.length > 1 ? "Editing..." : "Thinking..."}
279 />
280 </div>
281 ) : (
282 <CommandPrimitive.Input
283 className={cn(
284 "flex h-9 w-full min-w-0 border-input bg-transparent px-3 py-1 text-base outline-hidden transition-[color,box-shadow] placeholder:text-muted-foreground md:text-sm dark:bg-input/30",
285 "aria-invalid:border-destructive aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40",
286 "border-b focus-visible:ring-transparent",
287 )}
288 value={input}
289 onKeyDown={(event) => {
290 if (isHotkey("backspace")(event) && input.length === 0) {
291 event.preventDefault();
292 api.aiChat.hide();
293 }
294
295 if (isHotkey("enter")(event) && !event.shiftKey && !value) {
296 event.preventDefault();
297 void api.aiChat.submit(input);
298 setInput("");
299 }
300 }}
301 onValueChange={setInput}
302 placeholder="Ask AI anything..."
303 data-plate-focus
304 autoFocus
305 />
306 )}
307
308 {!isLoading && (
309 <CommandList>
310 <AIMenuItems
311 input={input}
312 setInput={setInput}
313 setValue={setValue}
314 />
315 </CommandList>
316 )}
317 </Command>
318 </PopoverContent>
319 </Popover>
320 );
321 }
322
323 const aiChatItems = {
324 accept: {
325 icon: <Check />,
326 label: "Accept",
327 value: "accept",
328 onSelect: ({ aiEditor, editor }: AIMenuItemSelectArgs) => {
329 const { mode, toolName } = editor.getOptions(AIChatPlugin);
330
331 if (mode === "chat" && toolName === "generate") {
332 void editor
333 .getTransforms(AIChatPlugin)
334 .aiChat.replaceSelection(aiEditor);
335 return;
336 }
337
338 editor.getTransforms(AIChatPlugin).aiChat.accept();
339 editor.tf.focus({ edge: "end" });
340 },
341 },
342 continueWrite: {
343 icon: <PenLine />,
344 label: "Continue writing",
345 value: "continueWrite",
346 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
347 const ancestorNode = editor.api.block({ highest: true });
348
349 if (!ancestorNode) {
350 return;
351 }
352
353 const isEmpty = NodeApi.string(ancestorNode[0]).trim().length === 0;
354
355 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
356 mode: "insert",
357 prompt: isEmpty
358 ? `<Document>
359 {editor}
360 </Document>
361 Start writing a new paragraph AFTER <Document> ONLY ONE SENTENCE`
362 : "Continue writing AFTER <Block> ONLY ONE SENTENCE. DONT REPEAT THE TEXT.",
363 toolName: "generate",
364 });
365 },
366 },
367 discard: {
368 icon: <X />,
369 label: "Discard",
370 shortcut: "Escape",
371 value: "discard",
372 onSelect: ({ editor }: AIMenuItemSelectArgs) => {
373 editor.getTransforms(AIPlugin)?.ai?.undo();
374 editor.getApi(AIChatPlugin).aiChat.hide();
375 },
376 },
377 emojify: {
378 icon: <SmileIcon />,
379 label: "Emojify",
380 value: "emojify",
381 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
382 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
383 prompt:
384 "Add a small number of contextually relevant emojis within each block only. Do not rewrite the meaning or modify markdown structure.",
385 toolName: "edit",
386 });
387 },
388 },
389 explain: {
390 icon: <BadgeHelp />,
391 label: "Explain",
392 value: "explain",
393 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
394 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
395 prompt: {
396 default: "Explain {editor}",
397 selecting: "Explain",
398 },
399 toolName: "generate",
400 });
401 },
402 },
403 fixSpelling: {
404 icon: <Check />,
405 label: "Fix spelling & grammar",
406 value: "fixSpelling",
407 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
408 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
409 prompt:
410 "Fix spelling, grammar, and punctuation errors within each block only, without changing meaning.",
411 toolName: "edit",
412 });
413 },
414 },
415 generateMarkdownSample: {
416 icon: <BookOpenCheck />,
417 label: "Generate Markdown sample",
418 value: "generateMarkdownSample",
419 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
420 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
421 prompt: "Generate a markdown sample",
422 toolName: "generate",
423 });
424 },
425 },
426 generateMdxSample: {
427 icon: <BookOpenCheck />,
428 label: "Generate MDX sample",
429 value: "generateMdxSample",
430 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
431 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
432 prompt: "Generate a mdx sample",
433 toolName: "generate",
434 });
435 },
436 },
437 improveWriting: {
438 icon: <Wand />,
439 label: "Improve writing",
440 value: "improveWriting",
441 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
442 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
443 prompt:
444 "Improve the writing for clarity and flow without changing meaning or adding new information.",
445 toolName: "edit",
446 });
447 },
448 },
449 insertBelow: {
450 icon: <ListEnd />,
451 label: "Insert below",
452 value: "insertBelow",
453 onSelect: ({ aiEditor, editor }: AIMenuItemSelectArgs) => {
454 void editor
455 .getTransforms(AIChatPlugin)
456 .aiChat.insertBelow(aiEditor, { format: "none" });
457 },
458 },
459 makeLonger: {
460 icon: <ListPlus />,
461 label: "Make longer",
462 value: "makeLonger",
463 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
464 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
465 prompt:
466 "Make the content longer by elaborating on existing ideas within each block only, without changing meaning.",
467 toolName: "edit",
468 });
469 },
470 },
471 makeShorter: {
472 icon: <ListMinus />,
473 label: "Make shorter",
474 value: "makeShorter",
475 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
476 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
477 prompt:
478 "Make the content shorter by reducing verbosity within each block only, without changing meaning.",
479 toolName: "edit",
480 });
481 },
482 },
483 replace: {
484 icon: <Check />,
485 label: "Replace selection",
486 value: "replace",
487 onSelect: ({ aiEditor, editor }: AIMenuItemSelectArgs) => {
488 void editor.getTransforms(AIChatPlugin).aiChat.replaceSelection(aiEditor);
489 },
490 },
491 simplifyLanguage: {
492 icon: <FeatherIcon />,
493 label: "Simplify language",
494 value: "simplifyLanguage",
495 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
496 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
497 prompt:
498 "Simplify the language using clearer and more straightforward wording within each block only, without changing meaning.",
499 toolName: "edit",
500 });
501 },
502 },
503 summarize: {
504 icon: <Album />,
505 label: "Add a summary",
506 value: "summarize",
507 onSelect: ({ editor, input }: AIMenuItemSelectArgs) => {
508 void editor.getApi(AIChatPlugin).aiChat.submit(input, {
509 mode: "insert",
510 prompt: {
511 default: "Summarize {editor}",
512 selecting: "Summarize",
513 },
514 toolName: "generate",
515 });
516 },
517 },
518 tryAgain: {
519 icon: <CornerUpLeft />,
520 label: "Try again",
521 value: "tryAgain",
522 onSelect: ({ editor }: AIMenuItemSelectArgs) => {
523 void editor.getApi(AIChatPlugin).aiChat.reload();
524 },
525 },
526 } satisfies Record<
527 string,
528 {
529 icon: React.ReactNode;
530 label: string;
531 value: string;
532 component?: React.ComponentType<{ menuState: EditorChatState }>;
533 filterItems?: boolean;
534 items?: { label: string; value: string }[];
535 shortcut?: string;
536 onSelect?: (args: AIMenuItemSelectArgs) => void;
537 }
538 >;
539
540 const menuStateItems: Record<
541 EditorChatState,
542 {
543 items: (typeof aiChatItems)[keyof typeof aiChatItems][];
544 heading?: string;
545 }[]
546 > = {
547 cursorCommand: [
548 {
549 items: [
550 aiChatItems.generateMdxSample,
551 aiChatItems.generateMarkdownSample,
552 aiChatItems.continueWrite,
553 aiChatItems.summarize,
554 aiChatItems.explain,
555 ],
556 },
557 ],
558 cursorSuggestion: [
559 {
560 items: [aiChatItems.accept, aiChatItems.discard, aiChatItems.tryAgain],
561 },
562 ],
563 selectionCommand: [
564 {
565 items: [
566 aiChatItems.improveWriting,
567 aiChatItems.emojify,
568 aiChatItems.makeLonger,
569 aiChatItems.makeShorter,
570 aiChatItems.fixSpelling,
571 aiChatItems.simplifyLanguage,
572 ],
573 },
574 ],
575 selectionSuggestion: [
576 {
577 items: [
578 aiChatItems.accept,
579 aiChatItems.discard,
580 aiChatItems.insertBelow,
581 aiChatItems.tryAgain,
582 ],
583 },
584 ],
585 };
586
587 const AIMenuItems = ({
588 input,
589 setInput,
590 setValue,
591 }: {
592 input: string;
593 setInput: (value: string) => void;
594 setValue: (value: string) => void;
595 }) => {
596 const editor = useEditorRef();
597 const { messages } = usePluginOption(AIChatPlugin, "chat");
598 const aiEditor = usePluginOption(AIChatPlugin, "aiEditor");
599 const isSelecting = useIsSelecting();
600
601 const menuState = React.useMemo(() => {
602 if (messages && messages.length > 0) {
603 return isSelecting ? "selectionSuggestion" : "cursorSuggestion";
604 }
605
606 return isSelecting ? "selectionCommand" : "cursorCommand";
607 }, [isSelecting, messages]);
608
609 const menuGroups = React.useMemo(() => {
610 return menuStateItems[menuState];
611 }, [menuState]);
612
613 React.useEffect(() => {
614 const defaultItem = menuGroups[0]?.items[0];
615 if (!defaultItem) {
616 return;
617 }
618
619 setValue(defaultItem.value);
620 }, [menuGroups, setValue]);
621
622 if (!aiEditor) {
623 return null;
624 }
625
626 return (
627 <>
628 {menuGroups.map((group, index) => (
629 <CommandGroup key={index} heading={group.heading}>
630 {group.items.map((menuItem) => (
631 <CommandItem
632 key={menuItem.value}
633 className="[&_svg]:text-muted-foreground"
634 value={menuItem.value}
635 onSelect={() => {
636 menuItem.onSelect?.({
637 aiEditor,
638 editor,
639 input,
640 });
641 setInput("");
642 }}
643 >
644 {menuItem.icon}
645 <span>{menuItem.label}</span>
646 </CommandItem>
647 ))}
648 </CommandGroup>
649 ))}
650 </>
651 );
652 };
653
654 export function AILoadingBar() {
655 const toolName = usePluginOption(AIChatPlugin, "toolName");
656 const chat = usePluginOption(AIChatPlugin, "chat");
657 const mode = usePluginOption(AIChatPlugin, "mode");
658 const { status } = chat;
659 const { api } = useEditorPlugin(AIChatPlugin);
660
661 const isLoading = status === "streaming" || status === "submitted";
662 const visible =
663 isLoading &&
664 (mode === "insert" || (toolName === "edit" && mode === "chat"));
665
666 if (!visible) {
667 return null;
668 }
669
670 return (
671 <div
672 className={cn(
673 "absolute bottom-4 left-1/2 z-20 flex -translate-x-1/2 items-center gap-3 rounded-md border border-border bg-muted px-3 py-1.5 text-sm text-muted-foreground shadow-md transition-all duration-300",
674 )}
675 >
676 <AILoadingLabel
677 label={status === "submitted" ? "Thinking..." : "Writing..."}
678 icon={
679 <span className="size-4 animate-spin rounded-full border-2 border-muted-foreground border-t-transparent" />
680 }
681 />
682 <Button
683 size="sm"
684 variant="ghost"
685 className="flex items-center gap-1 text-xs"
686 onClick={() => api.aiChat.stop()}
687 >
688 <PauseIcon className="size-4" />
689 Stop
690 <kbd className="ml-1 rounded bg-border px-1 font-mono text-[10px] text-muted-foreground shadow-xs">
691 Esc
692 </kbd>
693 </Button>
694 </div>
695 );
696 }
697
697 lines Plain Text