AI

PreviousNext

AI 驱动的写作辅助。

Plus
Loading…

功能特性

  • 上下文感知命令菜单，可适应光标、文本选择和块选择工作流。
  • 流式 Markdown/MDX 插入，支持表格、列和代码块，由 streamInsertChunk 驱动。
  • 插入和聊天审查模式，通过局部回滚的插入预览以及 withAIBatch 和 tf.ai.undo() 实现撤销安全批处理。
  • 块选择感知变换，使用 tf.aiChat.replaceSelection 和 tf.aiChat.insertBelow 替换或追加整个部分。
  • 与 @ai-sdk/react 直接集成，使 api.aiChat.submit 可以从 Vercel AI SDK 助手流式传输响应。
  • 建议和评论工具，可对 AI 编辑进行差异比较、接受/拒绝更改，并将 AI 反馈映射回文档范围。

Kit 使用方法

安装

添加 AI 功能最快的方式是使用 AIKit。它包含已配置的 AIPlugin、AIChatPlugin、Markdown 流式助手、光标覆盖层及其 Plate UI 组件。

'use client';
 
import cloneDeep from 'lodash/cloneDeep.js';
import { BaseAIPlugin, withAIBatch } from '@platejs/ai';
import {
  AIChatPlugin,
  AIPlugin,
  applyAISuggestions,
  getInsertPreviewStart,
  streamInsertChunk,
  useChatChunk,
} from '@platejs/ai/react';
import { ElementApi, getPluginType, KEYS, PathApi } from 'platejs';
import { usePluginOption } from 'platejs/react';
 
import { AILoadingBar, AIMenu } from '@/components/ui/ai-menu';
import { AIAnchorElement, AILeaf } from '@/components/ui/ai-node';
 
import { useChat } from '../use-chat';
import { CursorOverlayKit } from './cursor-overlay-kit';
import { MarkdownKit } from './markdown-kit';
 
export const aiChatPlugin = AIChatPlugin.extend({
  options: {
    chatOptions: {
      api: '/api/ai/command',
      body: {},
    },
  },
  render: {
    afterContainer: AILoadingBar,
    afterEditable: AIMenu,
    node: AIAnchorElement,
  },
  shortcuts: { show: { keys: 'mod+j' } },
  useHooks: ({ editor, getOption }) => {
    useChat();
 
    const mode = usePluginOption(AIChatPlugin, 'mode');
    const toolName = usePluginOption(AIChatPlugin, 'toolName');
    useChatChunk({
      onChunk: ({ chunk, isFirst, nodes, text: content }) => {
        if (isFirst && mode === 'insert') {
          const { startBlock, startInEmptyParagraph } =
            getInsertPreviewStart(editor);
 
          editor.getTransforms(BaseAIPlugin).ai.beginPreview({
            originalBlocks:
              startInEmptyParagraph &&
              startBlock &&
              ElementApi.isElement(startBlock)
                ? [cloneDeep(startBlock)]
                : [],
          });
 
          editor.tf.withoutSaving(() => {
            editor.tf.insertNodes(
              {
                children: [{ text: '' }],
                type: getPluginType(editor, KEYS.aiChat),
              },
              {
                at: PathApi.next(editor.selection!.focus.path.slice(0, 1)),
              }
            );
          });
          editor.setOption(AIChatPlugin, 'streaming', true);
        }
 
        if (mode === 'insert' && nodes.length > 0) {
          editor.tf.withoutSaving(() => {
            if (!getOption('streaming')) return;
 
            editor.tf.withScrolling(() => {
              streamInsertChunk(editor, chunk, {
                textProps: {
                  [getPluginType(editor, KEYS.ai)]: true,
                },
              });
            });
          });
        }
 
        if (toolName === 'edit' && mode === 'chat') {
          withAIBatch(
            editor,
            () => {
              applyAISuggestions(editor, content);
            },
            {
              split: isFirst,
            }
          );
        }
      },
      onFinish: () => {
        editor.getApi(AIChatPlugin).aiChat.stop();
      },
    });
  },
});
 
export const AIKit = [
  ...CursorOverlayKit,
  ...MarkdownKit,
  AIPlugin.withComponent(AILeaf),
  aiChatPlugin,
];
'use client';
 
import cloneDeep from 'lodash/cloneDeep.js';
import { BaseAIPlugin, withAIBatch } from '@platejs/ai';
import {
  AIChatPlugin,
  AIPlugin,
  applyAISuggestions,
  getInsertPreviewStart,
  streamInsertChunk,
  useChatChunk,
} from '@platejs/ai/react';
import { ElementApi, getPluginType, KEYS, PathApi } from 'platejs';
import { usePluginOption } from 'platejs/react';
 
import { AILoadingBar, AIMenu } from '@/components/ui/ai-menu';
import { AIAnchorElement, AILeaf } from '@/components/ui/ai-node';
 
import { useChat } from '../use-chat';
import { CursorOverlayKit } from './cursor-overlay-kit';
import { MarkdownKit } from './markdown-kit';
 
export const aiChatPlugin = AIChatPlugin.extend({
  options: {
    chatOptions: {
      api: '/api/ai/command',
      body: {},
    },
  },
  render: {
    afterContainer: AILoadingBar,
    afterEditable: AIMenu,
    node: AIAnchorElement,
  },
  shortcuts: { show: { keys: 'mod+j' } },
  useHooks: ({ editor, getOption }) => {
    useChat();
 
    const mode = usePluginOption(AIChatPlugin, 'mode');
    const toolName = usePluginOption(AIChatPlugin, 'toolName');
    useChatChunk({
      onChunk: ({ chunk, isFirst, nodes, text: content }) => {
        if (isFirst && mode === 'insert') {
          const { startBlock, startInEmptyParagraph } =
            getInsertPreviewStart(editor);
 
          editor.getTransforms(BaseAIPlugin).ai.beginPreview({
            originalBlocks:
              startInEmptyParagraph &&
              startBlock &&
              ElementApi.isElement(startBlock)
                ? [cloneDeep(startBlock)]
                : [],
          });
 
          editor.tf.withoutSaving(() => {
            editor.tf.insertNodes(
              {
                children: [{ text: '' }],
                type: getPluginType(editor, KEYS.aiChat),
              },
              {
                at: PathApi.next(editor.selection!.focus.path.slice(0, 1)),
              }
            );
          });
          editor.setOption(AIChatPlugin, 'streaming', true);
        }
 
        if (mode === 'insert' && nodes.length > 0) {
          editor.tf.withoutSaving(() => {
            if (!getOption('streaming')) return;
 
            editor.tf.withScrolling(() => {
              streamInsertChunk(editor, chunk, {
                textProps: {
                  [getPluginType(editor, KEYS.ai)]: true,
                },
              });
            });
          });
        }
 
        if (toolName === 'edit' && mode === 'chat') {
          withAIBatch(
            editor,
            () => {
              applyAISuggestions(editor, content);
            },
            {
              split: isFirst,
            }
          );
        }
      },
      onFinish: () => {
        editor.getApi(AIChatPlugin).aiChat.stop();
      },
    });
  },
});
 
export const AIKit = [
  ...CursorOverlayKit,
  ...MarkdownKit,
  AIPlugin.withComponent(AILeaf),
  aiChatPlugin,
];
  • AIMenu：用于提示、工具快捷方式和聊天审查的浮动命令界面。
  • AILoadingBar：在编辑器容器中显示流式状态。
  • AIAnchorElement：用于在流式传输期间定位浮动菜单的不可见锚点节点。
  • AILeaf：以微妙的样式渲染 AI 标记的文本。

添加 Kit

import { createPlateEditor } from 'platejs/react';
import { AIKit } from '@/components/editor/plugins/ai-kit';
 
const editor = createPlateEditor({
  plugins: [
    // ...otherPlugins,
    ...AIKit,
  ],
});
import { createPlateEditor } from 'platejs/react';
import { AIKit } from '@/components/editor/plugins/ai-kit';
 
const editor = createPlateEditor({
  plugins: [
    // ...otherPlugins,
    ...AIKit,
  ],
});

添加 API 路由

暴露一个流式命令端点，代理您的模型提供商：

import type {
  ChatMessage,
  ToolName,
} from '@/components/editor/use-chat';
import type { NextRequest } from 'next/server';
 
import { createGateway } from '@ai-sdk/gateway';
import {
  type LanguageModel,
  type UIMessageStreamWriter,
  createUIMessageStream,
  createUIMessageStreamResponse,
  generateText,
  Output,
  streamText,
  tool,
} from 'ai';
import { NextResponse } from 'next/server';
import { type SlateEditor, createSlateEditor, nanoid } from 'platejs';
import { z } from 'zod';
 
import { BaseEditorKit } from '@/components/editor/editor-base-kit';
import { markdownJoinerTransform } from '@/lib/markdown-joiner-transform';
 
import {
  buildEditTableMultiCellPrompt,
  getChooseToolPrompt,
  getCommentPrompt,
  getEditPrompt,
  getGeneratePrompt,
} from './prompt';
 
export async function POST(req: NextRequest) {
  const { apiKey: key, ctx, messages: messagesRaw, model } = await req.json();
 
  const { children, selection, toolName: toolNameParam } = ctx;
 
  const editor = createSlateEditor({
    plugins: BaseEditorKit,
    selection,
    value: children,
  });
 
  const apiKey = key || process.env.AI_GATEWAY_API_KEY;
 
  if (!apiKey) {
    return NextResponse.json(
      { error: 'Missing AI Gateway API key.' },
      { status: 401 }
    );
  }
 
  const isSelecting = editor.api.isExpanded();
 
  const gatewayProvider = createGateway({
    apiKey,
  });
 
  try {
    const stream = createUIMessageStream<ChatMessage>({
      execute: async ({ writer }) => {
        let toolName = toolNameParam;
 
        if (!toolName) {
          const prompt = getChooseToolPrompt({
            isSelecting,
            messages: messagesRaw,
          });
 
          const enumOptions = isSelecting
            ? ['generate', 'edit', 'comment']
            : ['generate', 'comment'];
          const modelId = model || 'google/gemini-2.5-flash';
 
          const { output: AIToolName } = await generateText({
            model: gatewayProvider(modelId),
            output: Output.choice({ options: enumOptions }),
            prompt,
          });
 
          writer.write({
            data: AIToolName as ToolName,
            type: 'data-toolName',
          });
 
          toolName = AIToolName;
        }
 
        const stream = streamText({
          experimental_transform: markdownJoinerTransform(),
          model: gatewayProvider(model || 'openai/gpt-4o-mini'),
          // Not used
          prompt: '',
          tools: {
            comment: getCommentTool(editor, {
              messagesRaw,
              model: gatewayProvider(model || 'google/gemini-2.5-flash'),
              writer,
            }),
            table: getTableTool(editor, {
              messagesRaw,
              model: gatewayProvider(model || 'google/gemini-2.5-flash'),
              writer,
            }),
          },
          prepareStep: async (step) => {
            if (toolName === 'comment') {
              return {
                ...step,
                toolChoice: { toolName: 'comment', type: 'tool' },
              };
            }
 
            if (toolName === 'edit') {
              const [editPrompt, editType] = getEditPrompt(editor, {
                isSelecting,
                messages: messagesRaw,
              });
 
              // Table editing uses the table tool
              if (editType === 'table') {
                return {
                  ...step,
                  toolChoice: { toolName: 'table', type: 'tool' },
                };
              }
 
              return {
                ...step,
                activeTools: [],
                model:
                  editType === 'selection'
                    ? //The selection task is more challenging, so we chose to use Gemini 2.5 Flash.
                      gatewayProvider(model || 'google/gemini-2.5-flash')
                    : gatewayProvider(model || 'openai/gpt-4o-mini'),
                messages: [
                  {
                    content: editPrompt,
                    role: 'user',
                  },
                ],
              };
            }
 
            if (toolName === 'generate') {
              const generatePrompt = getGeneratePrompt(editor, {
                isSelecting,
                messages: messagesRaw,
              });
 
              return {
                ...step,
                activeTools: [],
                messages: [
                  {
                    content: generatePrompt,
                    role: 'user',
                  },
                ],
                model: gatewayProvider(model || 'openai/gpt-4o-mini'),
              };
            }
          },
        });
 
        writer.merge(stream.toUIMessageStream({ sendFinish: false }));
      },
    });
 
    return createUIMessageStreamResponse({ stream });
  } catch {
    return NextResponse.json(
      { error: 'Failed to process AI request' },
      { status: 500 }
    );
  }
}
 
