+ "content": "import type {\n ChatMessage,\n ToolName,\n} from '@/registry/components/editor/use-chat';\nimport type { NextRequest } from 'next/server';\n\nimport { createGateway } from '@ai-sdk/gateway';\nimport {\n type LanguageModel,\n type UIMessageStreamWriter,\n createUIMessageStream,\n createUIMessageStreamResponse,\n generateObject,\n streamObject,\n streamText,\n tool,\n} from 'ai';\nimport { NextResponse } from 'next/server';\nimport { type SlateEditor, createSlateEditor, nanoid } 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\nimport {\n getChooseToolPrompt,\n getCommentPrompt,\n getEditPrompt,\n getGeneratePrompt,\n} from './prompts';\n\nexport async function POST(req: NextRequest) {\n const {\n apiKey: key,\n ctx,\n messages: messagesRaw = [],\n model,\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.AI_GATEWAY_API_KEY;\n\n if (!apiKey) {\n return NextResponse.json(\n { error: 'Missing AI Gateway API key.' },\n { status: 401 }\n );\n }\n\n const isSelecting = editor.api.isExpanded();\n\n const gatewayProvider = createGateway({\n apiKey,\n });\n\n try {\n const stream = createUIMessageStream<ChatMessage>({\n execute: async ({ writer }) => {\n let toolName = toolNameParam;\n\n if (!toolName) {\n const { object: AIToolName } = await generateObject({\n enum: isSelecting\n ? ['generate', 'edit', 'comment']\n : ['generate', 'comment'],\n model: gatewayProvider(model || 'google/gemini-2.5-flash'),\n output: 'enum',\n prompt: getChooseToolPrompt(messagesRaw),\n });\n\n writer.write({\n data: AIToolName as ToolName,\n type: 'data-toolName',\n });\n\n toolName = AIToolName;\n }\n\n const stream = streamText({\n experimental_transform: markdownJoinerTransform(),\n model: gatewayProvider(model || 'openai/gpt-4o-mini'),\n // Not used\n prompt: '',\n tools: {\n comment: getCommentTool(editor, {\n messagesRaw,\n model: gatewayProvider(model || 'google/gemini-2.5-flash'),\n writer,\n }),\n },\n prepareStep: async (step) => {\n if (toolName === 'comment') {\n return {\n ...step,\n toolChoice: { toolName: 'comment', type: 'tool' },\n };\n }\n\n if (toolName === 'edit') {\n const editPrompt = getEditPrompt(editor, {\n isSelecting,\n messages: messagesRaw,\n });\n\n return {\n ...step,\n activeTools: [],\n messages: [\n {\n content: editPrompt,\n role: 'user',\n },\n ],\n };\n }\n\n if (toolName === 'generate') {\n const generatePrompt = getGeneratePrompt(editor, {\n messages: messagesRaw,\n });\n\n return {\n ...step,\n activeTools: [],\n messages: [\n {\n content: generatePrompt,\n role: 'user',\n },\n ],\n model: gatewayProvider(model || 'openai/gpt-4o-mini'),\n };\n }\n },\n });\n\n writer.merge(stream.toUIMessageStream({ sendFinish: false }));\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 getCommentTool = (\n editor: SlateEditor,\n {\n messagesRaw,\n model,\n writer,\n }: {\n messagesRaw: ChatMessage[];\n model: LanguageModel;\n writer: UIMessageStreamWriter<ChatMessage>;\n }\n) => {\n return tool({\n description: 'Comment on the content',\n inputSchema: z.object({}),\n execute: async () => {\n const { elementStream } = streamObject({\n model,\n output: 'array',\n prompt: getCommentPrompt(editor, {\n messages: messagesRaw,\n }),\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('A brief comment or explanation for this fragment.'),\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 });\n\n for await (const comment of elementStream) {\n const commentDataId = nanoid();\n\n writer.write({\n id: commentDataId,\n data: {\n comment: comment,\n status: 'streaming',\n },\n type: 'data-comment',\n });\n }\n\n writer.write({\n id: nanoid(),\n data: {\n comment: null,\n status: 'finished',\n },\n type: 'data-comment',\n });\n },\n });\n};\n",
0 commit comments