+ "content": "import type {\n ChatMessage,\n ToolName,\n} from '@/registry/components/editor/use-chat';\nimport type { NextRequest } from 'next/server';\n\nimport { google } from '@ai-sdk/google';\nimport { createOpenAI } from '@ai-sdk/openai';\nimport { replacePlaceholders } from '@platejs/ai';\nimport {\n convertToModelMessages,\n createUIMessageStream,\n createUIMessageStreamResponse,\n generateObject,\n streamObject,\n streamText,\n} from 'ai';\nimport { NextResponse } from 'next/server';\nimport { type SlateEditor, createSlateEditor, nanoid, RangeApi } from 'platejs';\nimport { z } from 'zod';\n\nimport { BaseEditorKit } from '@/registry/components/editor/editor-base-kit';\nimport { markdownJoinerTransform } from '@/registry/lib/markdown-joiner-transform';\n\nexport async function POST(req: NextRequest) {\n const {\n apiKey: key,\n ctx,\n messages: messagesParam,\n prompt,\n system,\n } = await req.json();\n\n const { children, selection, toolName: toolNameParam } = ctx;\n\n const editor = createSlateEditor({\n plugins: BaseEditorKit,\n selection,\n value: children,\n });\n\n const apiKey = key || process.env.OPENAI_API_KEY;\n\n if (!apiKey) {\n return NextResponse.json(\n { error: 'Missing OpenAI API key.' },\n { status: 401 }\n );\n }\n\n const openai = createOpenAI({ apiKey });\n\n const isSelecting = editor.api.isExpanded();\n\n const isBlockSelecting = isSelectingAllBlocks(editor);\n\n try {\n const stream = createUIMessageStream<ChatMessage>({\n execute: async ({ writer }) => {\n const messages = replaceMessagePlaceholders(editor, messagesParam, {\n isSelecting,\n });\n\n const lastUserMessage = messages.findLast(\n (message: any) => message.role === 'user'\n );\n\n let toolName = toolNameParam;\n\n if (!toolName) {\n const { object: AIToolName } = await generateObject({\n enum: ['generate', 'edit', 'comment'],\n model: google('gemini-2.5-flash'),\n output: 'enum',\n prompt: `User message:\n ${JSON.stringify(lastUserMessage)}`,\n system: chooseToolSystem,\n });\n\n writer.write({\n data: AIToolName as ToolName,\n type: 'data-toolName',\n });\n\n toolName = AIToolName;\n }\n\n if (toolName === 'generate') {\n const generateSystem = replacePlaceholders(\n editor,\n systemTemplate({ isBlockSelecting, isSelecting }),\n {\n prompt: system,\n }\n );\n\n const gen = streamText({\n experimental_transform: markdownJoinerTransform(),\n maxOutputTokens: 2048,\n messages: convertToModelMessages(messages),\n model: google('gemini-2.5-flash'),\n system: generateSystem,\n });\n\n writer.merge(gen.toUIMessageStream({ sendFinish: false }));\n }\n\n if (toolName === 'edit') {\n const editSystem = replacePlaceholders(\n editor,\n systemTemplate({ isBlockSelecting, isSelecting }),\n {\n prompt: system,\n }\n );\n\n const edit = streamText({\n experimental_transform: markdownJoinerTransform(),\n maxOutputTokens: 2048,\n messages: convertToModelMessages(messages),\n model: google('gemini-2.5-flash'),\n system: editSystem,\n });\n\n writer.merge(edit.toUIMessageStream({ sendFinish: false }));\n }\n\n if (toolName === 'comment') {\n const commentPrompt = replacePlaceholders(\n editor,\n commentTemplate({ isSelecting }),\n {\n prompt,\n }\n );\n\n const { elementStream } = streamObject({\n maxOutputTokens: 2048,\n model: openai('gpt-4o'),\n output: 'array',\n prompt: commentPrompt,\n schema: z\n .object({\n blockId: z\n .string()\n .describe(\n 'The id of the starting block. If the comment spans multiple blocks, use the id of the first block.'\n ),\n comment: z\n .string()\n .describe(\n 'A brief comment or explanation for this fragment.'\n ),\n content: z\n .string()\n .describe(\n 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.`\n ),\n })\n .describe('A single comment'),\n system: commentSystem,\n });\n\n // Create a single message ID for the entire comment stream\n\n for await (const comment of elementStream) {\n const commentDataId = nanoid();\n // Send each comment as a delta\n\n writer.write({\n id: commentDataId,\n data: comment,\n type: 'data-comment',\n });\n }\n\n return;\n }\n },\n });\n\n return createUIMessageStreamResponse({ stream });\n } catch {\n return NextResponse.json(\n { error: 'Failed to process AI request' },\n { status: 500 }\n );\n }\n}\n\nconst systemTemplate = ({\n isBlockSelecting,\n isSelecting,\n}: {\n isBlockSelecting: boolean;\n isSelecting: boolean;\n}) => {\n return isBlockSelecting\n ? PROMPT_TEMPLATES.systemBlockSelecting\n : isSelecting\n ? PROMPT_TEMPLATES.systemSelecting\n : PROMPT_TEMPLATES.systemDefault;\n};\n\nconst promptTemplate = ({ isSelecting }: { isSelecting: boolean }) => {\n return isSelecting\n ? PROMPT_TEMPLATES.userSelecting\n : PROMPT_TEMPLATES.userDefault;\n};\n\nconst commentTemplate = ({ isSelecting }: { isSelecting: boolean }) => {\n return isSelecting\n ? PROMPT_TEMPLATES.commentSelecting\n : PROMPT_TEMPLATES.commentDefault;\n};\n\nconst chooseToolSystem = `You are a strict classifier. Classify the user's last request as \"generate\", \"edit\", or \"comment\".\n\nPriority rules:\n1. Default is \"generate\". Any open question, idea request, or creation request → \"generate\".\n2. Only return \"edit\" if the user provides original text (or a selection of text) AND asks to change, rephrase, translate, or shorten it.\n3. Only return \"comment\" if the user explicitly asks for comments, feedback, annotations, or review. Do not infer \"comment\" implicitly.\n\nReturn only one enum value with no explanation.`;\n\nconst commentSystem = `You are a document review assistant. \nYou will receive an MDX document wrapped in <block id=\"...\"> content </block> tags. \n\nYour task: \n- Read the content of all blocks and provide comments. \n- For each comment, generate a JSON object: \n - blockId: the id of the block being commented on.\n - content: the original document fragment that needs commenting.\n - comments: a brief comment or explanation for that fragment.\n\nRules:\n- IMPORTANT: If a comment spans multiple blocks, use the id of the **first** block.\n- The **content** field must be the original content inside the block tag. The returned content must not include the block tags, but should retain other MDX tags.\n- IMPORTANT: The **content** field must be flexible:\n - It can cover one full block, only part of a block, or multiple blocks. \n - If multiple blocks are included, separate them with two \\\\n\\\\n. \n - Do NOT default to using the entire block—use the smallest relevant span instead.\n- At least one comment must be provided.\n`;\n\nconst systemCommon = `\\\nYou are an advanced AI-powered note-taking assistant, designed to enhance productivity and creativity in note management.\nRespond directly to user prompts with clear, concise, and relevant content. Maintain a neutral, helpful tone.\n\nRules:\n- <Document> is the entire note the user is working on.\n- <Reminder> is a reminder of how you should reply to INSTRUCTIONS. It does not apply to questions.\n- Anything else is the user prompt.\n- Your response should be tailored to the user's prompt, providing precise assistance to optimize note management.\n- For INSTRUCTIONS: Follow the <Reminder> exactly. Provide ONLY the content to be inserted or replaced. No explanations or comments.\n- For QUESTIONS: Provide a helpful and concise answer. You may include brief explanations if necessary.\n- 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.\n- CRITICAL: Distinguish between INSTRUCTIONS and QUESTIONS. Instructions typically ask you to modify or add content. Questions ask for information or clarification.\n- CRITICAL: when asked to write in markdown, do not start with \\`\\`\\`markdown.\n- CRITICAL: When writing the column, such line breaks and indentation must be preserved.\n<column_group>\n <column>\n 1\n </column>\n <column>\n 2\n </column>\n <column>\n 3\n </column>\n</column_group>\n`;\n\nconst systemDefault = `\\\n${systemCommon}\n- <Block> is the current block of text the user is working on.\n- Ensure your output can seamlessly fit into the existing <Block> structure.\n\n<Block>\n{block}\n</Block>\n`;\n\nconst systemSelecting = `\\\n${systemCommon}\n- <Block> is the block of text containing the user's selection, providing context.\n- Ensure your output can seamlessly fit into the existing <Block> structure.\n- <Selection> is the specific text the user has selected in the block and wants to modify or ask about.\n- Consider the context provided by <Block>, but only modify <Selection>. Your response should be a direct replacement for <Selection>.\n<Block>\n{block}\n</Block>\n<Selection>\n{selection}\n</Selection>\n`;\n\nconst systemBlockSelecting = `\\\n${systemCommon}\n- <Selection> represents the full blocks of text the user has selected and wants to modify or ask about.\n- Your response should be a direct replacement for the entire <Selection>.\n- Maintain the overall structure and formatting of the selected blocks, unless explicitly instructed otherwise.\n- CRITICAL: Provide only the content to replace <Selection>. Do not add additional blocks or change the block structure unless specifically requested.\n<Selection>\n{block}\n</Selection>\n`;\n\nconst userDefault = `<Reminder>\nCRITICAL: NEVER write <Block>.\n</Reminder>\n{prompt}`;\nconst userSelecting = `<Reminder>\nIf this is a question, provide a helpful and concise answer about <Selection>.\nIf this is an instruction, provide ONLY the text to replace <Selection>. No explanations.\nEnsure it fits seamlessly within <Block>. If <Block> is empty, write ONE random sentence.\nNEVER write <Block> or <Selection>.\n</Reminder>\n{prompt} about <Selection>`;\n\nconst commentSelecting = `{prompt}:\n \n{blockWithBlockId}\n`;\n\nconst commentDefault = `{prompt}:\n \n{editorWithBlockId}\n`;\n\nconst PROMPT_TEMPLATES = {\n commentDefault,\n commentSelecting,\n systemBlockSelecting,\n systemDefault,\n systemSelecting,\n userDefault,\n userSelecting,\n};\n\nconst replaceMessagePlaceholders = (\n editor: SlateEditor,\n messages: ChatMessage[],\n { isSelecting }: { isSelecting: boolean }\n): ChatMessage[] => {\n const template = promptTemplate({ isSelecting });\n\n const lastUserIndex = messages\n .map((m) => (m as any).role)\n .lastIndexOf('user');\n\n if (lastUserIndex === -1) return messages;\n\n const targetMessage = messages[lastUserIndex];\n const parts = targetMessage.parts.map((part) => {\n if (part.type !== 'text' || !part.text) return part;\n\n const text = replacePlaceholders(editor, template, {\n prompt: part.text,\n });\n\n return { ...part, text } as typeof part;\n });\n\n const updatedMessage = { ...targetMessage, parts };\n\n return [\n ...messages.slice(0, lastUserIndex),\n updatedMessage,\n ...messages.slice(lastUserIndex + 1),\n ];\n};\n\n/** Check if the current selection fully covers all top-level blocks. */\nconst isSelectingAllBlocks = (editor: SlateEditor) => {\n const blocksRange = editor.api.nodesRange(\n editor.api.blocks({ mode: 'highest' })\n );\n\n return (\n !!blocksRange &&\n !!editor.selection &&\n RangeApi.equals(blocksRange, editor.selection)\n );\n};\n",
0 commit comments