const getCommentTool = (
  editor: SlateEditor,
  {
    messagesRaw,
    model,
    writer,
  }: {
    messagesRaw: ChatMessage[];
    model: LanguageModel;
    writer: UIMessageStreamWriter<ChatMessage>;
  }
) =>
  tool({
    description: 'Comment on the content',
    inputSchema: z.object({}),
    strict: true,
    execute: async () => {
      const commentSchema = z.object({
        blockId: z
          .string()
          .describe(
            'The id of the starting block. If the comment spans multiple blocks, use the id of the first block.'
          ),
        comment: z
          .string()
          .describe('A brief comment or explanation for this fragment.'),
        content: z
          .string()
          .describe(
            String.raw`The original document fragment to be commented on.It can be the entire block, a small part within a block, or span multiple blocks. If spanning multiple blocks, separate them with two \n\n.`
          ),
      });
 
      const { partialOutputStream } = streamText({
        model,
        output: Output.array({ element: commentSchema }),
        prompt: getCommentPrompt(editor, {
          messages: messagesRaw,
        }),
      });
 
      let lastLength = 0;
 
      for await (const partialArray of partialOutputStream) {
        for (let i = lastLength; i < partialArray.length; i++) {
          const comment = partialArray[i];
          const commentDataId = nanoid();
 
          writer.write({
            id: commentDataId,
            data: {
              comment,
              status: 'streaming',
            },
            type: 'data-comment',
          });
        }
 
        lastLength = partialArray.length;
      }
 
      writer.write({
        id: nanoid(),
        data: {
          comment: null,
          status: 'finished',
        },
        type: 'data-comment',
      });
    },
  });
 
const getTableTool = (
  editor: SlateEditor,
  {
    messagesRaw,
    model,
    writer,
  }: {
    messagesRaw: ChatMessage[];
    model: LanguageModel;
    writer: UIMessageStreamWriter<ChatMessage>;
  }
) =>
  tool({
    description: 'Edit table cells',
    inputSchema: z.object({}),
    strict: true,
    execute: async () => {
      const cellUpdateSchema = z.object({
        content: z
          .string()
          .describe(
            String.raw`The new content for the cell. Can contain multiple paragraphs separated by \n\n.`
          ),
        id: z.string().describe('The id of the table cell to update.'),
      });
 
      const { partialOutputStream } = streamText({
        model,
        output: Output.array({ element: cellUpdateSchema }),
        prompt: buildEditTableMultiCellPrompt(editor, messagesRaw),
      });
 
      let lastLength = 0;
 
      for await (const partialArray of partialOutputStream) {
        for (let i = lastLength; i < partialArray.length; i++) {
          const cellUpdate = partialArray[i];
 
          writer.write({
            id: nanoid(),
            data: {
              cellUpdate,
              status: 'streaming',
            },
            type: 'data-table',
          });
        }
 
        lastLength = partialArray.length;
      }
 
      writer.write({
        id: nanoid(),
        data: {
          cellUpdate: null,
          status: 'finished',
        },
        type: 'data-table',
      });
    },
  });
import type {
  ChatMessage,
  ToolName,
} from '@/components/editor/use-chat';
import type { NextRequest } from 'next/server';
 
import { createGateway } from '@ai-sdk/gateway';
import {
  type LanguageModel,
  type UIMessageStreamWriter,
  createUIMessageStream,
  createUIMessageStreamResponse,
  generateText,
  Output,
  streamText,
  tool,
} from 'ai';
import { NextResponse } from 'next/server';
import { type SlateEditor, createSlateEditor, nanoid } from 'platejs';
import { z } from 'zod';
 
import { BaseEditorKit } from '@/components/editor/editor-base-kit';
import { markdownJoinerTransform } from '@/lib/markdown-joiner-transform';
 
import {
  buildEditTableMultiCellPrompt,
  getChooseToolPrompt,
  getCommentPrompt,
  getEditPrompt,
  getGeneratePrompt,
} from './prompt';
 
export async function POST(req: NextRequest) {
  const { apiKey: key, ctx, messages: messagesRaw, model } = await req.json();
 
  const { children, selection, toolName: toolNameParam } = ctx;
 
  const editor = createSlateEditor({
    plugins: BaseEditorKit,
    selection,
    value: children,
  });
 
  const apiKey = key || process.env.AI_GATEWAY_API_KEY;
 
  if (!apiKey) {
    return NextResponse.json(
      { error: 'Missing AI Gateway API key.' },
      { status: 401 }
    );
  }
 
  const isSelecting = editor.api.isExpanded();
 
  const gatewayProvider = createGateway({
    apiKey,
  });
 
  try {
    const stream = createUIMessageStream<ChatMessage>({
      execute: async ({ writer }) => {
        let toolName = toolNameParam;
 
        if (!toolName) {
          const prompt = getChooseToolPrompt({
            isSelecting,
            messages: messagesRaw,
          });
 
          const enumOptions = isSelecting
            ? ['generate', 'edit', 'comment']
            : ['generate', 'comment'];
          const modelId = model || 'google/gemini-2.5-flash';
 
          const { output: AIToolName } = await generateText({
            model: gatewayProvider(modelId),
            output: Output.choice({ options: enumOptions }),
            prompt,
          });
 
          writer.write({
            data: AIToolName as ToolName,
            type: 'data-toolName',
          });
 
          toolName = AIToolName;
        }
 
        const stream = streamText({
          experimental_transform: markdownJoinerTransform(),
          model: gatewayProvider(model || 'openai/gpt-4o-mini'),
          // Not used
          prompt: '',
          tools: {
            comment: getCommentTool(editor, {
              messagesRaw,
              model: gatewayProvider(model || 'google/gemini-2.5-flash'),
              writer,
            }),
            table: getTableTool(editor, {
              messagesRaw,
              model: gatewayProvider(model || 'google/gemini-2.5-flash'),
              writer,
            }),
          },
          prepareStep: async (step) => {
            if (toolName === 'comment') {
              return {
                ...step,
                toolChoice: { toolName: 'comment', type: 'tool' },
              };
            }
 
            if (toolName === 'edit') {
              const [editPrompt, editType] = getEditPrompt(editor, {
                isSelecting,
                messages: messagesRaw,
              });
 
              // Table editing uses the table tool
              if (editType === 'table') {
                return {
                  ...step,
                  toolChoice: { toolName: 'table', type: 'tool' },
                };
              }
 
              return {
                ...step,
                activeTools: [],
                model:
                  editType === 'selection'
                    ? //The selection task is more challenging, so we chose to use Gemini 2.5 Flash.
                      gatewayProvider(model || 'google/gemini-2.5-flash')
                    : gatewayProvider(model || 'openai/gpt-4o-mini'),
                messages: [
                  {
                    content: editPrompt,
                    role: 'user',
                  },
                ],
              };
            }
 
            if (toolName === 'generate') {
              const generatePrompt = getGeneratePrompt(editor, {
                isSelecting,
                messages: messagesRaw,
              });
 
              return {
                ...step,
                activeTools: [],
                messages: [
                  {
                    content: generatePrompt,
                    role: 'user',
                  },
                ],
                model: gatewayProvider(model || 'openai/gpt-4o-mini'),
              };
            }
          },
        });
 
        writer.merge(stream.toUIMessageStream({ sendFinish: false }));
      },
    });
 
    return createUIMessageStreamResponse({ stream });
  } catch {
    return NextResponse.json(
      { error: 'Failed to process AI request' },
      { status: 500 }
    );
  }
}
 
const getCommentTool = (
  editor: SlateEditor,
  {
    messagesRaw,
    model,
    writer,
  }: {
    messagesRaw: ChatMessage[];
    model: LanguageModel;
    writer: UIMessageStreamWriter<ChatMessage>;
  }
) =>
  tool({
    description: 'Comment on the content',
    inputSchema: z.object({}),
    strict: true,
    execute: async () => {
      const commentSchema = z.object({
        blockId: z
          .string()
          .describe(
            'The id of the starting block. If the comment spans multiple blocks, use the id of the first block.'
          ),
        comment: z
          .string()
          .describe('A brief comment or explanation for this fragment.'),
        content: z
          .string()
          .describe(
            String.raw`The original document fragment to be commented on.It can be the entire block, a small part within a block, or span multiple blocks. If spanning multiple blocks, separate them with two \n\n.`
          ),
      });
 
      const { partialOutputStream } = streamText({
        model,
        output: Output.array({ element: commentSchema }),
        prompt: getCommentPrompt(editor, {
          messages: messagesRaw,
        }),
      });
 
      let lastLength = 0;
 
      for await (const partialArray of partialOutputStream) {
        for (let i = lastLength; i < partialArray.length; i++) {
          const comment = partialArray[i];
          const commentDataId = nanoid();
 
          writer.write({
            id: commentDataId,
            data: {
              comment,
              status: 'streaming',
            },
            type: 'data-comment',
          });
        }
 
        lastLength = partialArray.length;
      }
 
      writer.write({
        id: nanoid(),
        data: {
          comment: null,
          status: 'finished',
        },
        type: 'data-comment',
      });
    },
  });
 
const getTableTool = (
  editor: SlateEditor,
  {
    messagesRaw,
    model,
    writer,
  }: {
    messagesRaw: ChatMessage[];
    model: LanguageModel;
    writer: UIMessageStreamWriter<ChatMessage>;
  }
) =>
  tool({
    description: 'Edit table cells',
    inputSchema: z.object({}),
    strict: true,
    execute: async () => {
      const cellUpdateSchema = z.object({
        content: z
          .string()
          .describe(
            String.raw`The new content for the cell. Can contain multiple paragraphs separated by \n\n.`
          ),
        id: z.string().describe('The id of the table cell to update.'),
      });
 
      const { partialOutputStream } = streamText({
        model,
        output: Output.array({ element: cellUpdateSchema }),
        prompt: buildEditTableMultiCellPrompt(editor, messagesRaw),
      });
 
      let lastLength = 0;
 
      for await (const partialArray of partialOutputStream) {
        for (let i = lastLength; i < partialArray.length; i++) {
          const cellUpdate = partialArray[i];
 
          writer.write({
            id: nanoid(),
            data: {
              cellUpdate,
              status: 'streaming',
            },
            type: 'data-table',
          });
        }
 
        lastLength = partialArray.length;
      }
 
      writer.write({
        id: nanoid(),
        data: {
          cellUpdate: null,
          status: 'finished',
        },
        type: 'data-table',
      });
    },
  });

