| 1 | "use client"; |
| 2 | |
| 3 | import { useChat, type UseChatHelpers } from "@ai-sdk/react"; |
| 4 | import { withAIBatch } from "@platejs/ai"; |
| 5 | import { |
| 6 | AIChatPlugin, |
| 7 | AIPlugin, |
| 8 | applyAISuggestions, |
| 9 | streamInsertChunk, |
| 10 | useChatChunk, |
| 11 | type AIChatPluginConfig, |
| 12 | } from "@platejs/ai/react"; |
| 13 | import { DefaultChatTransport, type UIMessage } from "ai"; |
| 14 | import { getPluginType, KEYS, PathApi } from "platejs"; |
| 15 | import { usePluginOption } from "platejs/react"; |
| 16 | import { useEffect, useMemo } from "react"; |
| 17 | |
| 18 | import { AILoadingBar, AIMenu } from "@/components/plate/ui/ai-menu"; |
| 19 | import { AIAnchorElement, AILeaf } from "@/components/plate/ui/ai-node"; |
| 20 | import { CursorOverlayKit } from "./cursor-overlay-kit"; |
| 21 | import { MarkdownKit } from "./markdown-kit"; |
| 22 | |
| 23 | type PlateChatMessageData = { |
| 24 | toolName: "comment" | "edit" | "generate"; |
| 25 | comment?: { |
| 26 | blockId: string; |
| 27 | comment: string; |
| 28 | content: string; |
| 29 | }; |
| 30 | }; |
| 31 | |
| 32 | type PlateChatMessage = UIMessage<Record<string, never>, PlateChatMessageData>; |
| 33 | |
| 34 | const aiChatPlugin = AIChatPlugin.extend({ |
| 35 | render: { |
| 36 | afterContainer: AILoadingBar, |
| 37 | afterEditable: AIMenu, |
| 38 | node: AIAnchorElement, |
| 39 | }, |
| 40 | shortcuts: { show: { keys: "mod+j" } }, |
| 41 | useHooks: ({ editor, getOption }) => { |
| 42 | const rawChat = useChat<PlateChatMessage>({ |
| 43 | transport: new DefaultChatTransport({ |
| 44 | api: "/api/ai/command", |
| 45 | }), |
| 46 | }); |
| 47 | const chat = useMemo<UseChatHelpers<PlateChatMessage>>( |
| 48 | () => ({ |
| 49 | id: rawChat.id, |
| 50 | messages: rawChat.messages, |
| 51 | setMessages: rawChat.setMessages, |
| 52 | error: rawChat.error, |
| 53 | sendMessage: rawChat.sendMessage, |
| 54 | regenerate: rawChat.regenerate, |
| 55 | stop: rawChat.stop, |
| 56 | resumeStream: rawChat.resumeStream, |
| 57 | addToolResult: rawChat.addToolResult, |
| 58 | addToolOutput: rawChat.addToolOutput, |
| 59 | addToolApprovalResponse: rawChat.addToolApprovalResponse, |
| 60 | status: rawChat.status, |
| 61 | clearError: rawChat.clearError, |
| 62 | }), |
| 63 | [ |
| 64 | rawChat.id, |
| 65 | rawChat.messages, |
| 66 | rawChat.setMessages, |
| 67 | rawChat.error, |
| 68 | rawChat.sendMessage, |
| 69 | rawChat.regenerate, |
| 70 | rawChat.stop, |
| 71 | rawChat.resumeStream, |
| 72 | rawChat.addToolResult, |
| 73 | rawChat.addToolOutput, |
| 74 | rawChat.addToolApprovalResponse, |
| 75 | rawChat.status, |
| 76 | rawChat.clearError, |
| 77 | ], |
| 78 | ); |
| 79 | const mode = usePluginOption( |
| 80 | { key: KEYS.aiChat } as AIChatPluginConfig, |
| 81 | "mode", |
| 82 | ); |
| 83 | const toolName = usePluginOption( |
| 84 | { key: KEYS.aiChat } as AIChatPluginConfig, |
| 85 | "toolName", |
| 86 | ); |
| 87 | |
| 88 | useEffect(() => { |
| 89 | if (editor.getOption(AIChatPlugin, "chat") === chat) { |
| 90 | return; |
| 91 | } |
| 92 | |
| 93 | editor.setOption(AIChatPlugin, "chat", chat); |
| 94 | }, [chat, editor]); |
| 95 | |
| 96 | useChatChunk({ |
| 97 | onChunk: ({ chunk, isFirst, nodes, text }) => { |
| 98 | if (isFirst && mode == "insert") { |
| 99 | const focusPath = editor.selection?.focus.path; |
| 100 | const insertPath = focusPath |
| 101 | ? PathApi.next(focusPath.slice(0, 1)) |
| 102 | : [editor.children.length]; |
| 103 | |
| 104 | editor.tf.withoutSaving(() => { |
| 105 | editor.tf.insertNodes( |
| 106 | { |
| 107 | children: [{ text: "" }], |
| 108 | type: getPluginType(editor, KEYS.aiChat), |
| 109 | }, |
| 110 | { |
| 111 | at: insertPath, |
| 112 | }, |
| 113 | ); |
| 114 | }); |
| 115 | editor.setOption(AIChatPlugin, "streaming", true); |
| 116 | } |
| 117 | |
| 118 | if (mode === "insert" && nodes.length > 0) { |
| 119 | withAIBatch( |
| 120 | editor, |
| 121 | () => { |
| 122 | if (!getOption("streaming")) return; |
| 123 | editor.tf.withScrolling(() => { |
| 124 | streamInsertChunk(editor, chunk, { |
| 125 | textProps: { |
| 126 | [getPluginType(editor, KEYS.ai)]: true, |
| 127 | }, |
| 128 | }); |
| 129 | }); |
| 130 | }, |
| 131 | { split: isFirst }, |
| 132 | ); |
| 133 | } |
| 134 | |
| 135 | if (toolName === "edit" && mode === "chat") { |
| 136 | withAIBatch( |
| 137 | editor, |
| 138 | () => { |
| 139 | applyAISuggestions(editor, text); |
| 140 | }, |
| 141 | { split: isFirst }, |
| 142 | ); |
| 143 | } |
| 144 | }, |
| 145 | onFinish: () => { |
| 146 | editor.setOption(AIChatPlugin, "streaming", false); |
| 147 | editor.setOption(AIChatPlugin, "_blockChunks", ""); |
| 148 | editor.setOption(AIChatPlugin, "_blockPath", null); |
| 149 | editor.setOption(AIChatPlugin, "_mdxName", null); |
| 150 | }, |
| 151 | }); |
| 152 | }, |
| 153 | }); |
| 154 | |
| 155 | export const AIKit = [ |
| 156 | ...CursorOverlayKit, |
| 157 | ...MarkdownKit, |
| 158 | AIPlugin.withComponent(AILeaf), |
| 159 | aiChatPlugin, |
| 160 | ]; |
| 161 | |
| 162 | const systemCommon = `\ |
| 163 | You are an advanced AI-powered note-taking assistant, designed to enhance productivity and creativity in note management. |
| 164 | Respond directly to user prompts with clear, concise, and relevant content. Maintain a neutral, helpful tone. |
| 165 | |
| 166 | Rules: |
| 167 | - <Document> is the entire note the user is working on. |
| 168 | - <Reminder> is a reminder of how you should reply to INSTRUCTIONS. It does not apply to questions. |
| 169 | - Anything else is the user prompt. |
| 170 | - Your response should be tailored to the user's prompt, providing precise assistance to optimize note management. |
| 171 | - For INSTRUCTIONS: Follow the <Reminder> exactly. Provide ONLY the content to be inserted or replaced. No explanations or comments. |
| 172 | - For QUESTIONS: Provide a helpful and concise answer. You may include brief explanations if necessary. |
| 173 | - CRITICAL: DO NOT remove or modify the following custom MDX tags: <u>, <callout>, <kbd>, <toc>, <sub>, <sup>, <mark>, <del>, <date>, <span>, <column>, <column_group>, <file>, <audio>, <video> in <Selection> unless the user explicitly requests this change. |
| 174 | - CRITICAL: Distinguish between INSTRUCTIONS and QUESTIONS. Instructions typically ask you to modify or add content. Questions ask for information or clarification. |
| 175 | - CRITICAL: when asked to write in markdown, do not start with \`\`\`markdown. |
| 176 | `; |
| 177 | |
| 178 | const systemDefault = `\ |
| 179 | ${systemCommon} |
| 180 | - <Block> is the current block of text the user is working on. |
| 181 | - Ensure your output can seamlessly fit into the existing <Block> structure. |
| 182 | |
| 183 | <Block> |
| 184 | {block} |
| 185 | </Block> |
| 186 | `; |
| 187 | |
| 188 | const systemSelecting = `\ |
| 189 | ${systemCommon} |
| 190 | - <Block> is the block of text containing the user's selection, providing context. |
| 191 | - Ensure your output can seamlessly fit into the existing <Block> structure. |
| 192 | - <Selection> is the specific text the user has selected in the block and wants to modify or ask about. |
| 193 | - Consider the context provided by <Block>, but only modify <Selection>. Your response should be a direct replacement for <Selection>. |
| 194 | <Block> |
| 195 | {block} |
| 196 | </Block> |
| 197 | <Selection> |
| 198 | {selection} |
| 199 | </Selection> |
| 200 | `; |
| 201 | |
| 202 | const systemBlockSelecting = `\ |
| 203 | ${systemCommon} |
| 204 | - <Selection> represents the full blocks of text the user has selected and wants to modify or ask about. |
| 205 | - Your response should be a direct replacement for the entire <Selection>. |
| 206 | - Maintain the overall structure and formatting of the selected blocks, unless explicitly instructed otherwise. |
| 207 | - CRITICAL: Provide only the content to replace <Selection>. Do not add additional blocks or change the block structure unless specifically requested. |
| 208 | <Selection> |
| 209 | {block} |
| 210 | </Selection> |
| 211 | `; |
| 212 | |
| 213 | const userDefault = `<Reminder> |
| 214 | CRITICAL: NEVER write <Block>. |
| 215 | </Reminder> |
| 216 | {prompt}`; |
| 217 | const userSelecting = `<Reminder> |
| 218 | If this is a question, provide a helpful and concise answer about <Selection>. |
| 219 | If this is an instruction, provide ONLY the text to replace <Selection>. No explanations. |
| 220 | Ensure it fits seamlessly within <Block>. If <Block> is empty, write ONE random sentence. |
| 221 | NEVER write <Block> or <Selection>. |
| 222 | </Reminder> |
| 223 | {prompt} about <Selection>`; |
| 224 | |
| 225 | const userBlockSelecting = `<Reminder> |
| 226 | If this is a question, provide a helpful and concise answer about <Selection>. |
| 227 | If this is an instruction, provide ONLY the content to replace the entire <Selection>. No explanations. |
| 228 | Maintain the overall structure unless instructed otherwise. |
| 229 | NEVER write <Block> or <Selection>. |
| 230 | </Reminder> |
| 231 | {prompt} about <Selection>`; |
| 232 | |
| 233 | export const PROMPT_TEMPLATES = { |
| 234 | systemBlockSelecting, |
| 235 | systemDefault, |
| 236 | systemSelecting, |
| 237 | userBlockSelecting, |
| 238 | userDefault, |
| 239 | userSelecting, |
| 240 | }; |
| 241 |