Add canvas interface (#461)

app/(chat)/api/chat/route.ts

@@ -1,17 +1,50 @@
-import { convertToCoreMessages, Message, streamText } from 'ai';
+import {
+  convertToCoreMessages,
+  generateObject,
+  Message,
+  StreamData,
+  streamObject,
+  streamText,
+} from 'ai';
 import { z } from 'zod';
 
 import { customModel } from '@/ai';
+import { models } from '@/ai/models';
+import { canvasPrompt, regularPrompt } from '@/ai/prompts';
 import { auth } from '@/app/(auth)/auth';
-import { deleteChatById, getChatById, saveChat } from '@/db/queries';
-import { Model, models } from '@/lib/model';
+import {
+  deleteChatById,
+  getChatById,
+  getDocumentById,
+  saveChat,
+  saveDocument,
+  saveSuggestions,
+} from '@/db/queries';
+import { Suggestion } from '@/db/schema';
+import { generateUUID, sanitizeResponseMessages } from '@/lib/utils';
+
+export const maxDuration = 60;
+
+type AllowedTools =
+  | 'createDocument'
+  | 'updateDocument'
+  | 'requestSuggestions'
+  | 'getWeather';
+
+const canvasTools: AllowedTools[] = [
+  'createDocument',
+  'updateDocument',
+  'requestSuggestions',
+];
+
+const weatherTools: AllowedTools[] = ['getWeather'];
 
 export async function POST(request: Request) {
   const {
     id,
     messages,
-    model,
-  }: { id: string; messages: Array<Message>; model: Model['name'] } =
+    modelId,
+  }: { id: string; messages: Array<Message>; modelId: string } =
     await request.json();
 
   const session = await auth();
@@ -20,18 +53,22 @@ export async function POST(request: Request) {
     return new Response('Unauthorized', { status: 401 });
   }
 
-  if (!models.find((m) => m.name === model)) {
+  const model = models.find((model) => model.id === modelId);
+
+  if (!model) {
     return new Response('Model not found', { status: 404 });
   }
 
   const coreMessages = convertToCoreMessages(messages);
+  const streamingData = new StreamData();
 
   const result = await streamText({
-    model: customModel(model),
-    system:
-      'you are a friendly assistant! keep your responses concise and helpful.',
+    model: customModel(model.apiIdentifier),
+    system: modelId === 'gpt-4o-canvas' ? canvasPrompt : regularPrompt,
     messages: coreMessages,
     maxSteps: 5,
+    experimental_activeTools:
+      modelId === 'gpt-4o-canvas' ? canvasTools : weatherTools,
     tools: {
       getWeather: {
         description: 'Get the current weather at a location',
@@ -48,19 +85,233 @@ export async function POST(request: Request) {
           return weatherData;
         },
       },
+      createDocument: {
+        description: 'Create a document for a writing activity',
+        parameters: z.object({
+          title: z.string(),
+        }),
+        execute: async ({ title }) => {
+          const id = generateUUID();
+          let draftText: string = '';
+
+          streamingData.append({
+            type: 'id',
+            content: id,
+          });
+
+          streamingData.append({
+            type: 'title',
+            content: title,
+          });
+
+          streamingData.append({
+            type: 'clear',
+            content: '',
+          });
+
+          const { fullStream } = await streamText({
+            model: customModel(model.apiIdentifier),
+            system:
+              'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
+            prompt: title,
+          });
+
+          for await (const delta of fullStream) {
+            const { type } = delta;
+
+            if (type === 'text-delta') {
+              const { textDelta } = delta;
+
+              draftText += textDelta;
+              streamingData.append({
+                type: 'text-delta',
+                content: textDelta,
+              });
+            }
+          }
+
+          streamingData.append({ type: 'finish', content: '' });
+
+          if (session.user && session.user.id) {
+            await saveDocument({
+              id,
+              title,
+              content: draftText,
+              userId: session.user.id,
+            });
+          }
+
+          return {
+            id,
+            title,
+            content: `A document was created and is now visible to the user.`,
+          };
+        },
+      },
+      updateDocument: {
+        description: 'Update a document with the given description',
+        parameters: z.object({
+          id: z.string().describe('The ID of the document to update'),
+          description: z
+            .string()
+            .describe('The description of changes that need to be made'),
+        }),
+        execute: async ({ id, description }) => {
+          const document = await getDocumentById({ id });
+
+          if (!document) {
+            return {
+              error: 'Document not found',
+            };
+          }
+
+          const { content: currentContent } = document;
+          let draftText: string = '';
+
+          streamingData.append({
+            type: 'clear',
+            content: document.title,
+          });
+
+          const { fullStream } = await streamText({
+            model: customModel(model.apiIdentifier),
+            system:
+              'You are a helpful writing assistant. Based on the description, please update the piece of writing.',
+            messages: [
+              {
+                role: 'user',
+                content: description,
+              },
+              { role: 'user', content: currentContent },
+            ],
+          });
+
+          for await (const delta of fullStream) {
+            const { type } = delta;
+
+            if (type === 'text-delta') {
+              const { textDelta } = delta;
+
+              draftText += textDelta;
+              streamingData.append({
+                type: 'text-delta',
+                content: textDelta,
+              });
+            }
+          }
+
+          streamingData.append({ type: 'finish', content: '' });
+
+          if (session.user && session.user.id) {
+            await saveDocument({
+              id,
+              title: document.title,
+              content: draftText,
+              userId: session.user.id,
+            });
+          }
+
+          return {
+            id,
+            title: document.title,
+            content: 'The document has been updated successfully.',
+          };
+        },
+      },
+      requestSuggestions: {
+        description: 'Request suggestions for a document',
+        parameters: z.object({
+          documentId: z
+            .string()
+            .describe('The ID of the document to request edits'),
+        }),
+        execute: async ({ documentId }) => {
+          const document = await getDocumentById({ id: documentId });
+
+          if (!document || !document.content) {
+            return {
+              error: 'Document not found',
+            };
+          }
+
+          let suggestions: Array<
+            Omit<Suggestion, 'userId' | 'createdAt' | 'documentCreatedAt'>
+          > = [];
+
+          const { elementStream } = await streamObject({
+            model: customModel(model.apiIdentifier),
+            system:
+              'You are a help writing assistant. Given a piece of writing, please offer suggestions to improve the piece of writing and describe the change. It is very important for the edits to contain full sentences instead of just words.',
+            prompt: document.content,
+            output: 'array',
+            schema: z.object({
+              originalSentence: z.string().describe('The original sentence'),
+              suggestedSentence: z.string().describe('The suggested sentence'),
+              description: z
+                .string()
+                .describe('The description of the suggestion'),
+            }),
+          });
+
+          for await (const element of elementStream) {
+            const suggestion = {
+              originalText: element.originalSentence,
+              suggestedText: element.suggestedSentence,
+              description: element.description,
+              id: generateUUID(),
+              documentId: documentId,
+              isResolved: false,
+            };
+
+            streamingData.append({
+              type: 'suggestion',
+              content: suggestion,
+            });
+
+            suggestions.push(suggestion);
+          }
+
+          if (session.user && session.user.id) {
+            const userId = session.user.id;
+
+            await saveSuggestions({
+              suggestions: suggestions.map((suggestion) => ({
+                ...suggestion,
+                userId,
+                createdAt: new Date(),
+                documentCreatedAt: document.createdAt,
+              })),
+            });
+          }
+
+          return {
+            id: documentId,
+            title: document.title,
+            message: 'Suggestions have been added to the document',
+          };
+        },
+      },
     },
     onFinish: async ({ responseMessages }) => {
       if (session.user && session.user.id) {
         try {
+          const responseMessagesWithoutIncompleteToolCalls =
+            sanitizeResponseMessages(responseMessages);
+
           await saveChat({
             id,
-            messages: [...coreMessages, ...responseMessages],
+            messages: [
+              ...coreMessages,
+              ...responseMessagesWithoutIncompleteToolCalls,
+            ],
             userId: session.user.id,
           });
         } catch (error) {
           console.error('Failed to save chat');
         }
       }
+
+      streamingData.close();
     },
     experimental_telemetry: {
       isEnabled: true,
@@ -68,7 +319,9 @@ export async function POST(request: Request) {
     },
   });
 
-  return result.toDataStreamResponse({});
+  return result.toDataStreamResponse({
+    data: streamingData,
+  });
 }
 
 export async function DELETE(request: Request) {

app/(chat)/api/document/route.ts

@@ -0,0 +1,98 @@
+import { auth } from '@/app/(auth)/auth';
+import {
+  deleteDocumentsByIdAfterTimestamp,
+  getDocumentsById,
+  saveDocument,
+} from '@/db/queries';
+
+export async function GET(request: Request) {
+  const { searchParams } = new URL(request.url);
+  const id = searchParams.get('id');
+
+  if (!id) {
+    return new Response('Missing id', { status: 400 });
+  }
+
+  const session = await auth();
+
+  if (!session || !session.user) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  const documents = await getDocumentsById({ id });
+
+  const [document] = documents;
+
+  if (!document) {
+    return new Response('Not Found', { status: 404 });
+  }
+
+  if (document.userId !== session.user.id) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  return Response.json(documents, { status: 200 });
+}
+
+export async function POST(request: Request) {
+  const { searchParams } = new URL(request.url);
+  const id = searchParams.get('id');
+
+  if (!id) {
+    return new Response('Missing id', { status: 400 });
+  }
+
+  const session = await auth();
+
+  if (!session) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  const { content, title }: { content: string; title: string } =
+    await request.json();
+
+  if (session.user && session.user.id) {
+    const document = await saveDocument({
+      id,
+      content,
+      title,
+      userId: session.user.id,
+    });
+
+    return Response.json(document, { status: 200 });
+  } else {
+    return new Response('Unauthorized', { status: 401 });
+  }
+}
+
+export async function PATCH(request: Request) {
+  const { searchParams } = new URL(request.url);
+  const id = searchParams.get('id');
+
+  const { timestamp }: { timestamp: string } = await request.json();
+
+  if (!id) {
+    return new Response('Missing id', { status: 400 });
+  }
+
+  const session = await auth();
+
+  if (!session || !session.user) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  const documents = await getDocumentsById({ id });
+
+  const [document] = documents;
+
+  if (document.userId !== session.user.id) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  await deleteDocumentsByIdAfterTimestamp({
+    id,
+    timestamp: new Date(timestamp),
+  });
+
+  return new Response('Deleted', { status: 200 });
+}

app/(chat)/api/suggestions/route.ts

@@ -0,0 +1,33 @@
+import { auth } from '@/app/(auth)/auth';
+import { getSuggestionsByDocumentId } from '@/db/queries';
+
+export async function GET(request: Request) {
+  const { searchParams } = new URL(request.url);
+  const documentId = searchParams.get('documentId');
+
+  if (!documentId) {
+    return new Response('Not Found', { status: 404 });
+  }
+
+  const session = await auth();
+
+  if (!session || !session.user) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  const suggestions = await getSuggestionsByDocumentId({
+    documentId,
+  });
+
+  const [suggestion] = suggestions;
+
+  if (!suggestion) {
+    return Response.json([], { status: 200 });
+  }
+
+  if (suggestion.userId !== session.user.id) {
+    return new Response('Unauthorized', { status: 401 });
+  }
+
+  return Response.json(suggestions, { status: 200 });
+}