配置环境

在本地设置您的 AI 网关密钥（如果您不使用网关，请替换为您的提供商密钥）：

.env.local
AI_GATEWAY_API_KEY="your-api-key"
.env.local
AI_GATEWAY_API_KEY="your-api-key"

手动使用

安装

pnpm add @platejs/ai @platejs/markdown @platejs/selection @ai-sdk/react ai
pnpm add @platejs/ai @platejs/markdown @platejs/selection @ai-sdk/react ai

@platejs/suggestion 是可选的，但对于基于差异的编辑建议是必需的。

添加插件

import { createPlateEditor } from 'platejs/react';
import { AIChatPlugin, AIPlugin } from '@platejs/ai/react';
import { BlockSelectionPlugin } from '@platejs/selection/react';
import { MarkdownPlugin } from '@platejs/markdown';
 
export const editor = createPlateEditor({
  plugins: [
    BlockSelectionPlugin,
    MarkdownPlugin,
    AIPlugin,
    AIChatPlugin, // 在下一步中扩展
  ],
});
import { createPlateEditor } from 'platejs/react';
import { AIChatPlugin, AIPlugin } from '@platejs/ai/react';
import { BlockSelectionPlugin } from '@platejs/selection/react';
import { MarkdownPlugin } from '@platejs/markdown';
 
export const editor = createPlateEditor({
  plugins: [
    BlockSelectionPlugin,
    MarkdownPlugin,
    AIPlugin,
    AIChatPlugin, // 在下一步中扩展
  ],
});
  • BlockSelectionPlugin：启用 AIChatPlugin 依赖的多块选择，用于插入/替换变换。
  • MarkdownPlugin：提供流式工具使用的 Markdown 序列化。
  • AIPlugin：添加 AI 标记和用于撤销 AI 批处理的变换。
  • AIChatPlugin：提供 AI 组合框、API 助手和变换。

使用 AIPlugin.withComponent 配合您自己的元素（或 AILeaf）来高亮 AI 生成的文本。

配置 AIChatPlugin

扩展 AIChatPlugin 以连接流式传输和编辑。此示例镜像了 AIKit 的核心逻辑，同时保持 UI 无头。

import cloneDeep from 'lodash/cloneDeep';
import { BaseAIPlugin, withAIBatch } from '@platejs/ai';
import {
  AIChatPlugin,
  applyAISuggestions,
  getInsertPreviewStart,
  streamInsertChunk,
  useChatChunk,
} from '@platejs/ai/react';
import { ElementApi, getPluginType, KEYS, PathApi } from 'platejs';
import { usePluginOption } from 'platejs/react';
 
export const aiChatPlugin = AIChatPlugin.extend({
  options: {
    chatOptions: {
      api: '/api/ai/command',
      body: {
        model: 'openai/gpt-4o-mini',
      },
    },
    trigger: ' ',
    triggerPreviousCharPattern: /^\s?$/,
  },
  useHooks: ({ editor, getOption }) => {
    const mode = usePluginOption(AIChatPlugin, 'mode');
    const toolName = usePluginOption(AIChatPlugin, 'toolName');
 
    useChatChunk({
      onChunk: ({ chunk, isFirst, text }) => {
        if (isFirst && mode === 'insert') {
          const { startBlock, startInEmptyParagraph } =
            getInsertPreviewStart(editor);
 
          editor.getTransforms(BaseAIPlugin).ai.beginPreview({
            originalBlocks:
              startInEmptyParagraph &&
              startBlock &&
              ElementApi.isElement(startBlock)
                ? [cloneDeep(startBlock)]
                : [],
          });
 
          editor.setOption(AIChatPlugin, 'streaming', true);
 
          editor.tf.withoutSaving(() => {
            editor.tf.insertNodes(
              {
                children: [{ text: '' }],
                type: getPluginType(editor, KEYS.aiChat),
              },
              {
                at: PathApi.next(editor.selection!.focus.path.slice(0, 1)),
              }
            );
          });
        }
 
        if (mode === 'insert') {
          editor.tf.withoutSaving(() => {
            if (!getOption('streaming')) return;
 
            editor.tf.withScrolling(() => {
              streamInsertChunk(editor, chunk, {
                textProps: {
                  [getPluginType(editor, KEYS.ai)]: true,
                },
              });
            });
          });
        }
 
        if (toolName === 'edit' && mode === 'chat') {
          withAIBatch(
            editor,
            () => {
              applyAISuggestions(editor, text);
            },
            { split: isFirst }
          );
        }
      },
      onFinish: () => {
        editor.setOption(AIChatPlugin, 'streaming', false);
        editor.setOption(AIChatPlugin, '_blockChunks', '');
        editor.setOption(AIChatPlugin, '_blockPath', null);
        editor.setOption(AIChatPlugin, '_mdxName', null);
      },
    });
  },
});
import cloneDeep from 'lodash/cloneDeep';
import { BaseAIPlugin, withAIBatch } from '@platejs/ai';
import {
  AIChatPlugin,
  applyAISuggestions,
  getInsertPreviewStart,
  streamInsertChunk,
  useChatChunk,
} from '@platejs/ai/react';
import { ElementApi, getPluginType, KEYS, PathApi } from 'platejs';
import { usePluginOption } from 'platejs/react';
 
export const aiChatPlugin = AIChatPlugin.extend({
  options: {
    chatOptions: {
      api: '/api/ai/command',
      body: {
        model: 'openai/gpt-4o-mini',
      },
    },
    trigger: ' ',
    triggerPreviousCharPattern: /^\s?$/,
  },
  useHooks: ({ editor, getOption }) => {
    const mode = usePluginOption(AIChatPlugin, 'mode');
    const toolName = usePluginOption(AIChatPlugin, 'toolName');
 
    useChatChunk({
      onChunk: ({ chunk, isFirst, text }) => {
        if (isFirst && mode === 'insert') {
          const { startBlock, startInEmptyParagraph } =
            getInsertPreviewStart(editor);
 
          editor.getTransforms(BaseAIPlugin).ai.beginPreview({
            originalBlocks:
              startInEmptyParagraph &&
              startBlock &&
              ElementApi.isElement(startBlock)
                ? [cloneDeep(startBlock)]
                : [],
          });
 
          editor.setOption(AIChatPlugin, 'streaming', true);
 
          editor.tf.withoutSaving(() => {
            editor.tf.insertNodes(
              {
                children: [{ text: '' }],
                type: getPluginType(editor, KEYS.aiChat),
              },
              {
                at: PathApi.next(editor.selection!.focus.path.slice(0, 1)),
              }
            );
          });
        }
 
        if (mode === 'insert') {
          editor.tf.withoutSaving(() => {
            if (!getOption('streaming')) return;
 
            editor.tf.withScrolling(() => {
              streamInsertChunk(editor, chunk, {
                textProps: {
                  [getPluginType(editor, KEYS.ai)]: true,
                },
              });
            });
          });
        }
 
        if (toolName === 'edit' && mode === 'chat') {
          withAIBatch(
            editor,
            () => {
              applyAISuggestions(editor, text);
            },
            { split: isFirst }
          );
        }
      },
      onFinish: () => {
        editor.setOption(AIChatPlugin, 'streaming', false);
        editor.setOption(AIChatPlugin, '_blockChunks', '');
        editor.setOption(AIChatPlugin, '_blockPath', null);
        editor.setOption(AIChatPlugin, '_mdxName', null);
      },
    });
  },
});
  • useChatChunk：监视 UseChatHelpers 状态并生成增量块。
  • tf.ai.beginPreview：在写入首个未保存的预览块之前，捕获插入模式预览的回滚块切片和选区。
  • streamInsertChunk：将 Markdown/MDX 流式传输到文档中，尽可能重用现有块。
  • applyAISuggestions：当 toolName === 'edit' 时，将响应转换为临时建议节点。
  • withAIBatch：标记已保存的 AI 批次，使建议审查和已接受的 AI 更改保持可撤销。

扩展插件时，请提供您自己的 render 组件（工具栏按钮、浮动菜单等）。

构建 API 路由

在服务器上处理 api.aiChat.submit 请求。每个请求包含来自 @ai-sdk/react 的聊天 messages 和一个 ctx 有效载荷，其中包含编辑器 children、当前 selection 和最后的 toolName。 完整 API 示例

app/api/ai/command/route.ts
import { createGateway } from '@ai-sdk/gateway';
import { convertToCoreMessages, streamText } from 'ai';
import { createSlateEditor } from 'platejs';
 
import { BaseEditorKit } from '@/registry/components/editor/editor-base-kit';
import { markdownJoinerTransform } from '@/registry/lib/markdown-joiner-transform';
 
export async function POST(req: Request) {
  const { apiKey, ctx, messages, model } = await req.json();
 
  const editor = createSlateEditor({
    plugins: BaseEditorKit,
    selection: ctx.selection,
    value: ctx.children,
  });
 
  const gateway = createGateway({
    apiKey: apiKey ?? process.env.AI_GATEWAY_API_KEY!,
  });
 
  const result = streamText({
    experimental_transform: markdownJoinerTransform(),
    messages: convertToCoreMessages(messages),
    model: gateway(model ?? 'openai/gpt-4o-mini'),
    system: ctx.toolName === 'edit' ? 'You are an editor that rewrites user text.' : undefined,
  });
 
  return result.toDataStreamResponse();
}
app/api/ai/command/route.ts
import { createGateway } from '@ai-sdk/gateway';
import { convertToCoreMessages, streamText } from 'ai';
import { createSlateEditor } from 'platejs';
 
import { BaseEditorKit } from '@/registry/components/editor/editor-base-kit';
import { markdownJoinerTransform } from '@/registry/lib/markdown-joiner-transform';
 
export async function POST(req: Request) {
  const { apiKey, ctx, messages, model } = await req.json();
 
  const editor = createSlateEditor({
    plugins: BaseEditorKit,
    selection: ctx.selection,
    value: ctx.children,
  });
 
  const gateway = createGateway({
    apiKey: apiKey ?? process.env.AI_GATEWAY_API_KEY!,
  });
 
  const result = streamText({
    experimental_transform: markdownJoinerTransform(),
    messages: convertToCoreMessages(messages),
    model: gateway(model ?? 'openai/gpt-4o-mini'),
    system: ctx.toolName === 'edit' ? 'You are an editor that rewrites user text.' : undefined,
  });
 
  return result.toDataStreamResponse();
}
  • ctx.children 和 ctx.selection 被重新水合到 Slate 编辑器中，以便您可以构建丰富的提示（参见提示模板）。
  • 通过 chatOptions.body 转发提供商设置（model、apiKey、temperature、gateway 标志等）；您添加的所有内容都会在 JSON 有效载荷中原样传递，可以在调用 createGateway 之前读取。
  • 始终从服务器读取密钥。客户端应仅发送不透明标识符或短期令牌。
  • 返回流式响应，以便 useChat 和 useChatChunk 可以增量处理令牌。

连接 useChat

使用 @ai-sdk/react 桥接编辑器和您的模型端点。将助手存储在插件上，以便变换可以重新加载、停止或显示聊天状态。

import { useEffect } from 'react';
 
import { type UIMessage, DefaultChatTransport } from 'ai';
import { type UseChatHelpers, useChat } from '@ai-sdk/react';
import { AIChatPlugin } from '@platejs/ai/react';
import { useEditorPlugin } from 'platejs/react';
 
type ChatMessage = UIMessage<{}, { toolName: 'comment' | 'edit' | 'generate'; comment?: unknown }>;
 
export const useEditorAIChat = () => {
  const { editor, setOption } = useEditorPlugin(AIChatPlugin);
 
  const chat = useChat<ChatMessage>({
    id: 'editor',
    api: '/api/ai/command',
    transport: new DefaultChatTransport(),
    onData(data) {
      if (data.type === 'data-toolName') {
        editor.setOption(AIChatPlugin, 'toolName', data.data);
      }
    },
  });
 
  useEffect(() => {
    setOption('chat', chat as UseChatHelpers<ChatMessage>);
  }, [chat, setOption]);
 
  return chat;
};
import { useEffect } from 'react';
 
import { type UIMessage, DefaultChatTransport } from 'ai';
import { type UseChatHelpers, useChat } from '@ai-sdk/react';
import { AIChatPlugin } from '@platejs/ai/react';
import { useEditorPlugin } from 'platejs/react';
 
type ChatMessage = UIMessage<{}, { toolName: 'comment' | 'edit' | 'generate'; comment?: unknown }>;
 
export const useEditorAIChat = () => {
  const { editor, setOption } = useEditorPlugin(AIChatPlugin);
 
  const chat = useChat<ChatMessage>({
    id: 'editor',
    api: '/api/ai/command',
    transport: new DefaultChatTransport(),
    onData(data) {
      if (data.type === 'data-toolName') {
        editor.setOption(AIChatPlugin, 'toolName', data.data);
      }
    },
  });
 
  useEffect(() => {
    setOption('chat', chat as UseChatHelpers<ChatMessage>);
  }, [chat, setOption]);
 
  return chat;
};

将助手与 useEditorChat 结合使用，以保持浮动菜单正确锚定：

import { useEditorChat } from '@platejs/ai/react';
 
useEditorChat({
  onOpenChange: (open) => {
    if (!open) chat.stop?.();
  },
});
import { useEditorChat } from '@platejs/ai/react';
 
useEditorChat({
  onOpenChange: (open) => {
    if (!open) chat.stop?.();
  },
});

现在您可以以编程方式提交提示：

import { AIChatPlugin } from '@platejs/ai/react';
 
editor.getApi(AIChatPlugin).aiChat.submit('', {
  prompt: {
    default: 'Continue the document after {block}',
    selecting: 'Rewrite {selection} with a clearer tone',
  },
  toolName: 'generate',
});
import { AIChatPlugin } from '@platejs/ai/react';
 
editor.getApi(AIChatPlugin).aiChat.submit('', {
  prompt: {
    default: 'Continue the document after {block}',
    selecting: 'Rewrite {selection} with a clearer tone',
  },
  toolName: 'generate',
});

提示模板

客户端提示

  • api.aiChat.submit 接受一个 EditorPrompt。提供一个字符串、带有 default/selecting/blockSelecting 的对象，或一个接收 { editor, isSelecting, isBlockSelecting } 的函数。客户端中的助手 getEditorPrompt 将该值转换为最终字符串。
  • 将其与 replacePlaceholders(editor, template, { prompt }) 结合使用，以使用 @platejs/ai 生成的 Markdown 扩展 {editor}、{block}、{blockSelection} 和 {prompt}。
import { replacePlaceholders } from '@platejs/ai';
 
editor.getApi(AIChatPlugin).aiChat.submit('Improve tone', {
  prompt: ({ isSelecting }) =>
    isSelecting
      ? replacePlaceholders(editor, 'Rewrite {blockSelection} using a friendly tone.')
      : replacePlaceholders(editor, 'Continue {block} with two more sentences.'),
  toolName: 'generate',
});
import { replacePlaceholders } from '@platejs/ai';
 
editor.getApi(AIChatPlugin).aiChat.submit('Improve tone', {
  prompt: ({ isSelecting }) =>
    isSelecting
      ? replacePlaceholders(editor, 'Rewrite {blockSelection} using a friendly tone.')
      : replacePlaceholders(editor, 'Continue {block} with two more sentences.'),
  toolName: 'generate',
});

服务器端提示

apps/www/src/app/api/ai/command 中的演示后端从 ctx 重建编辑器并构建结构化提示：

  • getChooseToolPrompt 决定请求是 generate、edit 还是 comment。
  • getGeneratePrompt、getEditPrompt 和 getCommentPrompt 将当前编辑器状态转换为针对每种模式量身定制的指令。
  • 工具助手如 getMarkdown、getMarkdownWithSelection 和 buildStructuredPrompt（参见 apps/www/src/app/api/ai/command/prompts.ts）使将块 ID、选择和 MDX 标签嵌入 LLM 请求变得容易。

增强从客户端发送的有效载荷以微调服务器提示：

editor.setOption(aiChatPlugin, 'chatOptions', {
  api: '/api/ai/command',
  body: {
    model: 'openai/gpt-4o-mini',
    tone: 'playful',
    temperature: 0.4,
  },
});
editor.setOption(aiChatPlugin, 'chatOptions', {
  api: '/api/ai/command',
  body: {
    model: 'openai/gpt-4o-mini',
    tone: 'playful',
    temperature: 0.4,
  },
});

chatOptions.body 下的所有内容都会到达路由处理程序，让您可以交换提供商、传递用户特定的元数据或分支到不同的提示模板。

键盘快捷键

KeyDescription
Space在空块中打开 AI 菜单（光标模式）
Cmd + J显示 AI 菜单（通过 shortcuts.show 设置）
Escape隐藏 AI 菜单并停止流式传输

流式传输

流式工具在响应到达时保持复杂布局完整：

  • streamInsertChunk(editor, chunk, options) 反序列化 Markdown 块，就地更新当前块，并根据需要追加新块。使用 textProps/elementProps 标记流式节点（例如，标记 AI 文本）。
  • streamDeserializeMd 和 streamDeserializeInlineMd 提供更低级别的访问，如果您需要控制自定义节点类型的流式传输。
  • streamSerializeMd 镜像编辑器状态，以便您可以检测流式内容与响应缓冲区之间的漂移。

流式传输完成时重置内部 _blockChunks、_blockPath 和 _mdxName 选项，以从干净的状态开始下一个响应。

流式传输示例

Loading…

Plate Plus

Combobox menu with free-form prompt input

  • Additional trigger methods:
    • Block menu button
    • Slash command menu
  • Beautifully crafted UI

Hooks

useAIChatEditor

为聊天预览注册一个辅助编辑器，并使用块级记忆化反序列化 Markdown。

Parameters

    专用于聊天预览的编辑器实例。

    模型返回的 Markdown 内容。

    传递 parser 以在反序列化之前过滤令牌。
import { usePlateEditor } from 'platejs/react';
import { MarkdownPlugin } from '@platejs/markdown';
import { AIChatPlugin, useAIChatEditor } from '@platejs/ai/react';
 
const aiPreviewEditor = usePlateEditor({
  plugins: [MarkdownPlugin, AIChatPlugin],
});
 
useAIChatEditor(aiPreviewEditor, responseMarkdown, {
  parser: { exclude: ['space'] },
});
import { usePlateEditor } from 'platejs/react';
import { MarkdownPlugin } from '@platejs/markdown';
import { AIChatPlugin, useAIChatEditor } from '@platejs/ai/react';
 
const aiPreviewEditor = usePlateEditor({
  plugins: [MarkdownPlugin, AIChatPlugin],
});
 
useAIChatEditor(aiPreviewEditor, responseMarkdown, {
  parser: { exclude: ['space'] },
});

useEditorChat

将 UseChatHelpers 连接到编辑器状态，以便 AI 菜单知道是锚定到光标、选择还是块选择。

Parameters

    在块选择上打开菜单时调用。

    菜单打开或关闭时调用。

    在光标处打开菜单时调用。

    在文本选择上打开菜单时调用。

useChatChunk

逐块流式传输聊天响应，让您完全控制插入。

Parameters

    处理每个流式块。

    流式传输完成时调用。

工具函数

withAIBatch

将编辑器操作分组到单个历史批次中，并将其标记为 AI 生成，以便 tf.ai.undo() 安全地移除它。

Parameters

    目标编辑器。

    要运行的操作。

    设置 split: true 以开始新的历史批次。

applyAISuggestions

将 AI 输出与存储的 chatNodes 进行差异比较，并写入临时建议节点。需要 @platejs/suggestion。

Parameters

    要应用建议的编辑器。

    来自模型的 Markdown 响应。

补充助手允许您完成或放弃差异：

  • acceptAISuggestions(editor)：将临时建议节点转换为永久建议。
  • rejectAISuggestions(editor)：移除临时建议节点并清除建议标记。

aiCommentToRange

将流式评论元数据映射回文档范围，以便可以自动插入评论。

Parameters

    编辑器实例。

    用于定位范围的块 ID 和文本。

Returns{ start: BasePoint; end: BasePoint } | null

    匹配评论的范围，如果找不到则为 null。

findTextRangeInBlock

使用 LCS 在块内查找最接近匹配的模糊搜索助手。

Parameters

    要搜索的块节点。

    要定位的文本片段。

Returns{ start: { path: Path; offset: number }; end: { path: Path; offset: number } } | null

    匹配的范围或 null。

getEditorPrompt

生成尊重光标、选择或块选择状态的提示。

Parameters

    提供上下文的编辑器。

    描述提示的字符串、配置或函数。

Returnsstring

    上下文化的提示字符串。

replacePlaceholders

用序列化的 Markdown 替换 {editor}、{blockSelection} 和 {prompt} 等占位符。

Parameters

    提供内容的编辑器。

    模板文本。

    注入到 {prompt} 中的提示值。

Returnsstring

    占位符被 Markdown 替换后的模板。

插件

AIPlugin

向流式文本添加 ai 标记，并暴露变换以移除 AI 节点或撤销最后一个 AI 批次。使用 .withComponent 以自定义组件渲染 AI 标记的文本。

Options

    AI 内容存储在文本节点上。

    AI 标记是常规文本属性，而不是装饰。

AIChatPlugin

驱动 AI 菜单、聊天状态和变换的主插件。

Options

    打开命令菜单的字符。默认为 ' '。

    触发器前字符必须匹配的模式。默认为 /^\s?$/。

    返回 false 以在特定上下文中取消打开。

    存储来自 useChat 的助手，以便 API 调用可以访问它们。

    用于差异编辑建议的节点快照（内部管理）。

    提交提示前捕获的选择（内部管理）。

    控制响应是直接流式传输到文档还是打开审查面板。默认为 'insert'。

    AI 菜单是否可见。默认为 false。

    响应流式传输时为 true。默认为 false。

    用于解释响应的活动工具。

API

api.aiChat.submit(input, options?)

向您的模型提供商提交提示。当省略 mode 时，折叠光标默认为 'insert'，否则为 'chat'。

Parameters

    来自用户的原始输入。

    微调提交行为。

Optionsobject

    覆盖响应模式。

    转发到 chat.sendMessage（model、headers 等）。

    由 getEditorPrompt 处理的字符串、配置或函数。

    标记提交，以便 hooks 可以做出不同的响应。

api.aiChat.reset(options?)

清除聊天状态，移除 AI 节点，并可选择撤销最后一个 AI 批次。

Parameters

    传递 undo: false 以保留流式内容。

api.aiChat.node(options?)

检索与指定条件匹配的第一个 AI 节点。

Parameters

    设置 anchor: true 以获取锚点节点，或 streaming: true 以检索当前正在流式传输的节点。

ReturnsNodeEntry | undefined

    匹配的节点条目（如果找到）。

api.aiChat.reload()

使用存储的 UseChatHelpers 重放最后一个提示，在重新提交之前恢复原始选择或块选择。

api.aiChat.stop()

停止流式传输并调用 chat.stop。

api.aiChat.show()

打开 AI 菜单，清除之前的聊天消息，并重置工具状态。

api.aiChat.hide(options?)

关闭 AI 菜单，可选择撤销最后一个 AI 批次并重新聚焦编辑器。

Parameters

    设置 focus: false 以保持焦点在编辑器外，或 undo: false 以保留插入的内容。

变换

tf.aiChat.accept()

接受最新的响应。在插入模式下，它移除 AI 标记并将光标放置在流式内容的末尾。在聊天模式下，它应用待处理的建议。

tf.aiChat.insertBelow(sourceEditor, options?)

在当前选择或块选择下方插入聊天预览（sourceEditor）。

Parameters

    包含生成内容的编辑器。

    从源选择复制格式。默认为 'single'。

tf.aiChat.replaceSelection(sourceEditor, options?)

用聊天预览替换当前选择或块选择。

Parameters

    包含生成内容的编辑器。

    控制应应用原始选择的多少格式。

tf.aiChat.removeAnchor(options?)

移除用于定位 AI 菜单的临时锚点节点。

Parameters

    过滤要移除的节点。

tf.ai.insertNodes(nodes, options?)

在当前选择（或 options.target）处插入带有 AI 标记的节点。

tf.ai.removeMarks(options?)

从匹配的节点中清除 AI 标记。

tf.ai.removeNodes(options?)

移除标记为 AI 生成的文本节点。

tf.ai.beginPreview(options?)

捕获插入模式 AI 预览的回滚块切片和选区。在写入首个未保存的预览块之前调用一次。

Parameters

    预览将覆盖的顶层块。若预览是在现有内容之后插入，则传 []。

Returnsboolean

    当存储了新的预览回滚点时返回 true；若预览状态已存在则返回 false。

tf.ai.acceptPreview()

将当前预览作为一次新的可撤销批次提交，移除仅预览使用的标记，并清除预览状态。

Returnsboolean

    当活动预览被提交时返回 true。

tf.ai.cancelPreview()

恢复当前预览的回滚点，并清除预览状态。

Returnsboolean

    当活动预览被恢复时返回 true。

tf.ai.discardPreview()

清除预览状态而不恢复内容。当预览内容应该保留在文档中时使用它。

Returnsboolean

    当活动预览状态被清除时返回 true。

tf.ai.hasPreview()

报告当前是否存在插入模式预览的回滚点。

Returnsboolean

    当存在预览回滚状态时返回 true。

tf.ai.undo()

如果最新的 AI 历史条目由 withAIBatch 创建，则撤销它。若当前存在插入模式预览，则优先取消该预览，而不是回放每个流式块。两种情况下都会避免通过重做重新应用 AI 输出。

自定义

添加自定义 AI 命令

'use client';
 
import * as React from 'react';
 
import {
  AIChatPlugin,
  AIPlugin,
  useEditorChat,
  useLastAssistantMessage,
} from '@platejs/ai/react';
import { getTransientCommentKey } from '@platejs/comment';
import { BlockSelectionPlugin, useIsSelecting } from '@platejs/selection/react';
import { getTransientSuggestionKey } from '@platejs/suggestion';
import { Command as CommandPrimitive } from 'cmdk';
import {
  Album,
  BadgeHelp,
  BookOpenCheck,
  Check,
  CornerUpLeft,
  FeatherIcon,
  ListEnd,
  ListMinus,
  ListPlus,
  Loader2Icon,
  PauseIcon,
  PenLine,
  SmileIcon,
  Wand,
  X,
} from 'lucide-react';
import {
  type NodeEntry,
  type SlateEditor,
  isHotkey,
  KEYS,
  NodeApi,
  TextApi,
} from 'platejs';
import {
  useEditorPlugin,
  useFocusedLast,
  useHotkeys,
  usePluginOption,
} from 'platejs/react';
import { type PlateEditor, useEditorRef } from 'platejs/react';
 
import { Button } from '@/components/ui/button';
import {
  Command,
  CommandGroup,
  CommandItem,
  CommandList,
} from '@/components/ui/command';
import {
  Popover,
  PopoverAnchor,
  PopoverContent,
} from '@/components/ui/popover';
import { cn } from '@/lib/utils';
import { commentPlugin } from '@/components/editor/plugins/comment-kit';
 
import { AIChatEditor } from './ai-chat-editor';
 
export function AIMenu() {
  const { api, editor } = useEditorPlugin(AIChatPlugin);
  const mode = usePluginOption(AIChatPlugin, 'mode');
  const toolName = usePluginOption(AIChatPlugin, 'toolName');
 
  const streaming = usePluginOption(AIChatPlugin, 'streaming');
  const isSelecting = useIsSelecting();
  const isFocusedLast = useFocusedLast();
  const open = usePluginOption(AIChatPlugin, 'open') && isFocusedLast;
  const [value, setValue] = React.useState('');
 
  const [input, setInput] = React.useState('');
 
  const chat = usePluginOption(AIChatPlugin, 'chat');
 
  const { messages, status } = chat;
  const [anchorElement, setAnchorElement] = React.useState<HTMLElement | null>(
    null
  );
 
  const content = useLastAssistantMessage()?.parts.find(
    (part) => part.type === 'text'
  )?.text;
 
  React.useEffect(() => {
    if (!streaming) return;
 
    const anchorEntry = api.aiChat.node({ anchor: true });
    if (!anchorEntry) return;
 
    const anchorDom = editor.api.toDOMNode(anchorEntry[0])!;
    // eslint-disable-next-line react-hooks/set-state-in-effect -- Position the popover from editor DOM while the edit stream is active.
    setAnchorElement(anchorDom);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [streaming]);
 
  const setOpen = (open: boolean) => {
    if (open) {
      api.aiChat.show();
    } else {
      api.aiChat.hide();
    }
  };
 
  const show = (anchorElement: HTMLElement) => {
    setAnchorElement(anchorElement);
    setOpen(true);
  };
 
  useEditorChat({
    onOpenBlockSelection: (blocks: NodeEntry[]) => {
      show(editor.api.toDOMNode(blocks.at(-1)![0])!);
    },
    onOpenChange: (open) => {
      if (!open) {
        setAnchorElement(null);
        setInput('');
      }
    },
    onOpenCursor: () => {
      const [ancestor] = editor.api.block({ highest: true })!;
 
      if (!editor.api.isAt({ end: true }) && !editor.api.isEmpty(ancestor)) {
        editor
          .getApi(BlockSelectionPlugin)
          .blockSelection.set(ancestor.id as string);
      }
 
      show(editor.api.toDOMNode(ancestor)!);
    },
    onOpenSelection: () => {
      show(editor.api.toDOMNode(editor.api.blocks().at(-1)![0])!);
    },
  });
 
  useHotkeys('esc', () => {
    api.aiChat.stop();
 
    // remove when you implement the route /api/ai/command
    (chat as any)._abortFakeStream();
  });
 
  const isLoading = status === 'streaming' || status === 'submitted';
 
  React.useEffect(() => {
    if (toolName !== 'edit' || mode !== 'chat' || isLoading) return;
 
    let anchorNode = editor.api.node({
      at: [],
      reverse: true,
      match: (n) => !!n[KEYS.suggestion] && !!n[getTransientSuggestionKey()],
    });
 
    if (!anchorNode) {
      anchorNode = editor
        .getApi(BlockSelectionPlugin)
        .blockSelection.getNodes({ selectionFallback: true, sort: true })
        .at(-1);
    }
 
    if (!anchorNode) return;
 
    const block = editor.api.block({ at: anchorNode[1] });
    // eslint-disable-next-line react-hooks/set-state-in-effect -- Position the popover from editor DOM after the edit stream completes.
    setAnchorElement(editor.api.toDOMNode(block![0]!)!);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [isLoading]);
 
  if (isLoading && mode === 'insert') return null;
 
  if (toolName === 'comment') return null;
 
  if (toolName === 'edit' && mode === 'chat' && isLoading) return null;
 
  return (
    <Popover open={open} onOpenChange={setOpen} modal={false}>
      <PopoverAnchor virtualRef={{ current: anchorElement! }} />
 
      <PopoverContent
        className="border-none bg-transparent p-0 shadow-none"
        style={{
          width: anchorElement?.offsetWidth,
        }}
        onEscapeKeyDown={(e) => {
          e.preventDefault();
 
          api.aiChat.hide();
        }}
        align="center"
        side="bottom"
      >
        <Command
          className="w-full rounded-lg border shadow-md"
          value={value}
          onValueChange={setValue}
        >
          {mode === 'chat' &&
            isSelecting &&
            content &&
            toolName === 'generate' && <AIChatEditor content={content} />}
 
          {isLoading ? (
            <div className="flex grow select-none items-center gap-2 p-2 text-muted-foreground text-sm">
              <Loader2Icon className="size-4 animate-spin" />
              {messages.length > 1 ? 'Editing...' : 'Thinking...'}
            </div>
          ) : (
            <CommandPrimitive.Input
              className={cn(
                'flex h-9 w-full min-w-0 border-input bg-transparent px-3 py-1 text-base outline-none transition-[color,box-shadow] placeholder:text-muted-foreground md:text-sm dark:bg-input/30',
                'aria-invalid:border-destructive aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40',
                'border-b focus-visible:ring-transparent'
              )}
              value={input}
              onKeyDown={(e) => {
                if (isHotkey('backspace')(e) && input.length === 0) {
                  e.preventDefault();
                  api.aiChat.hide();
                }
                if (isHotkey('enter')(e) && !e.shiftKey && !value) {
                  e.preventDefault();
                  void api.aiChat.submit(input);
                  setInput('');
                }
              }}
              onValueChange={setInput}
              placeholder="Ask AI anything..."
              data-plate-focus
              autoFocus
            />
          )}
 
          {!isLoading && (
            <CommandList>
              <AIMenuItems
                input={input}
                setInput={setInput}
                setValue={setValue}
              />
            </CommandList>
          )}
        </Command>
      </PopoverContent>
    </Popover>
  );
}
 
type EditorChatState =
  | 'cursorCommand'
  | 'cursorSuggestion'
  | 'selectionCommand'
  | 'selectionSuggestion';
 
const AICommentIcon = () => (
  <svg
    fill="none"
    height="24"
    stroke="currentColor"
    strokeLinecap="round"
    strokeLinejoin="round"
    strokeWidth="2"
    viewBox="0 0 24 24"
    width="24"
    xmlns="http://www.w3.org/2000/svg"
  >
    <path d="M0 0h24v24H0z" fill="none" stroke="none" />
    <path d="M8 9h8" />
    <path d="M8 13h4.5" />
    <path d="M10 19l-1 -1h-3a3 3 0 0 1 -3 -3v-8a3 3 0 0 1 3 -3h12a3 3 0 0 1 3 3v4.5" />
    <path d="M17.8 20.817l-2.172 1.138a.392 .392 0 0 1 -.568 -.41l.415 -2.411l-1.757 -1.707a.389 .389 0 0 1 .217 -.665l2.428 -.352l1.086 -2.193a.392 .392 0 0 1 .702 0l1.086 2.193l2.428 .352a.39 .39 0 0 1 .217 .665l-1.757 1.707l.414 2.41a.39 .39 0 0 1 -.567 .411l-2.172 -1.138z" />
  </svg>
);
 
const aiChatItems = {
  accept: {
    icon: <Check />,
    label: 'Accept',
    value: 'accept',
    onSelect: ({ aiEditor, editor }) => {
      const { mode, toolName } = editor.getOptions(AIChatPlugin);
 
      if (mode === 'chat' && toolName === 'generate') {
        return editor
          .getTransforms(AIChatPlugin)
          .aiChat.replaceSelection(aiEditor);
      }
 
      editor.getTransforms(AIChatPlugin).aiChat.accept();
      editor.tf.focus({ edge: 'end' });
    },
  },
  comment: {
    icon: <AICommentIcon />,
    label: 'Comment',
    value: 'comment',
    onSelect: ({ editor, input }) => {
      editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt:
          'Please comment on the following content and provide reasonable and meaningful feedback.',
        toolName: 'comment',
      });
    },
  },
  continueWrite: {
    icon: <PenLine />,
    label: 'Continue writing',
    value: 'continueWrite',
    onSelect: ({ editor, input }) => {
      const ancestorNode = editor.api.block({ highest: true });
 
      if (!ancestorNode) return;
 
      const isEmpty = NodeApi.string(ancestorNode[0]).trim().length === 0;
 
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt: isEmpty
          ? `<Document>
{editor}
</Document>
Start writing a new paragraph AFTER <Document> ONLY ONE SENTENCE`
          : 'Continue writing AFTER <Block> ONLY ONE SENTENCE. DONT REPEAT THE TEXT.',
        toolName: 'generate',
      });
    },
  },
  discard: {
    icon: <X />,
    label: 'Discard',
    shortcut: 'Escape',
    value: 'discard',
    onSelect: ({ editor }) => {
      editor.getTransforms(AIPlugin).ai.undo();
      editor.getApi(AIChatPlugin).aiChat.hide();
    },
  },
  emojify: {
    icon: <SmileIcon />,
    label: 'Emojify',
    value: 'emojify',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Add a small number of contextually relevant emojis within each block only. You may insert emojis, but do not remove, replace, or rewrite existing text, and do not modify Markdown syntax, links, or line breaks.',
        toolName: 'edit',
      });
    },
  },
  explain: {
    icon: <BadgeHelp />,
    label: 'Explain',
    value: 'explain',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: {
          default: 'Explain {editor}',
          selecting: 'Explain',
        },
        toolName: 'generate',
      });
    },
  },
  fixSpelling: {
    icon: <Check />,
    label: 'Fix spelling & grammar',
    value: 'fixSpelling',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Fix spelling, grammar, and punctuation errors within each block only, without changing meaning, tone, or adding new information.',
        toolName: 'edit',
      });
    },
  },
  generateMarkdownSample: {
    icon: <BookOpenCheck />,
    label: 'Generate Markdown sample',
    value: 'generateMarkdownSample',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: 'Generate a markdown sample',
        toolName: 'generate',
      });
    },
  },
  generateMdxSample: {
    icon: <BookOpenCheck />,
    label: 'Generate MDX sample',
    value: 'generateMdxSample',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: 'Generate a mdx sample',
        toolName: 'generate',
      });
    },
  },
  improveWriting: {
    icon: <Wand />,
    label: 'Improve writing',
    value: 'improveWriting',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Improve the writing for clarity and flow, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  insertBelow: {
    icon: <ListEnd />,
    label: 'Insert below',
    value: 'insertBelow',
    onSelect: ({ aiEditor, editor }) => {
      /** Format: 'none' Fix insert table */
      void editor
        .getTransforms(AIChatPlugin)
        .aiChat.insertBelow(aiEditor, { format: 'none' });
    },
  },
  makeLonger: {
    icon: <ListPlus />,
    label: 'Make longer',
    value: 'makeLonger',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Make the content longer by elaborating on existing ideas within each block only, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  makeShorter: {
    icon: <ListMinus />,
    label: 'Make shorter',
    value: 'makeShorter',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Make the content shorter by reducing verbosity within each block only, without changing meaning or removing essential information.',
        toolName: 'edit',
      });
    },
  },
  replace: {
    icon: <Check />,
    label: 'Replace selection',
    value: 'replace',
    onSelect: ({ aiEditor, editor }) => {
      void editor.getTransforms(AIChatPlugin).aiChat.replaceSelection(aiEditor);
    },
  },
  simplifyLanguage: {
    icon: <FeatherIcon />,
    label: 'Simplify language',
    value: 'simplifyLanguage',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Simplify the language by using clearer and more straightforward wording within each block only, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  summarize: {
    icon: <Album />,
    label: 'Add a summary',
    value: 'summarize',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt: {
          default: 'Summarize {editor}',
          selecting: 'Summarize',
        },
        toolName: 'generate',
      });
    },
  },
  tryAgain: {
    icon: <CornerUpLeft />,
    label: 'Try again',
    value: 'tryAgain',
    onSelect: ({ editor }) => {
      void editor.getApi(AIChatPlugin).aiChat.reload();
    },
  },
} satisfies Record<
  string,
  {
    icon: React.ReactNode;
    label: string;
    value: string;
    component?: React.ComponentType<{ menuState: EditorChatState }>;
    filterItems?: boolean;
    items?: { label: string; value: string }[];
    shortcut?: string;
    onSelect?: ({
      aiEditor,
      editor,
      input,
    }: {
      aiEditor: SlateEditor;
      editor: PlateEditor;
      input: string;
    }) => void;
  }
>;
 
const menuStateItems: Record<
  EditorChatState,
  {
    items: (typeof aiChatItems)[keyof typeof aiChatItems][];
    heading?: string;
  }[]
> = {
  cursorCommand: [
    {
      items: [
        aiChatItems.comment,
        aiChatItems.generateMdxSample,
        aiChatItems.generateMarkdownSample,
        aiChatItems.continueWrite,
        aiChatItems.summarize,
        aiChatItems.explain,
      ],
    },
  ],
  cursorSuggestion: [
    {
      items: [aiChatItems.accept, aiChatItems.discard, aiChatItems.tryAgain],
    },
  ],
  selectionCommand: [
    {
      items: [
        aiChatItems.improveWriting,
        aiChatItems.comment,
        aiChatItems.emojify,
        aiChatItems.makeLonger,
        aiChatItems.makeShorter,
        aiChatItems.fixSpelling,
        aiChatItems.simplifyLanguage,
      ],
    },
  ],
  selectionSuggestion: [
    {
      items: [
        aiChatItems.accept,
        aiChatItems.discard,
        aiChatItems.insertBelow,
        aiChatItems.tryAgain,
      ],
    },
  ],
};
 
export const AIMenuItems = ({
  input,
  setInput,
  setValue,
}: {
  input: string;
  setInput: (value: string) => void;
  setValue: (value: string) => void;
}) => {
  const editor = useEditorRef();
  const { messages } = usePluginOption(AIChatPlugin, 'chat');
  const aiEditor = usePluginOption(AIChatPlugin, 'aiEditor')!;
  const isSelecting = useIsSelecting();
 
  const menuState = React.useMemo(() => {
    if (messages && messages.length > 0) {
      return isSelecting ? 'selectionSuggestion' : 'cursorSuggestion';
    }
 
    return isSelecting ? 'selectionCommand' : 'cursorCommand';
  }, [isSelecting, messages]);
 
  const menuGroups = React.useMemo(() => {
    const items = menuStateItems[menuState];
 
    return items;
  }, [menuState]);
 
  React.useEffect(() => {
    if (menuGroups.length > 0 && menuGroups[0].items.length > 0) {
      setValue(menuGroups[0].items[0].value);
    }
  }, [menuGroups, setValue]);
 
  return (
    <>
      {menuGroups.map((group, index) => (
        <CommandGroup key={index} heading={group.heading}>
          {group.items.map((menuItem) => (
            <CommandItem
              key={menuItem.value}
              className="[&_svg]:text-muted-foreground"
              value={menuItem.value}
              onSelect={() => {
                menuItem.onSelect?.({
                  aiEditor,
                  editor,
                  input,
                });
                setInput('');
              }}
            >
              {menuItem.icon}
              <span>{menuItem.label}</span>
            </CommandItem>
          ))}
        </CommandGroup>
      ))}
    </>
  );
};
 
export function AILoadingBar() {
  const editor = useEditorRef();
 
  const toolName = usePluginOption(AIChatPlugin, 'toolName');
  const chat = usePluginOption(AIChatPlugin, 'chat');
  const mode = usePluginOption(AIChatPlugin, 'mode');
 
  const { status } = chat;
 
  const { api } = useEditorPlugin(AIChatPlugin);
 
  const isLoading = status === 'streaming' || status === 'submitted';
 
  const handleComments = (type: 'accept' | 'reject') => {
    if (type === 'accept') {
      editor.tf.unsetNodes([getTransientCommentKey()], {
        at: [],
        match: (n) => TextApi.isText(n) && !!n[KEYS.comment],
      });
    }
 
    if (type === 'reject') {
      editor
        .getTransforms(commentPlugin)
        .comment.unsetMark({ transient: true });
    }
 
    api.aiChat.hide();
  };
 
  useHotkeys('esc', () => {
    api.aiChat.stop();
 
    // remove when you implement the route /api/ai/command
    (chat as any)._abortFakeStream();
  });
 
  if (
    isLoading &&
    (mode === 'insert' ||
      toolName === 'comment' ||
      (toolName === 'edit' && mode === 'chat'))
  ) {
    return (
      <div
        className={cn(
          '-translate-x-1/2 absolute bottom-4 left-1/2 z-20 flex items-center gap-3 rounded-md border border-border bg-muted px-3 py-1.5 text-muted-foreground text-sm shadow-md transition-all duration-300'
        )}
      >
        <span className="h-4 w-4 animate-spin rounded-full border-2 border-muted-foreground border-t-transparent" />
        <span>{status === 'submitted' ? 'Thinking...' : 'Writing...'}</span>
        <Button
          size="sm"
          variant="ghost"
          className="flex items-center gap-1 text-xs"
          onClick={() => api.aiChat.stop()}
        >
          <PauseIcon className="h-4 w-4" />
          Stop
          <kbd className="ml-1 rounded bg-border px-1 font-mono text-[10px] text-muted-foreground shadow-sm">
            Esc
          </kbd>
        </Button>
      </div>
    );
  }
 
  if (toolName === 'comment' && status === 'ready') {
    return (
      <div
        className={cn(
          '-translate-x-1/2 absolute bottom-4 left-1/2 z-50 flex flex-col items-center gap-0 rounded-xl border border-border/50 bg-popover p-1 text-muted-foreground text-sm shadow-xl backdrop-blur-sm',
          'p-3'
        )}
      >
        {/* Header with controls */}
        <div className="flex w-full items-center justify-between gap-3">
          <div className="flex items-center gap-5">
            <Button
              size="sm"
              disabled={isLoading}
              onClick={() => handleComments('accept')}
            >
              Accept
            </Button>
 
            <Button
              size="sm"
              disabled={isLoading}
              onClick={() => handleComments('reject')}
            >
              Reject
            </Button>
          </div>
        </div>
      </div>
    );
  }
 
  return null;
}
'use client';
 
import * as React from 'react';
 
import {
  AIChatPlugin,
  AIPlugin,
  useEditorChat,
  useLastAssistantMessage,
} from '@platejs/ai/react';
import { getTransientCommentKey } from '@platejs/comment';
import { BlockSelectionPlugin, useIsSelecting } from '@platejs/selection/react';
import { getTransientSuggestionKey } from '@platejs/suggestion';
import { Command as CommandPrimitive } from 'cmdk';
import {
  Album,
  BadgeHelp,
  BookOpenCheck,
  Check,
  CornerUpLeft,
  FeatherIcon,
  ListEnd,
  ListMinus,
  ListPlus,
  Loader2Icon,
  PauseIcon,
  PenLine,
  SmileIcon,
  Wand,
  X,
} from 'lucide-react';
import {
  type NodeEntry,
  type SlateEditor,
  isHotkey,
  KEYS,
  NodeApi,
  TextApi,
} from 'platejs';
import {
  useEditorPlugin,
  useFocusedLast,
  useHotkeys,
  usePluginOption,
} from 'platejs/react';
import { type PlateEditor, useEditorRef } from 'platejs/react';
 
import { Button } from '@/components/ui/button';
import {
  Command,
  CommandGroup,
  CommandItem,
  CommandList,
} from '@/components/ui/command';
import {
  Popover,
  PopoverAnchor,
  PopoverContent,
} from '@/components/ui/popover';
import { cn } from '@/lib/utils';
import { commentPlugin } from '@/components/editor/plugins/comment-kit';
 
import { AIChatEditor } from './ai-chat-editor';
 
export function AIMenu() {
  const { api, editor } = useEditorPlugin(AIChatPlugin);
  const mode = usePluginOption(AIChatPlugin, 'mode');
  const toolName = usePluginOption(AIChatPlugin, 'toolName');
 
  const streaming = usePluginOption(AIChatPlugin, 'streaming');
  const isSelecting = useIsSelecting();
  const isFocusedLast = useFocusedLast();
  const open = usePluginOption(AIChatPlugin, 'open') && isFocusedLast;
  const [value, setValue] = React.useState('');
 
  const [input, setInput] = React.useState('');
 
  const chat = usePluginOption(AIChatPlugin, 'chat');
 
  const { messages, status } = chat;
  const [anchorElement, setAnchorElement] = React.useState<HTMLElement | null>(
    null
  );
 
  const content = useLastAssistantMessage()?.parts.find(
    (part) => part.type === 'text'
  )?.text;
 
  React.useEffect(() => {
    if (!streaming) return;
 
    const anchorEntry = api.aiChat.node({ anchor: true });
    if (!anchorEntry) return;
 
    const anchorDom = editor.api.toDOMNode(anchorEntry[0])!;
    // eslint-disable-next-line react-hooks/set-state-in-effect -- Position the popover from editor DOM while the edit stream is active.
    setAnchorElement(anchorDom);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [streaming]);
 
  const setOpen = (open: boolean) => {
    if (open) {
      api.aiChat.show();
    } else {
      api.aiChat.hide();
    }
  };
 
  const show = (anchorElement: HTMLElement) => {
    setAnchorElement(anchorElement);
    setOpen(true);
  };
 
  useEditorChat({
    onOpenBlockSelection: (blocks: NodeEntry[]) => {
      show(editor.api.toDOMNode(blocks.at(-1)![0])!);
    },
    onOpenChange: (open) => {
      if (!open) {
        setAnchorElement(null);
        setInput('');
      }
    },
    onOpenCursor: () => {
      const [ancestor] = editor.api.block({ highest: true })!;
 
      if (!editor.api.isAt({ end: true }) && !editor.api.isEmpty(ancestor)) {
        editor
          .getApi(BlockSelectionPlugin)
          .blockSelection.set(ancestor.id as string);
      }
 
      show(editor.api.toDOMNode(ancestor)!);
    },
    onOpenSelection: () => {
      show(editor.api.toDOMNode(editor.api.blocks().at(-1)![0])!);
    },
  });
 
  useHotkeys('esc', () => {
    api.aiChat.stop();
 
    // remove when you implement the route /api/ai/command
    (chat as any)._abortFakeStream();
  });
 
  const isLoading = status === 'streaming' || status === 'submitted';
 
  React.useEffect(() => {
    if (toolName !== 'edit' || mode !== 'chat' || isLoading) return;
 
    let anchorNode = editor.api.node({
      at: [],
      reverse: true,
      match: (n) => !!n[KEYS.suggestion] && !!n[getTransientSuggestionKey()],
    });
 
    if (!anchorNode) {
      anchorNode = editor
        .getApi(BlockSelectionPlugin)
        .blockSelection.getNodes({ selectionFallback: true, sort: true })
        .at(-1);
    }
 
    if (!anchorNode) return;
 
    const block = editor.api.block({ at: anchorNode[1] });
    // eslint-disable-next-line react-hooks/set-state-in-effect -- Position the popover from editor DOM after the edit stream completes.
    setAnchorElement(editor.api.toDOMNode(block![0]!)!);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [isLoading]);
 
  if (isLoading && mode === 'insert') return null;
 
  if (toolName === 'comment') return null;
 
  if (toolName === 'edit' && mode === 'chat' && isLoading) return null;
 
  return (
    <Popover open={open} onOpenChange={setOpen} modal={false}>
      <PopoverAnchor virtualRef={{ current: anchorElement! }} />
 
      <PopoverContent
        className="border-none bg-transparent p-0 shadow-none"
        style={{
          width: anchorElement?.offsetWidth,
        }}
        onEscapeKeyDown={(e) => {
          e.preventDefault();
 
          api.aiChat.hide();
        }}
        align="center"
        side="bottom"
      >
        <Command
          className="w-full rounded-lg border shadow-md"
          value={value}
          onValueChange={setValue}
        >
          {mode === 'chat' &&
            isSelecting &&
            content &&
            toolName === 'generate' && <AIChatEditor content={content} />}
 
          {isLoading ? (
            <div className="flex grow select-none items-center gap-2 p-2 text-muted-foreground text-sm">
              <Loader2Icon className="size-4 animate-spin" />
              {messages.length > 1 ? 'Editing...' : 'Thinking...'}
            </div>
          ) : (
            <CommandPrimitive.Input
              className={cn(
                'flex h-9 w-full min-w-0 border-input bg-transparent px-3 py-1 text-base outline-none transition-[color,box-shadow] placeholder:text-muted-foreground md:text-sm dark:bg-input/30',
                'aria-invalid:border-destructive aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40',
                'border-b focus-visible:ring-transparent'
              )}
              value={input}
              onKeyDown={(e) => {
                if (isHotkey('backspace')(e) && input.length === 0) {
                  e.preventDefault();
                  api.aiChat.hide();
                }
                if (isHotkey('enter')(e) && !e.shiftKey && !value) {
                  e.preventDefault();
                  void api.aiChat.submit(input);
                  setInput('');
                }
              }}
              onValueChange={setInput}
              placeholder="Ask AI anything..."
              data-plate-focus
              autoFocus
            />
          )}
 
          {!isLoading && (
            <CommandList>
              <AIMenuItems
                input={input}
                setInput={setInput}
                setValue={setValue}
              />
            </CommandList>
          )}
        </Command>
      </PopoverContent>
    </Popover>
  );
}
 
type EditorChatState =
  | 'cursorCommand'
  | 'cursorSuggestion'
  | 'selectionCommand'
  | 'selectionSuggestion';
 
const AICommentIcon = () => (
  <svg
    fill="none"
    height="24"
    stroke="currentColor"
    strokeLinecap="round"
    strokeLinejoin="round"
    strokeWidth="2"
    viewBox="0 0 24 24"
    width="24"
    xmlns="http://www.w3.org/2000/svg"
  >
    <path d="M0 0h24v24H0z" fill="none" stroke="none" />
    <path d="M8 9h8" />
    <path d="M8 13h4.5" />
    <path d="M10 19l-1 -1h-3a3 3 0 0 1 -3 -3v-8a3 3 0 0 1 3 -3h12a3 3 0 0 1 3 3v4.5" />
    <path d="M17.8 20.817l-2.172 1.138a.392 .392 0 0 1 -.568 -.41l.415 -2.411l-1.757 -1.707a.389 .389 0 0 1 .217 -.665l2.428 -.352l1.086 -2.193a.392 .392 0 0 1 .702 0l1.086 2.193l2.428 .352a.39 .39 0 0 1 .217 .665l-1.757 1.707l.414 2.41a.39 .39 0 0 1 -.567 .411l-2.172 -1.138z" />
  </svg>
);
 
const aiChatItems = {
  accept: {
    icon: <Check />,
    label: 'Accept',
    value: 'accept',
    onSelect: ({ aiEditor, editor }) => {
      const { mode, toolName } = editor.getOptions(AIChatPlugin);
 
      if (mode === 'chat' && toolName === 'generate') {
        return editor
          .getTransforms(AIChatPlugin)
          .aiChat.replaceSelection(aiEditor);
      }
 
      editor.getTransforms(AIChatPlugin).aiChat.accept();
      editor.tf.focus({ edge: 'end' });
    },
  },
  comment: {
    icon: <AICommentIcon />,
    label: 'Comment',
    value: 'comment',
    onSelect: ({ editor, input }) => {
      editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt:
          'Please comment on the following content and provide reasonable and meaningful feedback.',
        toolName: 'comment',
      });
    },
  },
  continueWrite: {
    icon: <PenLine />,
    label: 'Continue writing',
    value: 'continueWrite',
    onSelect: ({ editor, input }) => {
      const ancestorNode = editor.api.block({ highest: true });
 
      if (!ancestorNode) return;
 
      const isEmpty = NodeApi.string(ancestorNode[0]).trim().length === 0;
 
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt: isEmpty
          ? `<Document>
{editor}
</Document>
Start writing a new paragraph AFTER <Document> ONLY ONE SENTENCE`
          : 'Continue writing AFTER <Block> ONLY ONE SENTENCE. DONT REPEAT THE TEXT.',
        toolName: 'generate',
      });
    },
  },
  discard: {
    icon: <X />,
    label: 'Discard',
    shortcut: 'Escape',
    value: 'discard',
    onSelect: ({ editor }) => {
      editor.getTransforms(AIPlugin).ai.undo();
      editor.getApi(AIChatPlugin).aiChat.hide();
    },
  },
  emojify: {
    icon: <SmileIcon />,
    label: 'Emojify',
    value: 'emojify',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Add a small number of contextually relevant emojis within each block only. You may insert emojis, but do not remove, replace, or rewrite existing text, and do not modify Markdown syntax, links, or line breaks.',
        toolName: 'edit',
      });
    },
  },
  explain: {
    icon: <BadgeHelp />,
    label: 'Explain',
    value: 'explain',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: {
          default: 'Explain {editor}',
          selecting: 'Explain',
        },
        toolName: 'generate',
      });
    },
  },
  fixSpelling: {
    icon: <Check />,
    label: 'Fix spelling & grammar',
    value: 'fixSpelling',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Fix spelling, grammar, and punctuation errors within each block only, without changing meaning, tone, or adding new information.',
        toolName: 'edit',
      });
    },
  },
  generateMarkdownSample: {
    icon: <BookOpenCheck />,
    label: 'Generate Markdown sample',
    value: 'generateMarkdownSample',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: 'Generate a markdown sample',
        toolName: 'generate',
      });
    },
  },
  generateMdxSample: {
    icon: <BookOpenCheck />,
    label: 'Generate MDX sample',
    value: 'generateMdxSample',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt: 'Generate a mdx sample',
        toolName: 'generate',
      });
    },
  },
  improveWriting: {
    icon: <Wand />,
    label: 'Improve writing',
    value: 'improveWriting',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Improve the writing for clarity and flow, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  insertBelow: {
    icon: <ListEnd />,
    label: 'Insert below',
    value: 'insertBelow',
    onSelect: ({ aiEditor, editor }) => {
      /** Format: 'none' Fix insert table */
      void editor
        .getTransforms(AIChatPlugin)
        .aiChat.insertBelow(aiEditor, { format: 'none' });
    },
  },
  makeLonger: {
    icon: <ListPlus />,
    label: 'Make longer',
    value: 'makeLonger',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Make the content longer by elaborating on existing ideas within each block only, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  makeShorter: {
    icon: <ListMinus />,
    label: 'Make shorter',
    value: 'makeShorter',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Make the content shorter by reducing verbosity within each block only, without changing meaning or removing essential information.',
        toolName: 'edit',
      });
    },
  },
  replace: {
    icon: <Check />,
    label: 'Replace selection',
    value: 'replace',
    onSelect: ({ aiEditor, editor }) => {
      void editor.getTransforms(AIChatPlugin).aiChat.replaceSelection(aiEditor);
    },
  },
  simplifyLanguage: {
    icon: <FeatherIcon />,
    label: 'Simplify language',
    value: 'simplifyLanguage',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        prompt:
          'Simplify the language by using clearer and more straightforward wording within each block only, without changing meaning or adding new information.',
        toolName: 'edit',
      });
    },
  },
  summarize: {
    icon: <Album />,
    label: 'Add a summary',
    value: 'summarize',
    onSelect: ({ editor, input }) => {
      void editor.getApi(AIChatPlugin).aiChat.submit(input, {
        mode: 'insert',
        prompt: {
          default: 'Summarize {editor}',
          selecting: 'Summarize',
        },
        toolName: 'generate',
      });
    },
  },
  tryAgain: {
    icon: <CornerUpLeft />,
    label: 'Try again',
    value: 'tryAgain',
    onSelect: ({ editor }) => {
      void editor.getApi(AIChatPlugin).aiChat.reload();
    },
  },
} satisfies Record<
  string,
  {
    icon: React.ReactNode;
    label: string;
    value: string;
    component?: React.ComponentType<{ menuState: EditorChatState }>;
    filterItems?: boolean;
    items?: { label: string; value: string }[];
    shortcut?: string;
    onSelect?: ({
      aiEditor,
      editor,
      input,
    }: {
      aiEditor: SlateEditor;
      editor: PlateEditor;
      input: string;
    }) => void;
  }
>;
 
const menuStateItems: Record<
  EditorChatState,
  {
    items: (typeof aiChatItems)[keyof typeof aiChatItems][];
    heading?: string;
  }[]
> = {
  cursorCommand: [
    {
      items: [
        aiChatItems.comment,
        aiChatItems.generateMdxSample,
        aiChatItems.generateMarkdownSample,
        aiChatItems.continueWrite,
        aiChatItems.summarize,
        aiChatItems.explain,
      ],
    },
  ],
  cursorSuggestion: [
    {
      items: [aiChatItems.accept, aiChatItems.discard, aiChatItems.tryAgain],
    },
  ],
  selectionCommand: [
    {
      items: [
        aiChatItems.improveWriting,
        aiChatItems.comment,
        aiChatItems.emojify,
        aiChatItems.makeLonger,
        aiChatItems.makeShorter,
        aiChatItems.fixSpelling,
        aiChatItems.simplifyLanguage,
      ],
    },
  ],
  selectionSuggestion: [
    {
      items: [
        aiChatItems.accept,
        aiChatItems.discard,
        aiChatItems.insertBelow,
        aiChatItems.tryAgain,
      ],
    },
  ],
};
 
export const AIMenuItems = ({
  input,
  setInput,
  setValue,
}: {
  input: string;
  setInput: (value: string) => void;
  setValue: (value: string) => void;
}) => {
  const editor = useEditorRef();
  const { messages } = usePluginOption(AIChatPlugin, 'chat');
  const aiEditor = usePluginOption(AIChatPlugin, 'aiEditor')!;
  const isSelecting = useIsSelecting();
 
  const menuState = React.useMemo(() => {
    if (messages && messages.length > 0) {
      return isSelecting ? 'selectionSuggestion' : 'cursorSuggestion';
    }
 
    return isSelecting ? 'selectionCommand' : 'cursorCommand';
  }, [isSelecting, messages]);
 
  const menuGroups = React.useMemo(() => {
    const items = menuStateItems[menuState];
 
    return items;
  }, [menuState]);
 
  React.useEffect(() => {
    if (menuGroups.length > 0 && menuGroups[0].items.length > 0) {
      setValue(menuGroups[0].items[0].value);
    }
  }, [menuGroups, setValue]);
 
  return (
    <>
      {menuGroups.map((group, index) => (
        <CommandGroup key={index} heading={group.heading}>
          {group.items.map((menuItem) => (
            <CommandItem
              key={menuItem.value}
              className="[&_svg]:text-muted-foreground"
              value={menuItem.value}
              onSelect={() => {
                menuItem.onSelect?.({
                  aiEditor,
                  editor,
                  input,
                });
                setInput('');
              }}
            >
              {menuItem.icon}
              <span>{menuItem.label}</span>
            </CommandItem>
          ))}
        </CommandGroup>
      ))}
    </>
  );
};
 
export function AILoadingBar() {
  const editor = useEditorRef();
 
  const toolName = usePluginOption(AIChatPlugin, 'toolName');
  const chat = usePluginOption(AIChatPlugin, 'chat');
  const mode = usePluginOption(AIChatPlugin, 'mode');
 
  const { status } = chat;
 
  const { api } = useEditorPlugin(AIChatPlugin);
 
  const isLoading = status === 'streaming' || status === 'submitted';
 
  const handleComments = (type: 'accept' | 'reject') => {
    if (type === 'accept') {
      editor.tf.unsetNodes([getTransientCommentKey()], {
        at: [],
        match: (n) => TextApi.isText(n) && !!n[KEYS.comment],
      });
    }
 
    if (type === 'reject') {
      editor
        .getTransforms(commentPlugin)
        .comment.unsetMark({ transient: true });
    }
 
    api.aiChat.hide();
  };
 
  useHotkeys('esc', () => {
    api.aiChat.stop();
 
    // remove when you implement the route /api/ai/command
    (chat as any)._abortFakeStream();
  });
 
  if (
    isLoading &&
    (mode === 'insert' ||
      toolName === 'comment' ||
      (toolName === 'edit' && mode === 'chat'))
  ) {
    return (
      <div
        className={cn(
          '-translate-x-1/2 absolute bottom-4 left-1/2 z-20 flex items-center gap-3 rounded-md border border-border bg-muted px-3 py-1.5 text-muted-foreground text-sm shadow-md transition-all duration-300'
        )}
      >
        <span className="h-4 w-4 animate-spin rounded-full border-2 border-muted-foreground border-t-transparent" />
        <span>{status === 'submitted' ? 'Thinking...' : 'Writing...'}</span>
        <Button
          size="sm"
          variant="ghost"
          className="flex items-center gap-1 text-xs"
          onClick={() => api.aiChat.stop()}
        >
          <PauseIcon className="h-4 w-4" />
          Stop
          <kbd className="ml-1 rounded bg-border px-1 font-mono text-[10px] text-muted-foreground shadow-sm">
            Esc
          </kbd>
        </Button>
      </div>
    );
  }
 
  if (toolName === 'comment' && status === 'ready') {
    return (
      <div
        className={cn(
          '-translate-x-1/2 absolute bottom-4 left-1/2 z-50 flex flex-col items-center gap-0 rounded-xl border border-border/50 bg-popover p-1 text-muted-foreground text-sm shadow-xl backdrop-blur-sm',
          'p-3'
        )}
      >
        {/* Header with controls */}
        <div className="flex w-full items-center justify-between gap-3">
          <div className="flex items-center gap-5">
            <Button
              size="sm"
              disabled={isLoading}
              onClick={() => handleComments('accept')}
            >
              Accept
            </Button>
 
            <Button
              size="sm"
              disabled={isLoading}
              onClick={() => handleComments('reject')}
            >
              Reject
            </Button>
          </div>
        </div>
      </div>
    );
  }
 
  return null;
}

扩展 aiChatItems 映射以添加新命令。每个命令接收 { aiEditor, editor, input } 并可以使用自定义提示或变换调度 api.aiChat.submit。

简单自定义命令

summarizeInBullets: {
  icon: <ListIcon />,
  label: 'Summarize in bullets',
  value: 'summarizeInBullets',
  onSelect: ({ editor }) => {
    void editor.getApi(AIChatPlugin).aiChat.submit('', {
      prompt: 'Summarize the current selection using bullet points',
      toolName: 'generate',
    });
  },
},
summarizeInBullets: {
  icon: <ListIcon />,
  label: 'Summarize in bullets',
  value: 'summarizeInBullets',
  onSelect: ({ editor }) => {
    void editor.getApi(AIChatPlugin).aiChat.submit('', {
      prompt: 'Summarize the current selection using bullet points',
      toolName: 'generate',
    });
  },
},

带复杂逻辑的命令

generateTOC: {
  icon: <BookIcon />,
  label: 'Generate table of contents',
  value: 'generateTOC',
  onSelect: ({ editor }) => {
    const headings = editor.api.nodes({
      match: (n) => ['h1', 'h2', 'h3'].includes(n.type as string),
    });
 
    const prompt =
      headings.length === 0
        ? 'Create a realistic table of contents for this document'
        : 'Generate a table of contents that reflects the existing headings';
 
    void editor.getApi(AIChatPlugin).aiChat.submit('', {
      mode: 'insert',
      prompt,
      toolName: 'generate',
    });
  },
},
generateTOC: {
  icon: <BookIcon />,
  label: 'Generate table of contents',
  value: 'generateTOC',
  onSelect: ({ editor }) => {
    const headings = editor.api.nodes({
      match: (n) => ['h1', 'h2', 'h3'].includes(n.type as string),
    });
 
    const prompt =
      headings.length === 0
        ? 'Create a realistic table of contents for this document'
        : 'Generate a table of contents that reflects the existing headings';
 
    void editor.getApi(AIChatPlugin).aiChat.submit('', {
      mode: 'insert',
      prompt,
      toolName: 'generate',
    });
  },
},

菜单自动在命令和建议状态之间切换：

  • cursorCommand：光标折叠且尚无响应。
  • selectionCommand：文本被选中且尚无响应。
  • cursorSuggestion / selectionSuggestion：存在响应，因此显示接受、重试或在下方插入等操作。

使用 toolName（'generate' | 'edit' | 'comment'）来控制流式 hooks 如何处理响应。例如，'edit' 启用基于差异的建议，'comment' 允许您使用 aiCommentToRange 将流式评论转换为讨论线程。