Add canvas interface (#461)

ai/custom-middleware.ts

@@ -1,3 +1,3 @@
-import { Experimental_LanguageModelV1Middleware } from "ai";
+import { Experimental_LanguageModelV1Middleware } from 'ai';
 
 export const customMiddleware: Experimental_LanguageModelV1Middleware = {};

ai/index.ts

@@ -1,13 +1,11 @@
 import { openai } from '@ai-sdk/openai';
 import { experimental_wrapLanguageModel as wrapLanguageModel } from 'ai';
 
-import { type Model } from '@/lib/model';
-
 import { customMiddleware } from './custom-middleware';
 
-export const customModel = (modelName: Model['name']) => {
+export const customModel = (apiIdentifier: string) => {
   return wrapLanguageModel({
-    model: openai(modelName),
+    model: openai(apiIdentifier),
     middleware: customMiddleware,
   });
 };

ai/models.ts

@@ -0,0 +1,31 @@
+// Define your models here.
+
+export interface Model {
+  id: string;
+  label: string;
+  apiIdentifier: string;
+  description: string;
+}
+
+export const models: Array<Model> = [
+  {
+    id: 'gpt-4o-mini',
+    label: 'GPT 4o mini',
+    apiIdentifier: 'gpt-4o-mini',
+    description: 'Small model for fast, lightweight tasks',
+  },
+  {
+    id: 'gpt-4o',
+    label: 'GPT 4o',
+    apiIdentifier: 'gpt-4o',
+    description: 'For complex, multi-step tasks',
+  },
+  {
+    id: 'gpt-4o-canvas',
+    label: 'GPT 4o with Canvas',
+    apiIdentifier: 'gpt-4o',
+    description: 'Collaborate with writing',
+  },
+] as const;
+
+export const DEFAULT_MODEL_NAME: string = 'gpt-4o-mini';

ai/prompts.ts

@@ -0,0 +1,26 @@
+export const canvasPrompt = `
+  Canvas is a special user interface mode that helps users with writing, editing, and other content creation tasks. When canvas is open, it is on the right side of the screen, while the conversation is on the left side. When creating or updating documents, changes are reflected in real-time on the canvas and visible to the user.
+
+  This is a guide for using canvas tools: \`createDocument\` and \`updateDocument\`, which render content on a canvas beside the conversation.
+
+  **When to use \`createDocument\`:**
+  - For substantial content (>10 lines)
+  - For content users will likely save/reuse (emails, code, essays, etc.)
+  - When explicitly requested to create a document
+
+  **When NOT to use \`createDocument\`:**
+  - For short content (<10 lines)
+  - For informational/explanatory content
+  - For conversational responses
+  - When asked to keep it in chat
+
+  **Using \`updateDocument\`:**
+  - Default to full document rewrites for major changes
+  - Use targeted updates only for specific, isolated changes
+  - Follow user instructions for which parts to modify
+
+  Do not update document right after creating it. Wait for user feedback or request to update it.
+  `;
+
+export const regularPrompt =
+  'You are a friendly assistant! Keep your responses concise and helpful.';

app/(chat)/actions.ts

@@ -2,7 +2,7 @@
 
 import { cookies } from 'next/headers';
 
-export async function saveModel(model: string) {
+export async function saveModelId(model: string) {
   const cookieStore = await cookies();
-  cookieStore.set('model', model);
+  cookieStore.set('model-id', model);
 }

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 {
-import { Model, models } from '@/lib/model';
+  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,
+    modelId,
-  }: { id: string; messages: Array<Message>; model: Model['name'] } =
+  }: { 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),
+    model: customModel(model.apiIdentifier),
-    system:
+    system: modelId === 'gpt-4o-canvas' ? canvasPrompt : regularPrompt,
-      'you are a friendly assistant! keep your responses concise and helpful.',
     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 });
+}

app/(chat)/chat/[id]/page.tsx

@@ -2,11 +2,11 @@ import { CoreMessage } from 'ai';
 import { cookies } from 'next/headers';
 import { notFound } from 'next/navigation';
 
+import { DEFAULT_MODEL_NAME, models } from '@/ai/models';
 import { auth } from '@/app/(auth)/auth';
 import { Chat as PreviewChat } from '@/components/custom/chat';
 import { getChatById } from '@/db/queries';
 import { Chat } from '@/db/schema';
-import { DEFAULT_MODEL_NAME, models } from '@/lib/model';
 import { convertToUIMessages } from '@/lib/utils';
 
 export default async function Page(props: { params: Promise<any> }) {
@@ -35,15 +35,16 @@ export default async function Page(props: { params: Promise<any> }) {
   }
 
   const cookieStore = await cookies();
-  const value = cookieStore.get('model')?.value;
+  const modelIdFromCookie = cookieStore.get('model-id')?.value;
-  const selectedModelName =
+  const selectedModelId =
-    models.find((m) => m.name === value)?.name || DEFAULT_MODEL_NAME;
+    models.find((model) => model.id === modelIdFromCookie)?.id ||
+    DEFAULT_MODEL_NAME;
 
   return (
     <PreviewChat
       id={chat.id}
       initialMessages={chat.messages}
-      selectedModelName={selectedModelName}
+      selectedModelId={selectedModelId}
     />
   );
 }

app/(chat)/page.tsx

@@ -1,23 +1,25 @@
 import { cookies } from 'next/headers';
 
+import { DEFAULT_MODEL_NAME, models } from '@/ai/models';
 import { Chat } from '@/components/custom/chat';
-import { DEFAULT_MODEL_NAME, models } from '@/lib/model';
 import { generateUUID } from '@/lib/utils';
 
 export default async function Page() {
   const id = generateUUID();
 
   const cookieStore = await cookies();
-  const value = cookieStore.get('model')?.value;
+  const modelIdFromCookie = cookieStore.get('model-id')?.value;
-  const selectedModelName =
+
-    models.find((m) => m.name === value)?.name || DEFAULT_MODEL_NAME;
+  const selectedModelId =
+    models.find((model) => model.id === modelIdFromCookie)?.id ||
+    DEFAULT_MODEL_NAME;
 
   return (
     <Chat
       key={id}
       id={id}
       initialMessages={[]}
-      selectedModelName={selectedModelName}
+      selectedModelId={selectedModelId}
     />
   );
 }

app/globals.css

@@ -136,3 +136,7 @@
     @apply bg-foreground/20 animate-pulse;
   }
 }
+
+.ProseMirror {
+  outline: none;
+}

components/custom/canvas-stream-handler.tsx

@@ -0,0 +1,43 @@
+import { JSONValue } from 'ai';
+import { Dispatch, memo, SetStateAction } from 'react';
+
+import { UICanvas } from './canvas';
+import { useCanvasStream } from './use-canvas-stream';
+
+interface CanvasStreamHandlerProps {
+  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  streamingData: JSONValue[] | undefined;
+}
+
+export function PureCanvasStreamHandler({
+  setCanvas,
+  streamingData,
+}: CanvasStreamHandlerProps) {
+  useCanvasStream({
+    streamingData,
+    setCanvas,
+  });
+
+  return null;
+}
+
+function areEqual(
+  prevProps: CanvasStreamHandlerProps,
+  nextProps: CanvasStreamHandlerProps
+) {
+  if (!prevProps.streamingData && !nextProps.streamingData) {
+    return true;
+  }
+
+  if (!prevProps.streamingData || !nextProps.streamingData) {
+    return false;
+  }
+
+  if (prevProps.streamingData.length !== nextProps.streamingData.length) {
+    return false;
+  }
+
+  return true;
+}
+
+export const CanvasStreamHandler = memo(PureCanvasStreamHandler, areEqual);

components/custom/canvas.tsx

@@ -0,0 +1,545 @@
+import { Attachment, ChatRequestOptions, CreateMessage, Message } from 'ai';
+import cx from 'classnames';
+import { formatDistance, isAfter } from 'date-fns';
+import { AnimatePresence, motion } from 'framer-motion';
+import {
+  Dispatch,
+  SetStateAction,
+  useCallback,
+  useEffect,
+  useState,
+} from 'react';
+import useSWR, { useSWRConfig } from 'swr';
+
+import { Document, Suggestion } from '@/db/schema';
+import { fetcher } from '@/lib/utils';
+
+import { DiffView } from './diffview';
+import { DocumentSkeleton } from './document-skeleton';
+import { Editor } from './editor';
+import { CrossIcon, DeltaIcon, RedoIcon, UndoIcon } from './icons';
+import { Message as PreviewMessage } from './message';
+import { MultimodalInput } from './multimodal-input';
+import { Toolbar } from './toolbar';
+import { useDebounce } from './use-debounce';
+import { useScrollToBottom } from './use-scroll-to-bottom';
+import useWindowSize from './use-window-size';
+import { Button } from '../ui/button';
+
+export interface UICanvas {
+  title: string;
+  documentId: string;
+  content: string;
+  isVisible: boolean;
+  status: 'streaming' | 'idle';
+  boundingBox: {
+    top: number;
+    left: number;
+    width: number;
+    height: number;
+  };
+}
+
+export function Canvas({
+  input,
+  setInput,
+  handleSubmit,
+  isLoading,
+  stop,
+  attachments,
+  setAttachments,
+  append,
+  canvas,
+  setCanvas,
+  messages,
+  setMessages,
+}: {
+  input: string;
+  setInput: (input: string) => void;
+  isLoading: boolean;
+  stop: () => void;
+  attachments: Array<Attachment>;
+  setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
+  canvas: UICanvas;
+  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+  messages: Array<Message>;
+  setMessages: Dispatch<SetStateAction<Array<Message>>>;
+  append: (
+    message: Message | CreateMessage,
+    chatRequestOptions?: ChatRequestOptions
+  ) => Promise<string | null | undefined>;
+  handleSubmit: (
+    event?: {
+      preventDefault?: () => void;
+    },
+    chatRequestOptions?: ChatRequestOptions
+  ) => void;
+}) {
+  const [messagesContainerRef, messagesEndRef] =
+    useScrollToBottom<HTMLDivElement>();
+
+  const {
+    data: documents,
+    isLoading: isDocumentsFetching,
+    isValidating: isDocumentsValidating,
+    mutate: mutateDocuments,
+  } = useSWR<Array<Document>>(
+    canvas && canvas.status !== 'streaming'
+      ? `/api/document?id=${canvas.documentId}`
+      : null,
+    fetcher
+  );
+
+  const { data: suggestions } = useSWR<Array<Suggestion>>(
+    documents && canvas && canvas.status !== 'streaming'
+      ? `/api/suggestions?documentId=${canvas.documentId}`
+      : null,
+    fetcher,
+    {
+      dedupingInterval: 5000,
+    }
+  );
+
+  const [mode, setMode] = useState<'edit' | 'diff'>('edit');
+  const [document, setDocument] = useState<Document | null>(null);
+  const [currentVersionIndex, setCurrentVersionIndex] = useState(-1);
+
+  useEffect(() => {
+    if (documents && documents.length > 0) {
+      const mostRecentDocument = documents.at(-1);
+
+      if (mostRecentDocument) {
+        setDocument(mostRecentDocument);
+        setCurrentVersionIndex(documents.length - 1);
+        setCanvas((currentCanvas) =>
+          currentCanvas
+            ? {
+                ...currentCanvas,
+                content: mostRecentDocument.content ?? '',
+              }
+            : null
+        );
+      }
+    }
+  }, [documents, setCanvas]);
+
+  useEffect(() => {
+    mutateDocuments();
+  }, [canvas.status, mutateDocuments]);
+
+  const { mutate } = useSWRConfig();
+  const [isContentDirty, setIsContentDirty] = useState(false);
+
+  const handleContentChange = useCallback(
+    (updatedContent: string) => {
+      if (!canvas) return;
+
+      mutate<Array<Document>>(
+        `/api/document?id=${canvas.documentId}`,
+        async (currentDocuments) => {
+          if (!currentDocuments) return undefined;
+
+          const currentDocument = currentDocuments.at(-1);
+
+          if (!currentDocument || !currentDocument.content) {
+            setIsContentDirty(false);
+            return currentDocuments;
+          }
+
+          if (currentDocument.content !== updatedContent) {
+            await fetch(`/api/document?id=${canvas.documentId}`, {
+              method: 'POST',
+              body: JSON.stringify({
+                title: canvas.title,
+                content: updatedContent,
+              }),
+            });
+
+            setIsContentDirty(false);
+
+            const newDocument = {
+              ...currentDocument,
+              content: updatedContent,
+              createdAt: new Date(),
+            };
+
+            return [...currentDocuments, newDocument];
+          } else {
+            return currentDocuments;
+          }
+        },
+        { revalidate: false }
+      );
+    },
+    [canvas, mutate]
+  );
+
+  const debouncedHandleEditorChange = useCallback(
+    useDebounce(handleContentChange, 4000),
+    [handleContentChange]
+  );
+
+  const handleEditorChange = useCallback(
+    (updatedContent: string) => {
+      if (document && updatedContent !== document.content) {
+        debouncedHandleEditorChange(updatedContent);
+        setIsContentDirty(true);
+      }
+    },
+    [document, debouncedHandleEditorChange]
+  );
+
+  function getDocumentContentById(index: number) {
+    if (!documents) return '';
+    if (!documents[index]) return '';
+    return documents[index].content ?? '';
+  }
+
+  function getDocumentTimestampById(index: number) {
+    if (!documents) return '';
+    return documents[index]?.createdAt ?? '';
+  }
+
+  const handleVersionChange = (type: 'next' | 'prev' | 'toggle' | 'latest') => {
+    if (!documents) return;
+
+    if (type === 'latest') {
+      setCurrentVersionIndex(documents.length - 1);
+      setMode('edit');
+    }
+
+    if (type === 'toggle') {
+      setMode((mode) => (mode === 'edit' ? 'diff' : 'edit'));
+    }
+
+    if (type === 'prev') {
+      if (currentVersionIndex > 0) {
+        setCurrentVersionIndex((index) => index - 1);
+      }
+    } else if (type === 'next') {
+      if (currentVersionIndex < documents.length - 1) {
+        setCurrentVersionIndex((index) => index + 1);
+      }
+    }
+  };
+
+  const [isToolbarVisible, setIsToolbarVisible] = useState(false);
+
+  /*
+   * NOTE: if there are no documents, or if
+   * the documents are being fetched, then
+   * we mark it as the current version.
+   */
+
+  const isCurrentVersion =
+    isDocumentsFetching || isDocumentsValidating
+      ? true
+      : documents && documents.length > 0
+        ? currentVersionIndex === documents.length - 1
+        : true;
+
+  const { width: windowWidth, height: windowHeight } = useWindowSize();
+  const isMobile = windowWidth ? windowWidth < 768 : false;
+
+  return (
+    <motion.div
+      className="flex flex-row h-dvh w-dvw fixed top-0 left-0 z-50 bg-muted"
+      initial={{ opacity: 1 }}
+      animate={{ opacity: 1 }}
+      exit={{ opacity: 0, transition: { delay: 0.4 } }}
+    >
+      {!isMobile && (
+        <motion.div
+          className="relative w-[400px] bg-muted dark:bg-background h-dvh shrink-0"
+          initial={{ opacity: 0, x: 10, scale: 1 }}
+          animate={{
+            opacity: 1,
+            x: 0,
+            scale: 1,
+            transition: {
+              delay: 0.2,
+              type: 'spring',
+              stiffness: 200,
+              damping: 30,
+            },
+          }}
+          exit={{ opacity: 0, x: 10, scale: 0.95, transition: { delay: 0 } }}
+        >
+          <div className="-right-12 w-12 bg-muted absolute h-dvh top-0" />
+          <div className="flex flex-col h-full justify-between items-center gap-4">
+            <div
+              ref={messagesContainerRef}
+              className="flex flex-col gap-4 h-full items-center overflow-y-scroll px-4 pt-20"
+            >
+              {messages.map((message) => (
+                <PreviewMessage
+                  key={message.id}
+                  role={message.role}
+                  content={message.content}
+                  attachments={message.experimental_attachments}
+                  toolInvocations={message.toolInvocations}
+                  canvas={canvas}
+                  setCanvas={setCanvas}
+                />
+              ))}
+
+              <div
+                ref={messagesEndRef}
+                className="shrink-0 min-w-[24px] min-h-[24px]"
+              />
+            </div>
+
+            <form className="flex flex-row gap-2 relative items-end w-full px-4 pb-4">
+              <MultimodalInput
+                input={input}
+                setInput={setInput}
+                handleSubmit={handleSubmit}
+                isLoading={isLoading}
+                stop={stop}
+                attachments={attachments}
+                setAttachments={setAttachments}
+                messages={messages}
+                append={append}
+                className="bg-background dark:bg-muted"
+                setMessages={setMessages}
+              />
+            </form>
+          </div>
+        </motion.div>
+      )}
+
+      <motion.div
+        className="fixed dark:bg-muted bg-background h-dvh flex flex-col shadow-xl overflow-y-scroll"
+        initial={
+          isMobile
+            ? {
+                opacity: 0,
+                x: 0,
+                y: 0,
+                width: windowWidth,
+                height: windowHeight,
+                borderRadius: 50,
+              }
+            : {
+                opacity: 0,
+                x: canvas.boundingBox.left,
+                y: canvas.boundingBox.top,
+                height: canvas.boundingBox.height,
+                width: canvas.boundingBox.width,
+                borderRadius: 50,
+              }
+        }
+        animate={
+          isMobile
+            ? {
+                opacity: 1,
+                x: 0,
+                y: 0,
+                width: windowWidth,
+                height: '100dvh',
+                borderRadius: 0,
+                transition: {
+                  delay: 0,
+                  type: 'spring',
+                  stiffness: 200,
+                  damping: 30,
+                },
+              }
+            : {
+                opacity: 1,
+                x: 400,
+                y: 0,
+                height: windowHeight,
+                width: windowWidth ? windowWidth - 400 : 'calc(100dvw-400px)',
+                borderRadius: 0,
+                transition: {
+                  delay: 0,
+                  type: 'spring',
+                  stiffness: 200,
+                  damping: 30,
+                },
+              }
+        }
+        exit={{
+          opacity: 0,
+          scale: 0.5,
+          transition: {
+            delay: 0.1,
+            type: 'spring',
+            stiffness: 600,
+            damping: 30,
+          },
+        }}
+      >
+        <div className="p-2 flex flex-row justify-between items-start">
+          <div className="flex flex-row gap-4 items-start">
+            <div
+              className="cursor-pointer hover:bg-muted p-2 rounded-lg text-muted-foreground"
+              onClick={() => {
+                setCanvas(null);
+              }}
+            >
+              <CrossIcon size={18} />
+            </div>
+
+            <div className="flex flex-col pt-1">
+              <div className="font-medium">
+                {document?.title ?? canvas.title}
+              </div>
+
+              {isContentDirty ? (
+                <div className="text-sm text-muted-foreground">
+                  Saving changes...
+                </div>
+              ) : document ? (
+                <div className="text-sm text-muted-foreground">
+                  {`Updated ${formatDistance(
+                    new Date(document.createdAt),
+                    new Date(),
+                    {
+                      addSuffix: true,
+                    }
+                  )}`}
+                </div>
+              ) : null}
+            </div>
+          </div>
+
+          <div className="flex flex-row gap-1">
+            <div
+              className="cursor-pointer hover:bg-muted p-2 rounded-lg text-muted-foreground"
+              onClick={() => {
+                handleVersionChange('prev');
+              }}
+            >
+              <UndoIcon size={18} />
+            </div>
+            <div
+              className="cursor-pointer hover:bg-muted p-2 rounded-lg text-muted-foreground"
+              onClick={() => {
+                handleVersionChange('next');
+              }}
+            >
+              <RedoIcon size={18} />
+            </div>
+            <div
+              className={cx(
+                'cursor-pointer hover:bg-muted p-2 rounded-lg text-muted-foreground',
+                { 'bg-muted': mode === 'diff' }
+              )}
+              onClick={() => {
+                handleVersionChange('toggle');
+              }}
+            >
+              <DeltaIcon size={18} />
+            </div>
+          </div>
+        </div>
+
+        <div className="prose dark:prose-invert dark:bg-muted bg-background h-full overflow-y-scroll px-4 py-8 md:p-20 !max-w-full pb-40 items-center">
+          <div className="flex flex-row max-w-[600px] mx-auto">
+            {isDocumentsFetching && !canvas.content ? (
+              <DocumentSkeleton />
+            ) : mode === 'edit' ? (
+              <Editor
+                content={
+                  isCurrentVersion
+                    ? canvas.content
+                    : getDocumentContentById(currentVersionIndex)
+                }
+                currentVersionIndex={currentVersionIndex}
+                status={canvas.status ?? 'idle'}
+                onChange={handleEditorChange}
+                suggestions={isCurrentVersion ? (suggestions ?? []) : []}
+              />
+            ) : (
+              <DiffView
+                oldContent={getDocumentContentById(currentVersionIndex - 1)}
+                newContent={getDocumentContentById(currentVersionIndex)}
+              />
+            )}
+
+            {suggestions ? (
+              <div className="md:hidden h-dvh w-12 shrink-0" />
+            ) : null}
+          </div>
+        </div>
+
+        <AnimatePresence>
+          {!isCurrentVersion && (
+            <motion.div
+              className="absolute flex flex-col gap-4 lg:flex-row bottom-0 bg-background p-4 w-full border-t z-50 justify-between"
+              initial={{ y: isMobile ? 200 : 77 }}
+              animate={{ y: 0 }}
+              exit={{ y: isMobile ? 200 : 77 }}
+              transition={{ type: 'spring', stiffness: 140, damping: 20 }}
+            >
+              <div>
+                <div>You are viewing a previous version</div>
+                <div className="text-muted-foreground text-sm">
+                  Restore this version to make edits
+                </div>
+              </div>
+
+              <div className="flex flex-row gap-4">
+                <Button
+                  onClick={async () => {
+                    mutate(
+                      `/api/document?id=${canvas.documentId}`,
+                      await fetch(`/api/document?id=${canvas.documentId}`, {
+                        method: 'PATCH',
+                        body: JSON.stringify({
+                          timestamp:
+                            getDocumentTimestampById(currentVersionIndex),
+                        }),
+                      }),
+                      {
+                        optimisticData: documents
+                          ? [
+                              ...documents.filter((d) =>
+                                isAfter(
+                                  new Date(d.createdAt),
+                                  new Date(
+                                    getDocumentTimestampById(
+                                      currentVersionIndex
+                                    )
+                                  )
+                                )
+                              ),
+                            ]
+                          : [],
+                      }
+                    );
+                  }}
+                >
+                  Restore this version
+                </Button>
+                <Button
+                  variant="outline"
+                  onClick={() => {
+                    handleVersionChange('latest');
+                  }}
+                >
+                  Back to latest version
+                </Button>
+              </div>
+            </motion.div>
+          )}
+        </AnimatePresence>
+      </motion.div>
+
+      <AnimatePresence>
+        {isCurrentVersion && (
+          <Toolbar
+            isToolbarVisible={isToolbarVisible}
+            setIsToolbarVisible={setIsToolbarVisible}
+            append={append}
+            isLoading={isLoading}
+            stop={stop}
+            setMessages={setMessages}
+          />
+        )}
+      </AnimatePresence>
+    </motion.div>
+  );
+}

components/custom/chat-header.tsx

@@ -5,13 +5,8 @@ import { ModelSelector } from '@/components/custom/model-selector';
 import { SidebarToggle } from '@/components/custom/sidebar-toggle';
 import { Button } from '@/components/ui/button';
 import { BetterTooltip } from '@/components/ui/tooltip';
-import { Model } from '@/lib/model';
 
-export function ChatHeader({
+export function ChatHeader({ selectedModelId }: { selectedModelId: string }) {
-  selectedModelName,
-}: {
-  selectedModelName: Model['name'];
-}) {
   return (
     <header className="flex h-16 sticky top-0 bg-background md:h-12 items-center px-2 md:px-2 z-10">
       <SidebarToggle />
@@ -28,7 +23,7 @@ export function ChatHeader({
         </Button>
       </BetterTooltip>
       <ModelSelector
-        selectedModelName={selectedModelName}
+        selectedModelId={selectedModelId}
         className="order-1 md:order-2"
       />
     </header>

components/custom/chat.tsx

@@ -2,42 +2,57 @@
 
 import { Attachment, Message } from 'ai';
 import { useChat } from 'ai/react';
+import { AnimatePresence } from 'framer-motion';
 import { useState } from 'react';
 
+import { Model } from '@/ai/models';
 import { ChatHeader } from '@/components/custom/chat-header';
 import { Message as PreviewMessage } from '@/components/custom/message';
 import { useScrollToBottom } from '@/components/custom/use-scroll-to-bottom';
-import { Model } from '@/lib/model';
 
+import { Canvas, UICanvas } from './canvas';
+import { CanvasStreamHandler } from './canvas-stream-handler';
 import { MultimodalInput } from './multimodal-input';
 import { Overview } from './overview';
 
 export function Chat({
   id,
   initialMessages,
-  selectedModelName,
+  selectedModelId,
 }: {
   id: string;
   initialMessages: Array<Message>;
-  selectedModelName: Model['name'];
+  selectedModelId: string;
 }) {
-  const { messages, handleSubmit, input, setInput, append, isLoading, stop } =
+  const {
-    useChat({
+    messages,
-      body: { id, model: selectedModelName },
+    setMessages,
+    handleSubmit,
+    input,
+    setInput,
+    append,
+    isLoading,
+    stop,
+    data: streamingData,
+  } = useChat({
+    body: { id, modelId: selectedModelId },
     initialMessages,
     onFinish: () => {
       window.history.replaceState({}, '', `/chat/${id}`);
     },
   });
 
+  const [canvas, setCanvas] = useState<UICanvas | null>(null);
+
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>();
 
   const [attachments, setAttachments] = useState<Array<Attachment>>([]);
 
   return (
+    <>
       <div className="flex flex-col min-w-0 h-dvh bg-background">
-      <ChatHeader selectedModelName={selectedModelName} />
+        <ChatHeader selectedModelId={selectedModelId} />
         <div
           ref={messagesContainerRef}
           className="flex flex-col min-w-0 gap-6 flex-1 overflow-y-scroll"
@@ -51,6 +66,8 @@ export function Chat({
               content={message.content}
               attachments={message.experimental_attachments}
               toolInvocations={message.toolInvocations}
+              canvas={canvas}
+              setCanvas={setCanvas}
             />
           ))}
 
@@ -69,9 +86,35 @@ export function Chat({
             attachments={attachments}
             setAttachments={setAttachments}
             messages={messages}
+            setMessages={setMessages}
             append={append}
           />
         </form>
       </div>
+
+      <AnimatePresence>
+        {canvas && canvas.isVisible && (
+          <Canvas
+            input={input}
+            setInput={setInput}
+            handleSubmit={handleSubmit}
+            isLoading={isLoading}
+            stop={stop}
+            attachments={attachments}
+            setAttachments={setAttachments}
+            append={append}
+            canvas={canvas}
+            setCanvas={setCanvas}
+            messages={messages}
+            setMessages={setMessages}
+          />
+        )}
+      </AnimatePresence>
+
+      <CanvasStreamHandler
+        streamingData={streamingData}
+        setCanvas={setCanvas}
+      />
+    </>
   );
 }

components/custom/diffview.tsx

@@ -0,0 +1,100 @@
+import OrderedMap from 'orderedmap';
+import {
+  Schema,
+  Node as ProsemirrorNode,
+  MarkSpec,
+  DOMParser,
+} from 'prosemirror-model';
+import { schema } from 'prosemirror-schema-basic';
+import { addListNodes } from 'prosemirror-schema-list';
+import { EditorState } from 'prosemirror-state';
+import { EditorView } from 'prosemirror-view';
+import React, { useEffect, useRef } from 'react';
+import { renderToString } from 'react-dom/server';
+import ReactMarkdown from 'react-markdown';
+
+import { diffEditor, DiffType } from '@/lib/editor/index';
+
+const diffSchema = new Schema({
+  nodes: addListNodes(schema.spec.nodes, 'paragraph block*', 'block'),
+  marks: OrderedMap.from({
+    ...schema.spec.marks.toObject(),
+    diffMark: {
+      attrs: { type: { default: '' } },
+      toDOM(mark) {
+        let className = '';
+
+        switch (mark.attrs.type) {
+          case DiffType.Inserted:
+            className =
+              'bg-green-100 text-green-700 dark:bg-green-500/70 dark:text-green-300';
+            break;
+          case DiffType.Deleted:
+            className =
+              'bg-red-100 line-through text-red-600 dark:bg-red-500/70 dark:text-red-300';
+            break;
+          default:
+            className = '';
+        }
+        return ['span', { class: className }, 0];
+      },
+    } as MarkSpec,
+  }),
+});
+
+function computeDiff(oldDoc: ProsemirrorNode, newDoc: ProsemirrorNode) {
+  return diffEditor(diffSchema, oldDoc.toJSON(), newDoc.toJSON());
+}
+
+type DiffEditorProps = {
+  oldContent: string;
+  newContent: string;
+};
+
+export const DiffView = ({ oldContent, newContent }: DiffEditorProps) => {
+  const editorRef = useRef<HTMLDivElement>(null);
+  const viewRef = useRef<EditorView | null>(null);
+
+  useEffect(() => {
+    if (editorRef.current && !viewRef.current) {
+      const parser = DOMParser.fromSchema(diffSchema);
+
+      const oldHtmlContent = renderToString(
+        <ReactMarkdown>{oldContent}</ReactMarkdown>
+      );
+      const newHtmlContent = renderToString(
+        <ReactMarkdown>{newContent}</ReactMarkdown>
+      );
+
+      const oldContainer = document.createElement('div');
+      oldContainer.innerHTML = oldHtmlContent;
+
+      const newContainer = document.createElement('div');
+      newContainer.innerHTML = newHtmlContent;
+
+      const oldDoc = parser.parse(oldContainer);
+      const newDoc = parser.parse(newContainer);
+
+      const diffedDoc = computeDiff(oldDoc, newDoc);
+
+      const state = EditorState.create({
+        doc: diffedDoc,
+        plugins: [],
+      });
+
+      viewRef.current = new EditorView(editorRef.current, {
+        state,
+        editable: () => false,
+      });
+    }
+
+    return () => {
+      if (viewRef.current) {
+        viewRef.current.destroy();
+        viewRef.current = null;
+      }
+    };
+  }, [oldContent, newContent]);
+
+  return <div className="diff-editor" ref={editorRef} />;
+};

components/custom/document-skeleton.tsx

@@ -0,0 +1,15 @@
+'use client';
+
+export const DocumentSkeleton = () => {
+  return (
+    <div className="flex flex-col gap-4 w-full">
+      <div className="animate-pulse rounded-lg h-12 bg-muted-foreground/20 w-1/2" />
+      <div className="animate-pulse rounded-lg h-5 bg-muted-foreground/20 w-full" />
+      <div className="animate-pulse rounded-lg h-5 bg-muted-foreground/20 w-full" />
+      <div className="animate-pulse rounded-lg h-5 bg-muted-foreground/20 w-1/3" />
+      <div className="animate-pulse rounded-lg h-5 bg-transparent w-52" />
+      <div className="animate-pulse rounded-lg h-8 bg-muted-foreground/20 w-52" />
+      <div className="animate-pulse rounded-lg h-5 bg-muted-foreground/20 w-2/3" />
+    </div>
+  );
+};

components/custom/document.tsx

@@ -0,0 +1,111 @@
+import { SetStateAction } from 'react';
+
+import { UICanvas } from './canvas';
+import { FileIcon, LoaderIcon, MessageIcon, PencilEditIcon } from './icons';
+
+const getActionText = (type: 'create' | 'update' | 'request-suggestions') => {
+  switch (type) {
+    case 'create':
+      return 'Creating';
+    case 'update':
+      return 'Updating';
+    case 'request-suggestions':
+      return 'Adding suggestions';
+    default:
+      return null;
+  }
+};
+
+interface DocumentToolResultProps {
+  type: 'create' | 'update' | 'request-suggestions';
+  result: any;
+  canvas: UICanvas | null;
+  setCanvas: (value: SetStateAction<UICanvas | null>) => void;
+}
+
+export function DocumentToolResult({
+  type,
+  result,
+  canvas,
+  setCanvas,
+}: DocumentToolResultProps) {
+  return (
+    <div
+      className="cursor-pointer border py-2 px-3 rounded-xl w-fit flex flex-row gap-3 items-start"
+      onClick={(event) => {
+        const rect = event.currentTarget.getBoundingClientRect();
+
+        const boundingBox = {
+          top: rect.top,
+          left: rect.left,
+          width: rect.width,
+          height: rect.height,
+        };
+
+        if (!canvas) {
+          setCanvas({
+            documentId: result.id,
+            content: '',
+            title: result.title,
+            isVisible: true,
+            status: 'idle',
+            boundingBox,
+          });
+        } else {
+          if (canvas.documentId !== result.id) {
+            setCanvas({
+              documentId: result.id,
+              content: '',
+              title: result.title,
+              isVisible: true,
+              status: 'idle',
+              boundingBox,
+            });
+          }
+        }
+      }}
+    >
+      <div className="text-muted-foreground mt-1">
+        {type === 'create' ? (
+          <FileIcon />
+        ) : type === 'update' ? (
+          <PencilEditIcon />
+        ) : type === 'request-suggestions' ? (
+          <MessageIcon />
+        ) : null}
+      </div>
+      <div className="">
+        {getActionText(type)} {result.title}
+      </div>
+    </div>
+  );
+}
+
+interface DocumentToolCallProps {
+  type: 'create' | 'update' | 'request-suggestions';
+  args: any;
+}
+
+export function DocumentToolCall({ type, args }: DocumentToolCallProps) {
+  return (
+    <div className="w-fit border py-2 px-3 rounded-xl flex flex-row items-start justify-between gap-3">
+      <div className="flex flex-row gap-3 items-start">
+        <div className="text-zinc-500 mt-1">
+          {type === 'create' ? (
+            <FileIcon />
+          ) : type === 'update' ? (
+            <PencilEditIcon />
+          ) : type === 'request-suggestions' ? (
+            <MessageIcon />
+          ) : null}
+        </div>
+
+        <div className="">
+          {getActionText(type)} {args.title}
+        </div>
+      </div>
+
+      <div className="animate-spin mt-1">{<LoaderIcon />}</div>
+    </div>
+  );
+}

components/custom/editor.tsx

@@ -0,0 +1,193 @@
+'use client';
+
+import { exampleSetup } from 'prosemirror-example-setup';
+import { inputRules, textblockTypeInputRule } from 'prosemirror-inputrules';
+import { defaultMarkdownSerializer } from 'prosemirror-markdown';
+import { Schema, DOMParser } from 'prosemirror-model';
+import { schema } from 'prosemirror-schema-basic';
+import { addListNodes } from 'prosemirror-schema-list';
+import { EditorState } from 'prosemirror-state';
+import { Decoration, DecorationSet, EditorView } from 'prosemirror-view';
+import React, { memo, useEffect, useMemo, useRef, useState } from 'react';
+import { renderToString } from 'react-dom/server';
+
+import { Suggestion } from '@/db/schema';
+import {
+  createSuggestionWidget,
+  projectWithHighlights,
+  suggestionsPlugin,
+  suggestionsPluginKey,
+  UISuggestion,
+} from '@/lib/editor/suggestions';
+
+import { Markdown } from './markdown';
+
+const mySchema = new Schema({
+  nodes: addListNodes(schema.spec.nodes, 'paragraph block*', 'block'),
+  // @ts-expect-error: TODO need to fix type mismatch
+  marks: {
+    ...schema.spec.marks,
+  },
+});
+
+function headingRule(level: number) {
+  return textblockTypeInputRule(
+    new RegExp(`^(#{1,${level}})\\s$`),
+    mySchema.nodes.heading,
+    () => ({ level })
+  );
+}
+
+interface WidgetRoot {
+  destroy: () => void;
+}
+
+type EditorProps = {
+  content: string;
+  onChange: (updatedContent: string) => void;
+  status: 'streaming' | 'idle';
+  currentVersionIndex: number;
+  suggestions: Array<Suggestion>;
+};
+
+function PureEditor({
+  content,
+  onChange,
+  suggestions: suggestionsWithoutHighlights,
+}: EditorProps) {
+  const editorRef = useRef<HTMLDivElement>(null);
+  const viewRef = useRef<EditorView | null>(null);
+  const widgetRootsRef = useRef<Map<string, WidgetRoot>>(new Map());
+
+  useEffect(() => {
+    if (editorRef.current && !viewRef.current) {
+      if (mySchema) {
+        const parser = DOMParser.fromSchema(mySchema);
+
+        const htmlContent = renderToString(<Markdown>{content}</Markdown>);
+
+        const container = document.createElement('div');
+        container.innerHTML = htmlContent;
+
+        const state = EditorState.create({
+          doc: parser.parse(container),
+          plugins: [
+            ...exampleSetup({ schema: mySchema, menuBar: false }),
+            inputRules({
+              rules: [
+                headingRule(1),
+                headingRule(2),
+                headingRule(3),
+                headingRule(4),
+                headingRule(5),
+                headingRule(6),
+              ],
+            }),
+            suggestionsPlugin,
+          ],
+        });
+
+        viewRef.current = new EditorView(editorRef.current, {
+          state,
+          dispatchTransaction: (transaction) => {
+            const newState = viewRef.current!.state.apply(transaction);
+            viewRef.current!.updateState(newState);
+
+            if (transaction.docChanged) {
+              const content = defaultMarkdownSerializer.serialize(newState.doc);
+              onChange(content);
+            }
+          },
+        });
+      } else {
+        console.error('Schema is not properly initialized');
+      }
+    }
+
+    return () => {
+      if (viewRef.current) {
+        viewRef.current.destroy();
+        viewRef.current = null;
+      }
+    };
+  }, [content, onChange]);
+
+  useEffect(() => {
+    if (viewRef.current && viewRef.current.state.doc && content) {
+      const computedSuggestions = projectWithHighlights(
+        viewRef.current.state.doc,
+        suggestionsWithoutHighlights
+      ).filter(
+        (suggestion) =>
+          suggestion.selectionStart !== 0 && suggestion.selectionEnd !== 0
+      );
+
+      const decorations = createDecorations(
+        computedSuggestions,
+        viewRef.current
+      );
+      const transaction = viewRef.current.state.tr;
+      transaction.setMeta(suggestionsPluginKey, { decorations });
+      viewRef.current.dispatch(transaction);
+    }
+  }, [suggestionsWithoutHighlights, content]);
+
+  const createDecorations = (suggestions: UISuggestion[], view: EditorView) => {
+    const decorations: Decoration[] = [];
+
+    suggestions.forEach((suggestion) => {
+      decorations.push(
+        Decoration.inline(suggestion.selectionStart, suggestion.selectionEnd, {
+          class:
+            'suggestion-highlight bg-yellow-100 hover:bg-yellow-200 dark:hover:bg-yellow-400/50 dark:text-yellow-50 dark:bg-yellow-400/40',
+        })
+      );
+
+      decorations.push(
+        Decoration.widget(suggestion.selectionStart, (view) => {
+          const { dom, destroy } = createSuggestionWidget(suggestion, view);
+          const key = `widget-${suggestion.id}`;
+          widgetRootsRef.current.set(key, { destroy });
+          return dom;
+        })
+      );
+    });
+
+    return DecorationSet.create(view.state.doc, decorations);
+  };
+
+  useEffect(() => {
+    return () => {
+      widgetRootsRef.current.forEach((root) => {
+        root.destroy();
+      });
+
+      widgetRootsRef.current.clear();
+
+      if (viewRef.current) {
+        viewRef.current.destroy();
+        viewRef.current = null;
+      }
+    };
+  }, []);
+
+  return <div className="relative prose dark:prose-invert" ref={editorRef} />;
+}
+
+function areEqual(prevProps: EditorProps, nextProps: EditorProps) {
+  if (prevProps.suggestions !== nextProps.suggestions) {
+    return false;
+  } else if (prevProps.content === '' && nextProps.content !== '') {
+    return false;
+  } else if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex) {
+    return false;
+  } else if (prevProps.onChange !== nextProps.onChange) {
+    return false;
+  } else if (prevProps.status === 'idle') {
+    return true;
+  } else {
+    return prevProps.content === nextProps.content;
+  }
+}
+
+export const Editor = memo(PureEditor, areEqual);

components/custom/icons.tsx

@@ -609,3 +609,105 @@ export const MessageIcon = ({ size = 16 }: { size?: number }) => {
     </svg>
   );
 };
+
+export const CrossIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M12.4697 13.5303L13 14.0607L14.0607 13L13.5303 12.4697L9.06065 7.99999L13.5303 3.53032L14.0607 2.99999L13 1.93933L12.4697 2.46966L7.99999 6.93933L3.53032 2.46966L2.99999 1.93933L1.93933 2.99999L2.46966 3.53032L6.93933 7.99999L2.46966 12.4697L1.93933 13L2.99999 14.0607L3.53032 13.5303L7.99999 9.06065L12.4697 13.5303Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const UndoIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M13.5 8C13.5 4.96643 11.0257 2.5 7.96452 2.5C5.42843 2.5 3.29365 4.19393 2.63724 6.5H5.25H6V8H5.25H0.75C0.335787 8 0 7.66421 0 7.25V2.75V2H1.5V2.75V5.23347C2.57851 2.74164 5.06835 1 7.96452 1C11.8461 1 15 4.13001 15 8C15 11.87 11.8461 15 7.96452 15C5.62368 15 3.54872 13.8617 2.27046 12.1122L1.828 11.5066L3.03915 10.6217L3.48161 11.2273C4.48831 12.6051 6.12055 13.5 7.96452 13.5C11.0257 13.5 13.5 11.0336 13.5 8Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const RedoIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M2.5 8C2.5 4.96643 4.97431 2.5 8.03548 2.5C10.5716 2.5 12.7064 4.19393 13.3628 6.5H10.75H10V8H10.75H15.25C15.6642 8 16 7.66421 16 7.25V2.75V2H14.5V2.75V5.23347C13.4215 2.74164 10.9316 1 8.03548 1C4.1539 1 1 4.13001 1 8C1 11.87 4.1539 15 8.03548 15C10.3763 15 12.4513 13.8617 13.7295 12.1122L14.172 11.5066L12.9609 10.6217L12.5184 11.2273C11.5117 12.6051 9.87945 13.5 8.03548 13.5C4.97431 13.5 2.5 11.0336 2.5 8Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const DeltaIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M2.67705 15H1L1.75 13.5L6.16147 4.67705L6.15836 4.67082L6.16667 4.66667L7.16147 2.67705L8 1L8.83853 2.67705L14.25 13.5L15 15H13.3229H2.67705ZM7 6.3541L10.5729 13.5H3.42705L7 6.3541Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const PenIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M8.75 0.189331L9.28033 0.719661L15.2803 6.71966L15.8107 7.24999L15.2803 7.78032L13.7374 9.32322C13.1911 9.8696 12.3733 9.97916 11.718 9.65188L9.54863 13.5568C8.71088 15.0648 7.12143 16 5.39639 16H0.75H0V15.25V10.6036C0 8.87856 0.935237 7.28911 2.4432 6.45136L6.34811 4.28196C6.02084 3.62674 6.13039 2.80894 6.67678 2.26255L8.21967 0.719661L8.75 0.189331ZM7.3697 5.43035L10.5696 8.63029L8.2374 12.8283C7.6642 13.8601 6.57668 14.5 5.39639 14.5H2.56066L5.53033 11.5303L4.46967 10.4697L1.5 13.4393V10.6036C1.5 9.42331 2.1399 8.33579 3.17166 7.76259L7.3697 5.43035ZM12.6768 8.26256C12.5791 8.36019 12.4209 8.36019 12.3232 8.26255L12.0303 7.96966L8.03033 3.96966L7.73744 3.67677C7.63981 3.57914 7.63981 3.42085 7.73744 3.32321L8.75 2.31065L13.6893 7.24999L12.6768 8.26256Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const SummarizeIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: "currentcolor" }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M1.75 12H1V10.5H1.75H5.25H6V12H5.25H1.75ZM1.75 7.75H1V6.25H1.75H4.25H5V7.75H4.25H1.75ZM1.75 3.5H1V2H1.75H7.25H8V3.5H7.25H1.75ZM12.5303 14.7803C12.2374 15.0732 11.7626 15.0732 11.4697 14.7803L9.21967 12.5303L8.68934 12L9.75 10.9393L10.2803 11.4697L11.25 12.4393V2.75V2H12.75V2.75V12.4393L13.7197 11.4697L14.25 10.9393L15.3107 12L14.7803 12.5303L12.5303 14.7803Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);

components/custom/markdown.tsx

@@ -63,6 +63,48 @@ const NonMemoizedMarkdown = ({ children }: { children: string }) => {
         </Link>
       );
     },
+    h1: ({ node, children, ...props }: any) => {
+      return (
+        <h1 className="text-3xl font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h1>
+      );
+    },
+    h2: ({ node, children, ...props }: any) => {
+      return (
+        <h2 className="text-2xl font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h2>
+      );
+    },
+    h3: ({ node, children, ...props }: any) => {
+      return (
+        <h3 className="text-xl font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h3>
+      );
+    },
+    h4: ({ node, children, ...props }: any) => {
+      return (
+        <h4 className="text-lg font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h4>
+      );
+    },
+    h5: ({ node, children, ...props }: any) => {
+      return (
+        <h5 className="text-base font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h5>
+      );
+    },
+    h6: ({ node, children, ...props }: any) => {
+      return (
+        <h6 className="text-sm font-semibold mt-6 mb-2" {...props}>
+          {children}
+        </h6>
+      );
+    },
   };
 
   return (

components/custom/message.tsx

@@ -1,10 +1,13 @@
 'use client';
 
 import { Attachment, ToolInvocation } from 'ai';
+import cx from 'classnames';
 import { motion } from 'framer-motion';
 import { Sparkles } from 'lucide-react';
-import { ReactNode } from 'react';
+import { Dispatch, ReactNode, SetStateAction } from 'react';
 
+import { UICanvas } from './canvas';
+import { DocumentToolCall, DocumentToolResult } from './document';
 import { Markdown } from './markdown';
 import { PreviewAttachment } from './preview-attachment';
 import { Weather } from './weather';
@@ -14,11 +17,15 @@ export const Message = ({
   content,
   toolInvocations,
   attachments,
+  canvas,
+  setCanvas,
 }: {
   role: string;
   content: string | ReactNode;
   toolInvocations: Array<ToolInvocation> | undefined;
   attachments?: Array<Attachment>;
+  canvas: UICanvas | null;
+  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
 }) => {
   return (
     <motion.div
@@ -27,7 +34,16 @@ export const Message = ({
       animate={{ y: 0, opacity: 1 }}
       data-role={role}
     >
-      <div className="flex gap-4 group-data-[role=user]/message:px-5 w-full group-data-[role=user]/message:w-fit group-data-[role=user]/message:ml-auto group-data-[role=user]/message:max-w-2xl group-data-[role=user]/message:py-3.5 group-data-[role=user]/message:bg-muted rounded-xl">
+      <div
+        className={cx(
+          'flex gap-4 group-data-[role=user]/message:px-5 w-full group-data-[role=user]/message:w-fit group-data-[role=user]/message:ml-auto group-data-[role=user]/message:max-w-2xl group-data-[role=user]/message:py-3.5 rounded-xl',
+          {
+            'group-data-[role=user]/message:bg-muted': !canvas,
+            'group-data-[role=user]/message:bg-zinc-300 dark:group-data-[role=user]/message:bg-zinc-800':
+              canvas,
+          }
+        )}
+      >
         {role === 'assistant' && (
           <div className="size-8 flex items-center rounded-full justify-center ring-1 shrink-0 ring-border">
             <Sparkles className="size-4" />
@@ -40,10 +56,10 @@ export const Message = ({
             </div>
           )}
 
-          {toolInvocations && toolInvocations.length > 0 ? (
+          {toolInvocations && toolInvocations.length > 0 && (
             <div className="flex flex-col gap-4">
               {toolInvocations.map((toolInvocation) => {
-                const { toolName, toolCallId, state } = toolInvocation;
+                const { toolName, toolCallId, state, args } = toolInvocation;
 
                 if (state === 'result') {
                   const { result } = toolInvocation;
@@ -52,19 +68,58 @@ export const Message = ({
                     <div key={toolCallId}>
                       {toolName === 'getWeather' ? (
                         <Weather weatherAtLocation={result} />
-                      ) : null}
+                      ) : toolName === 'createDocument' ? (
+                        <DocumentToolResult
+                          type="create"
+                          result={result}
+                          canvas={canvas}
+                          setCanvas={setCanvas}
+                        />
+                      ) : toolName === 'updateDocument' ? (
+                        <DocumentToolResult
+                          type="update"
+                          result={result}
+                          canvas={canvas}
+                          setCanvas={setCanvas}
+                        />
+                      ) : toolName === 'requestSuggestions' ? (
+                        <DocumentToolResult
+                          type="request-suggestions"
+                          result={result}
+                          canvas={canvas}
+                          setCanvas={setCanvas}
+                        />
+                      ) : (
+                        <pre>{JSON.stringify(result, null, 2)}</pre>
+                      )}
                     </div>
                   );
                 } else {
                   return (
-                    <div key={toolCallId} className="skeleton">
+                    <div
-                      {toolName === 'getWeather' ? <Weather /> : null}
+                      key={toolCallId}
+                      className={cx({
+                        skeleton: ['getWeather'].includes(toolName),
+                      })}
+                    >
+                      {toolName === 'getWeather' ? (
+                        <Weather />
+                      ) : toolName === 'createDocument' ? (
+                        <DocumentToolCall type="create" args={args} />
+                      ) : toolName === 'updateDocument' ? (
+                        <DocumentToolCall type="update" args={args} />
+                      ) : toolName === 'requestSuggestions' ? (
+                        <DocumentToolCall
+                          type="request-suggestions"
+                          args={args}
+                        />
+                      ) : null}
                     </div>
                   );
                 }
               })}
             </div>
-          ) : null}
+          )}
 
           {attachments && (
             <div className="flex flex-row gap-2">

components/custom/model-selector.tsx

@@ -3,7 +3,8 @@
 import { Check, ChevronDown } from 'lucide-react';
 import { startTransition, useMemo, useOptimistic, useState } from 'react';
 
-import { saveModel } from '@/app/(chat)/actions';
+import { models } from '@/ai/models';
+import { saveModelId } from '@/app/(chat)/actions';
 import { Button } from '@/components/ui/button';
 import {
   DropdownMenu,
@@ -11,22 +12,21 @@ import {
   DropdownMenuItem,
   DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu';
-import { type Model, models } from '@/lib/model';
 import { cn } from '@/lib/utils';
 
 export function ModelSelector({
-  selectedModelName,
+  selectedModelId,
   className,
 }: {
-  selectedModelName: Model['name'];
+  selectedModelId: string;
 } & React.ComponentProps<typeof Button>) {
   const [open, setOpen] = useState(false);
-  const [optimisticModelName, setOptimisticModelName] =
+  const [optimisticModelId, setOptimisticModelId] =
-    useOptimistic(selectedModelName);
+    useOptimistic(selectedModelId);
 
   const selectModel = useMemo(
-    () => models.find((model) => model.name === optimisticModelName),
+    () => models.find((model) => model.id === optimisticModelId),
-    [optimisticModelName]
+    [optimisticModelId]
   );
 
   return (
@@ -46,17 +46,17 @@ export function ModelSelector({
       <DropdownMenuContent align="start" className="min-w-[300px]">
         {models.map((model) => (
           <DropdownMenuItem
-            key={model.name}
+            key={model.id}
             onSelect={() => {
               setOpen(false);
 
               startTransition(() => {
-                setOptimisticModelName(model.name);
+                setOptimisticModelId(model.id);
-                saveModel(model.name);
+                saveModelId(model.id);
               });
             }}
             className="gap-4 group/item"
-            data-active={model.name === optimisticModelName}
+            data-active={model.id === optimisticModelId}
           >
             <div className="flex flex-col gap-1 items-start">
               {model.label}

components/custom/multimodal-input.tsx

@@ -1,6 +1,7 @@
 'use client';
 
 import { Attachment, ChatRequestOptions, CreateMessage, Message } from 'ai';
+import cx from 'classnames';
 import { motion } from 'framer-motion';
 import React, {
   useRef,
@@ -13,6 +14,8 @@ import React, {
 } from 'react';
 import { toast } from 'sonner';
 
+import { sanitizeUIMessages } from '@/lib/utils';
+
 import { ArrowUpIcon, PaperclipIcon, StopIcon } from './icons';
 import { PreviewAttachment } from './preview-attachment';
 import useWindowSize from './use-window-size';
@@ -26,9 +29,9 @@ const suggestedActions = [
     action: 'What is the weather in San Francisco?',
   },
   {
-    title: "Answer like I'm 5,",
+    title: 'Help me draft an essay',
-    label: 'why is the sky blue?',
+    label: 'about Silicon Valley',
-    action: "Answer like I'm 5, why is the sky blue?",
+    action: 'Help me draft an essay about Silicon Valley',
   },
 ];
 
@@ -40,8 +43,10 @@ export function MultimodalInput({
   attachments,
   setAttachments,
   messages,
+  setMessages,
   append,
   handleSubmit,
+  className,
 }: {
   input: string;
   setInput: (value: string) => void;
@@ -50,6 +55,7 @@ export function MultimodalInput({
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
   messages: Array<Message>;
+  setMessages: Dispatch<SetStateAction<Array<Message>>>;
   append: (
     message: Message | CreateMessage,
     chatRequestOptions?: ChatRequestOptions
@@ -60,6 +66,7 @@ export function MultimodalInput({
     },
     chatRequestOptions?: ChatRequestOptions
   ) => void;
+  className?: string;
 }) {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { width } = useWindowSize();
@@ -220,7 +227,10 @@ export function MultimodalInput({
         placeholder="Send a message..."
         value={input}
         onChange={handleInput}
-        className="min-h-[24px] overflow-hidden resize-none rounded-xl p-4 focus-visible:ring-0 focus-visible:ring-offset-0 text-base bg-muted border-none"
+        className={cx(
+          'min-h-[24px] overflow-hidden resize-none rounded-lg text-base bg-muted',
+          className
+        )}
         rows={3}
         onKeyDown={(event) => {
           if (event.key === 'Enter' && !event.shiftKey) {
@@ -241,6 +251,7 @@ export function MultimodalInput({
           onClick={(event) => {
             event.preventDefault();
             stop();
+            setMessages((messages) => sanitizeUIMessages(messages));
           }}
         >
           <StopIcon size={14} />

components/custom/suggestion.tsx

@@ -0,0 +1,59 @@
+'use client';
+
+import { motion } from 'framer-motion';
+import { useState } from 'react';
+
+import { UISuggestion } from '@/lib/editor/suggestions';
+
+import { CrossIcon, MessageIcon } from './icons';
+import { Button } from '../ui/button';
+
+export const Suggestion = ({
+  suggestion,
+  onApply,
+}: {
+  suggestion: UISuggestion;
+  onApply: () => void;
+}) => {
+  const [isExpanded, setIsExpanded] = useState(false);
+
+  return !isExpanded ? (
+    <div
+      className="absolute cursor-pointer text-muted-foreground mt-1 -right-8"
+      onClick={() => {
+        setIsExpanded(true);
+      }}
+    >
+      <MessageIcon size={14} />
+    </div>
+  ) : (
+    <motion.div
+      className="absolute bg-background p-3 flex flex-col gap-3 rounded-2xl border text-sm w-56 shadow-xl z-50 -right-12 md:-right-20"
+      transition={{ type: 'spring', stiffness: 500, damping: 30 }}
+      whileHover={{ scale: 1.05 }}
+    >
+      <div className="flex flex-row items-center justify-between">
+        <div className="flex flex-row items-center gap-2">
+          <div className="size-4 bg-muted-foreground/25 rounded-full" />
+          <div className="font-medium">Assistant</div>
+        </div>
+        <div
+          className="text-xs text-gray-500 cursor-pointer"
+          onClick={() => {
+            setIsExpanded(false);
+          }}
+        >
+          <CrossIcon size={12} />
+        </div>
+      </div>
+      <div>{suggestion.description}</div>
+      <Button
+        variant="outline"
+        className="w-fit py-1.5 px-3 rounded-full"
+        onClick={onApply}
+      >
+        Apply
+      </Button>
+    </motion.div>
+  );
+};

components/custom/toolbar.tsx

@@ -0,0 +1,413 @@
+'use client';
+
+import { ChatRequestOptions, CreateMessage, Message } from 'ai';
+import cx from 'classnames';
+import {
+  AnimatePresence,
+  motion,
+  useMotionValue,
+  useTransform,
+} from 'framer-motion';
+import { Dispatch, SetStateAction, useEffect, useRef, useState } from 'react';
+
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from '@/components/ui/tooltip';
+import { sanitizeUIMessages } from '@/lib/utils';
+
+import {
+  ArrowUpIcon,
+  MessageIcon,
+  PenIcon,
+  StopIcon,
+  SummarizeIcon,
+} from './icons';
+import useClickOutside from './use-click-outside';
+
+type ToolProps = {
+  type: 'final-polish' | 'request-suggestions' | 'adjust-reading-level';
+  description: string;
+  icon: JSX.Element;
+  selectedTool: string | null;
+  setSelectedTool: Dispatch<SetStateAction<string | null>>;
+  isToolbarVisible?: boolean;
+  setIsToolbarVisible?: Dispatch<SetStateAction<boolean>>;
+  isAnimating: boolean;
+  append: (
+    message: Message | CreateMessage,
+    chatRequestOptions?: ChatRequestOptions
+  ) => Promise<string | null | undefined>;
+};
+
+const Tool = ({
+  type,
+  description,
+  icon,
+  selectedTool,
+  setSelectedTool,
+  isToolbarVisible,
+  setIsToolbarVisible,
+  isAnimating,
+  append,
+}: ToolProps) => {
+  const [isHovered, setIsHovered] = useState(false);
+
+  useEffect(() => {
+    if (selectedTool !== type) {
+      setIsHovered(false);
+    }
+  }, [selectedTool, type]);
+
+  return (
+    <Tooltip open={isHovered && !isAnimating}>
+      <TooltipTrigger>
+        <motion.div
+          className={cx('p-3 rounded-full', {
+            'bg-foreground !text-background': selectedTool === type,
+          })}
+          onHoverStart={() => {
+            setIsHovered(true);
+          }}
+          onHoverEnd={() => {
+            if (selectedTool !== type) setIsHovered(false);
+          }}
+          initial={{ scale: 1, opacity: 0 }}
+          animate={{ opacity: 1, transition: { delay: 0.1 } }}
+          whileHover={{ scale: 1.1 }}
+          whileTap={{ scale: 0.95 }}
+          exit={{
+            scale: 0.9,
+            opacity: 0,
+            transition: { duration: 0.1 },
+          }}
+          onClick={() => {
+            if (!isToolbarVisible && setIsToolbarVisible) {
+              setIsToolbarVisible(true);
+              return;
+            }
+
+            if (!selectedTool) {
+              setIsHovered(true);
+              setSelectedTool(type);
+              return;
+            }
+
+            if (selectedTool !== type) {
+              setSelectedTool(type);
+            } else {
+              if (type === 'final-polish') {
+                append({
+                  role: 'user',
+                  content:
+                    'Please add final polish and check for grammar, add section titles for better structure, and ensure everything reads smoothly.',
+                });
+
+                setSelectedTool(null);
+              } else if (type === 'request-suggestions') {
+                append({
+                  role: 'user',
+                  content:
+                    'Please add suggestions you have that could improve the writing.',
+                });
+
+                setSelectedTool(null);
+              }
+            }
+          }}
+        >
+          {selectedTool === type ? <ArrowUpIcon /> : icon}
+        </motion.div>
+      </TooltipTrigger>
+      <TooltipContent
+        side="left"
+        sideOffset={16}
+        className="bg-foreground text-background rounded-2xl p-3 px-4"
+      >
+        {description}
+      </TooltipContent>
+    </Tooltip>
+  );
+};
+
+const ReadingLevelSelector = ({
+  setSelectedTool,
+  append,
+  isAnimating,
+}: {
+  setSelectedTool: Dispatch<SetStateAction<string | null>>;
+  isAnimating: boolean;
+  append: (
+    message: Message | CreateMessage,
+    chatRequestOptions?: ChatRequestOptions
+  ) => Promise<string | null | undefined>;
+}) => {
+  const LEVELS = [
+    'Elementary',
+    'Middle School',
+    'Keep current level',
+    'High School',
+    'College',
+    'Graduate',
+  ];
+
+  const y = useMotionValue(-40 * 2);
+  const dragConstraints = 5 * 40 + 2;
+  const yToLevel = useTransform(y, [0, -dragConstraints], [0, 5]);
+
+  const [currentLevel, setCurrentLevel] = useState(2);
+  const [hasUserSelectedLevel, setHasUserSelectedLevel] =
+    useState<boolean>(false);
+
+  useEffect(() => {
+    const unsubscribe = yToLevel.on('change', (latest) => {
+      const level = Math.min(5, Math.max(0, Math.round(Math.abs(latest))));
+      setCurrentLevel(level);
+    });
+
+    return () => unsubscribe();
+  }, [yToLevel]);
+
+  return (
+    <div className="relative flex flex-col justify-end items-center">
+      {[...Array(6)].map((_, index) => (
+        <motion.div
+          key={`dot-${index}`}
+          className="size-[40px] flex flex-row items-center justify-center"
+          initial={{ opacity: 0 }}
+          animate={{ opacity: 1 }}
+          exit={{ opacity: 0 }}
+          transition={{ delay: 0.1 }}
+        >
+          <div className="size-2 rounded-full bg-muted-foreground/40" />
+        </motion.div>
+      ))}
+
+      <TooltipProvider>
+        <Tooltip open={!isAnimating}>
+          <TooltipTrigger asChild>
+            <motion.div
+              className={cx(
+                'absolute bg-background p-3 border rounded-full flex flex-row items-center',
+                {
+                  'bg-foreground text-background': currentLevel !== 2,
+                  'bg-background text-foreground': currentLevel === 2,
+                }
+              )}
+              style={{ y }}
+              drag="y"
+              dragElastic={0}
+              dragMomentum={false}
+              whileHover={{ scale: 1.05 }}
+              whileTap={{ scale: 0.95 }}
+              transition={{ duration: 0.1 }}
+              dragConstraints={{ top: -dragConstraints, bottom: 0 }}
+              onDragStart={() => {
+                setHasUserSelectedLevel(false);
+              }}
+              onDragEnd={() => {
+                if (currentLevel === 2) {
+                  setSelectedTool(null);
+                } else {
+                  setHasUserSelectedLevel(true);
+                }
+              }}
+              onClick={() => {
+                if (currentLevel !== 2 && hasUserSelectedLevel) {
+                  append({
+                    role: 'user',
+                    content: `Please adjust the reading level to ${LEVELS[currentLevel]} level.`,
+                  });
+
+                  setSelectedTool(null);
+                }
+              }}
+            >
+              {currentLevel === 2 ? <SummarizeIcon /> : <ArrowUpIcon />}
+            </motion.div>
+          </TooltipTrigger>
+          <TooltipContent
+            side="left"
+            sideOffset={16}
+            className="bg-foreground text-background text-sm rounded-2xl p-3 px-4"
+          >
+            {LEVELS[currentLevel]}
+          </TooltipContent>
+        </Tooltip>
+      </TooltipProvider>
+    </div>
+  );
+};
+
+export const Toolbar = ({
+  isToolbarVisible,
+  setIsToolbarVisible,
+  append,
+  isLoading,
+  stop,
+  setMessages,
+}: {
+  isToolbarVisible: boolean;
+  setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
+  isLoading: boolean;
+  append: (
+    message: Message | CreateMessage,
+    chatRequestOptions?: ChatRequestOptions
+  ) => Promise<string | null | undefined>;
+  stop: () => void;
+  setMessages: Dispatch<SetStateAction<Message[]>>;
+}) => {
+  const toolbarRef = useRef<HTMLDivElement>(null);
+  const timeoutRef = useRef<NodeJS.Timeout>();
+
+  const [selectedTool, setSelectedTool] = useState<string | null>(null);
+  const [isAnimating, setIsAnimating] = useState(false);
+
+  useClickOutside(toolbarRef, () => {
+    setIsToolbarVisible(false);
+    setSelectedTool(null);
+  });
+
+  const Tools = (
+    <>
+      <AnimatePresence>
+        {isToolbarVisible && (
+          <>
+            <Tool
+              type="adjust-reading-level"
+              description="Adjust reading level"
+              icon={<SummarizeIcon />}
+              selectedTool={selectedTool}
+              setSelectedTool={setSelectedTool}
+              append={append}
+              isAnimating={isAnimating}
+            />
+
+            <Tool
+              type="request-suggestions"
+              description="Request suggestions"
+              icon={<MessageIcon />}
+              selectedTool={selectedTool}
+              setSelectedTool={setSelectedTool}
+              append={append}
+              isAnimating={isAnimating}
+            />
+          </>
+        )}
+      </AnimatePresence>
+
+      <Tool
+        type="final-polish"
+        description="Add final polish"
+        icon={<PenIcon />}
+        selectedTool={selectedTool}
+        setSelectedTool={setSelectedTool}
+        isToolbarVisible={isToolbarVisible}
+        setIsToolbarVisible={setIsToolbarVisible}
+        append={append}
+        isAnimating={isAnimating}
+      />
+    </>
+  );
+
+  const startCloseTimer = () => {
+    if (timeoutRef.current) {
+      clearTimeout(timeoutRef.current);
+    }
+
+    timeoutRef.current = setTimeout(() => {
+      setSelectedTool(null);
+      setIsToolbarVisible(false);
+    }, 2000);
+  };
+
+  const cancelCloseTimer = () => {
+    if (timeoutRef.current) {
+      clearTimeout(timeoutRef.current);
+    }
+  };
+
+  useEffect(() => {
+    return () => {
+      if (timeoutRef.current) {
+        clearTimeout(timeoutRef.current);
+      }
+    };
+  }, []);
+
+  useEffect(() => {
+    if (isLoading) {
+      setIsToolbarVisible(false);
+    }
+  }, [isLoading, setIsToolbarVisible]);
+
+  return (
+    <TooltipProvider delayDuration={0}>
+      <motion.div
+        className="cursor-pointer fixed right-6 bottom-6 p-1.5 border rounded-full shadow-lg bg-background flex flex-col justify-end"
+        initial={{ opacity: 0, y: -20, scale: 1 }}
+        animate={
+          isToolbarVisible
+            ? selectedTool === 'adjust-reading-level'
+              ? {
+                  opacity: 1,
+                  y: 0,
+                  height: 6 * 43,
+                  transition: { delay: 0 },
+                  scale: 0.95,
+                }
+              : {
+                  opacity: 1,
+                  y: 0,
+                  height: 3 * 45,
+                  transition: { delay: 0 },
+                  scale: 1,
+                }
+            : { opacity: 1, y: 0, height: 54, transition: { delay: 0 } }
+        }
+        exit={{ opacity: 0, y: -20, transition: { duration: 0.1 } }}
+        transition={{ type: 'spring', stiffness: 300, damping: 25 }}
+        onHoverStart={() => {
+          if (isLoading) return;
+
+          cancelCloseTimer();
+          setIsToolbarVisible(true);
+        }}
+        onHoverEnd={() => {
+          if (isLoading) return;
+
+          startCloseTimer();
+        }}
+        onAnimationStart={() => {
+          setIsAnimating(true);
+        }}
+        onAnimationComplete={() => {
+          setIsAnimating(false);
+        }}
+        ref={toolbarRef}
+      >
+        {isLoading ? (
+          <div
+            className="p-3"
+            onClick={() => {
+              stop();
+              setMessages((messages) => sanitizeUIMessages(messages));
+            }}
+          >
+            <StopIcon />
+          </div>
+        ) : selectedTool === 'adjust-reading-level' ? (
+          <ReadingLevelSelector
+            append={append}
+            setSelectedTool={setSelectedTool}
+            isAnimating={isAnimating}
+          />
+        ) : (
+          Tools
+        )}
+      </motion.div>
+    </TooltipProvider>
+  );
+};

components/custom/use-canvas-stream.tsx

@@ -0,0 +1,103 @@
+import { JSONValue } from 'ai';
+import { Dispatch, SetStateAction, useEffect, useState } from 'react';
+import { useSWRConfig } from 'swr';
+
+import { Suggestion } from '@/db/schema';
+
+import { UICanvas } from './canvas';
+
+type StreamingDelta = {
+  type: 'text-delta' | 'title' | 'id' | 'suggestion' | 'clear' | 'finish';
+  content: string | Suggestion;
+};
+
+export function useCanvasStream({
+  streamingData,
+  setCanvas,
+}: {
+  streamingData: JSONValue[] | undefined;
+  setCanvas: Dispatch<SetStateAction<UICanvas | null>>;
+}) {
+  const { mutate } = useSWRConfig();
+  const [optimisticSuggestions, setOptimisticSuggestions] = useState<
+    Array<Suggestion>
+  >([]);
+
+  useEffect(() => {
+    if (optimisticSuggestions && optimisticSuggestions.length > 0) {
+      const [optimisticSuggestion] = optimisticSuggestions;
+      const url = `/api/suggestions?documentId=${optimisticSuggestion.documentId}`;
+      mutate(url, optimisticSuggestions, false);
+    }
+  }, [optimisticSuggestions, mutate]);
+
+  useEffect(() => {
+    const mostRecentDelta = streamingData?.at(-1);
+    if (!mostRecentDelta) return;
+
+    const delta = mostRecentDelta as StreamingDelta;
+
+    setCanvas((draftCanvas) => {
+      if (!draftCanvas) {
+        return {
+          content: '',
+          title: '',
+          isVisible: false,
+          documentId: delta.type === 'id' ? (delta.content as string) : '',
+          status: 'idle',
+          boundingBox: {
+            top: 0,
+            left: 0,
+            width: 0,
+            height: 0,
+          },
+        };
+      }
+
+      switch (delta.type) {
+        case 'text-delta':
+          return {
+            ...draftCanvas,
+            content: draftCanvas.content + (delta.content as string),
+            isVisible:
+              draftCanvas.status === 'streaming'
+                ? true
+                : draftCanvas.content.length > 200,
+            status: 'streaming',
+          };
+
+        case 'title':
+          return {
+            ...draftCanvas,
+            title: delta.content as string,
+          };
+
+        case 'suggestion':
+          setTimeout(() => {
+            setOptimisticSuggestions((currentSuggestions) => [
+              ...currentSuggestions,
+              delta.content as Suggestion,
+            ]);
+          }, 0);
+
+          return draftCanvas;
+
+        case 'clear':
+          return {
+            ...draftCanvas,
+            content: '',
+            status: 'streaming',
+          };
+
+        case 'finish':
+          return {
+            ...draftCanvas,
+            status: 'idle',
+          };
+
+        default:
+          return draftCanvas;
+      }
+    });
+  }, [streamingData, setCanvas]);
+}

components/custom/use-click-outside.tsx

@@ -0,0 +1,31 @@
+import { useEffect, RefObject } from 'react';
+
+type Handler = (event: MouseEvent | TouchEvent) => void;
+
+function useClickOutside<T extends HTMLElement = HTMLElement>(
+  ref: RefObject<T>,
+  handler: Handler,
+  mouseEvent: 'mousedown' | 'mouseup' = 'mousedown'
+): void {
+  useEffect(() => {
+    const listener = (event: MouseEvent | TouchEvent) => {
+      const el = ref?.current;
+
+      if (!el || el.contains((event?.target as Node) || null)) {
+        return;
+      }
+
+      handler(event);
+    };
+
+    document.addEventListener(mouseEvent, listener);
+    document.addEventListener('touchstart', listener);
+
+    return () => {
+      document.removeEventListener(mouseEvent, listener);
+      document.removeEventListener('touchstart', listener);
+    };
+  }, [ref, handler, mouseEvent]);
+}
+
+export default useClickOutside;

components/custom/use-debounce.tsx

@@ -0,0 +1,21 @@
+import { useCallback, useRef } from "react";
+
+export function useDebounce<T extends (...args: any[]) => any>(
+  callback: T,
+  delay: number,
+): (...args: Parameters<T>) => void {
+  const timeoutRef = useRef<NodeJS.Timeout | null>(null);
+
+  return useCallback(
+    (...args: Parameters<T>) => {
+      if (timeoutRef.current) {
+        clearTimeout(timeoutRef.current);
+      }
+
+      timeoutRef.current = setTimeout(() => {
+        callback(...args);
+      }, delay);
+    },
+    [callback, delay],
+  );
+}

db/queries.ts

@@ -1,11 +1,11 @@
 "server-only";
 
 import { genSaltSync, hashSync } from "bcrypt-ts";
-import { desc, eq } from "drizzle-orm";
+import { and, asc, desc, eq, gt } from "drizzle-orm";
 import { drizzle } from "drizzle-orm/postgres-js";
 import postgres from "postgres";
 
-import { user, chat, User } from "./schema";
+import { user, chat, User, document, Suggestion } from "./schema";
 
 // Optionally, if not using email/pass login, you can
 // use the Drizzle adapter for Auth.js / NextAuth
@@ -98,3 +98,108 @@ export async function getChatById({ id }: { id: string }) {
     throw error;
   }
 }
+
+export async function saveDocument({
+  id,
+  title,
+  content,
+  userId,
+}: {
+  id: string;
+  title: string;
+  content: string;
+  userId: string;
+}) {
+  try {
+    return await db.insert(document).values({
+      id,
+      title,
+      content,
+      userId,
+      createdAt: new Date(),
+    });
+  } catch (error) {
+    console.error("Failed to save document in database");
+    throw error;
+  }
+}
+
+export async function getDocumentsById({ id }: { id: string }) {
+  try {
+    const documents = await db
+      .select()
+      .from(document)
+      .where(eq(document.id, id))
+      .orderBy(asc(document.createdAt));
+
+    return documents;
+  } catch (error) {
+    console.error("Failed to get document by id from database");
+    throw error;
+  }
+}
+
+export async function getDocumentById({ id }: { id: string }) {
+  try {
+    const [selectedDocument] = await db
+      .select()
+      .from(document)
+      .where(eq(document.id, id))
+      .orderBy(desc(document.createdAt));
+
+    return selectedDocument;
+  } catch (error) {
+    console.error("Failed to get document by id from database");
+    throw error;
+  }
+}
+
+export async function deleteDocumentsByIdAfterTimestamp({
+  id,
+  timestamp,
+}: {
+  id: string;
+  timestamp: Date;
+}) {
+  try {
+    return await db
+      .delete(document)
+      .where(and(eq(document.id, id), gt(document.createdAt, timestamp)));
+  } catch (error) {
+    console.error(
+      "Failed to delete documents by id after timestamp from database",
+    );
+    throw error;
+  }
+}
+
+export async function saveSuggestions({
+  suggestions,
+}: {
+  suggestions: Array<Suggestion>;
+}) {
+  try {
+    return await db.insert(Suggestion).values(suggestions);
+  } catch (error) {
+    console.error("Failed to save suggestions in database");
+    throw error;
+  }
+}
+
+export async function getSuggestionsByDocumentId({
+  documentId,
+}: {
+  documentId: string;
+}) {
+  try {
+    return await db
+      .select()
+      .from(Suggestion)
+      .where(and(eq(Suggestion.documentId, documentId)));
+  } catch (error) {
+    console.error(
+      "Failed to get suggestions by document version from database",
+    );
+    throw error;
+  }
+}

db/schema.ts

@@ -1,24 +1,80 @@
-import { Message } from "ai";
+import { Message } from 'ai';
-import { InferSelectModel } from "drizzle-orm";
+import { InferSelectModel } from 'drizzle-orm';
-import { pgTable, varchar, timestamp, json, uuid } from "drizzle-orm/pg-core";
+import {
+  pgTable,
+  varchar,
+  timestamp,
+  json,
+  uuid,
+  text,
+  primaryKey,
+  foreignKey,
+  boolean,
+} from 'drizzle-orm/pg-core';
 
-export const user = pgTable("User", {
+export const user = pgTable('User', {
-  id: uuid("id").primaryKey().notNull().defaultRandom(),
+  id: uuid('id').primaryKey().notNull().defaultRandom(),
-  email: varchar("email", { length: 64 }).notNull(),
+  email: varchar('email', { length: 64 }).notNull(),
-  password: varchar("password", { length: 64 }),
+  password: varchar('password', { length: 64 }),
 });
 
 export type User = InferSelectModel<typeof user>;
 
-export const chat = pgTable("Chat", {
+export const chat = pgTable('Chat', {
-  id: uuid("id").primaryKey().notNull().defaultRandom(),
+  id: uuid('id').primaryKey().notNull().defaultRandom(),
-  createdAt: timestamp("createdAt").notNull(),
+  createdAt: timestamp('createdAt').notNull(),
-  messages: json("messages").notNull(),
+  messages: json('messages').notNull(),
-  userId: uuid("userId")
+  userId: uuid('userId')
     .notNull()
     .references(() => user.id),
 });
 
-export type Chat = Omit<InferSelectModel<typeof chat>, "messages"> & {
+export type Chat = Omit<InferSelectModel<typeof chat>, 'messages'> & {
   messages: Array<Message>;
 };
+
+export const document = pgTable(
+  'Document',
+  {
+    id: uuid('id').notNull().defaultRandom(),
+    createdAt: timestamp('createdAt').notNull(),
+    title: text('title').notNull(),
+    content: text('content'),
+    userId: uuid('userId')
+      .notNull()
+      .references(() => user.id),
+  },
+  (table) => {
+    return {
+      pk: primaryKey({ columns: [table.id, table.createdAt] }),
+    };
+  }
+);
+
+export type Document = InferSelectModel<typeof document>;
+
+export const Suggestion = pgTable(
+  'Suggestion',
+  {
+    id: uuid('id').notNull().defaultRandom(),
+    documentId: uuid('documentId').notNull(),
+    documentCreatedAt: timestamp('documentCreatedAt').notNull(),
+    originalText: text('originalText').notNull(),
+    suggestedText: text('suggestedText').notNull(),
+    description: text('description'),
+    isResolved: boolean('isResolved').notNull().default(false),
+    userId: uuid('userId')
+      .notNull()
+      .references(() => user.id),
+    createdAt: timestamp('createdAt').notNull(),
+  },
+  (table) => ({
+    pk: primaryKey({ columns: [table.id] }),
+    documentRef: foreignKey({
+      columns: [table.documentId, table.documentCreatedAt],
+      foreignColumns: [document.id, document.createdAt],
+    }),
+  })
+);
+
+export type Suggestion = InferSelectModel<typeof Suggestion>;

lib/drizzle/0001_sparkling_blue_marvel.sql

@@ -0,0 +1,39 @@
+CREATE TABLE IF NOT EXISTS "Suggestion" (
+	"id" uuid DEFAULT gen_random_uuid() NOT NULL,
+	"documentId" uuid NOT NULL,
+	"documentCreatedAt" timestamp NOT NULL,
+	"originalText" text NOT NULL,
+	"suggestedText" text NOT NULL,
+	"description" text,
+	"isResolved" boolean DEFAULT false NOT NULL,
+	"userId" uuid NOT NULL,
+	"createdAt" timestamp NOT NULL,
+	CONSTRAINT "Suggestion_id_pk" PRIMARY KEY("id")
+);
+--> statement-breakpoint
+CREATE TABLE IF NOT EXISTS "Document" (
+	"id" uuid DEFAULT gen_random_uuid() NOT NULL,
+	"createdAt" timestamp NOT NULL,
+	"title" text NOT NULL,
+	"content" text,
+	"userId" uuid NOT NULL,
+	CONSTRAINT "Document_id_createdAt_pk" PRIMARY KEY("id","createdAt")
+);
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Suggestion" ADD CONSTRAINT "Suggestion_userId_User_id_fk" FOREIGN KEY ("userId") REFERENCES "public"."User"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Suggestion" ADD CONSTRAINT "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk" FOREIGN KEY ("documentId","documentCreatedAt") REFERENCES "public"."Document"("id","createdAt") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Document" ADD CONSTRAINT "Document_userId_User_id_fk" FOREIGN KEY ("userId") REFERENCES "public"."User"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;

lib/drizzle/meta/0001_snapshot.json

@@ -0,0 +1,259 @@
+{
+  "id": "f3d3437c-4735-4c91-80af-1014048a904e",
+  "prevId": "715ec9ec-6715-4d0f-9f6c-9b5c7f09827c",
+  "version": "7",
+  "dialect": "postgresql",
+  "tables": {
+    "public.Suggestion": {
+      "name": "Suggestion",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "documentId": {
+          "name": "documentId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "documentCreatedAt": {
+          "name": "documentCreatedAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "originalText": {
+          "name": "originalText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "suggestedText": {
+          "name": "suggestedText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "description": {
+          "name": "description",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "isResolved": {
+          "name": "isResolved",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true,
+          "default": false
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Suggestion_userId_User_id_fk": {
+          "name": "Suggestion_userId_User_id_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk": {
+          "name": "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "Document",
+          "columnsFrom": [
+            "documentId",
+            "documentCreatedAt"
+          ],
+          "columnsTo": [
+            "id",
+            "createdAt"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Suggestion_id_pk": {
+          "name": "Suggestion_id_pk",
+          "columns": [
+            "id"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Chat": {
+      "name": "Chat",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "messages": {
+          "name": "messages",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Chat_userId_User_id_fk": {
+          "name": "Chat_userId_User_id_fk",
+          "tableFrom": "Chat",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Document": {
+      "name": "Document",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "title": {
+          "name": "title",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "content": {
+          "name": "content",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Document_userId_User_id_fk": {
+          "name": "Document_userId_User_id_fk",
+          "tableFrom": "Document",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Document_id_createdAt_pk": {
+          "name": "Document_id_createdAt_pk",
+          "columns": [
+            "id",
+            "createdAt"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.User": {
+      "name": "User",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "email": {
+          "name": "email",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "password": {
+          "name": "password",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    }
+  },
+  "enums": {},
+  "schemas": {},
+  "sequences": {},
+  "_meta": {
+    "columns": {},
+    "schemas": {},
+    "tables": {}
+  }
+}

lib/drizzle/meta/_journal.json

@@ -8,6 +8,13 @@
       "when": 1728598022383,
       "tag": "0000_keen_devos",
       "breakpoints": true
+    },
+    {
+      "idx": 1,
+      "version": "7",
+      "when": 1730207363999,
+      "tag": "0001_sparkling_blue_marvel",
+      "breakpoints": true
     }
   ]
 }

lib/editor/DiffType.js

@@ -0,0 +1,7 @@
+// Taken from https://github.com/hamflx/prosemirror-diff/blob/master/src/DiffType.js
+
+export const DiffType = {
+  Unchanged: 0,
+  Deleted: -1,
+  Inserted: 1,
+};

lib/editor/diff.js

@@ -0,0 +1,511 @@
+// Taken from https://github.com/hamflx/prosemirror-diff/blob/master/src/diff.js
+
+import { diff_match_patch } from 'diff-match-patch';
+import { Fragment, Node } from 'prosemirror-model';
+
+import { DiffType } from './DiffType.js';
+
+export const patchDocumentNode = (schema, oldNode, newNode) => {
+  assertNodeTypeEqual(oldNode, newNode);
+
+  const finalLeftChildren = [];
+  const finalRightChildren = [];
+
+  const oldChildren = normalizeNodeContent(oldNode);
+  const newChildren = normalizeNodeContent(newNode);
+  const oldChildLen = oldChildren.length;
+  const newChildLen = newChildren.length;
+  const minChildLen = Math.min(oldChildLen, newChildLen);
+
+  let left = 0;
+  let right = 0;
+
+  // console.log("==> searching same left");
+  for (; left < minChildLen; left++) {
+    const oldChild = oldChildren[left];
+    const newChild = newChildren[left];
+    if (!isNodeEqual(oldChild, newChild)) {
+      break;
+    }
+    finalLeftChildren.push(...ensureArray(oldChild));
+  }
+
+  // console.log("==> searching same right");
+  for (; right + left + 1 < minChildLen; right++) {
+    const oldChild = oldChildren[oldChildLen - right - 1];
+    const newChild = newChildren[newChildLen - right - 1];
+    if (!isNodeEqual(oldChild, newChild)) {
+      break;
+    }
+    finalRightChildren.unshift(...ensureArray(oldChild));
+  }
+
+  // console.log(
+  //   `==> eq left:${left}, right:${right}`,
+  //   [...finalLeftChildren],
+  //   [...finalRightChildren],
+  // );
+
+  const diffOldChildren = oldChildren.slice(left, oldChildLen - right);
+  const diffNewChildren = newChildren.slice(left, newChildLen - right);
+  // console.log(
+  //   "==> diff children",
+  //   diffOldChildren.length,
+  //   diffNewChildren.length,
+  // );
+  if (diffOldChildren.length && diffNewChildren.length) {
+    const matchedNodes = matchNodes(
+      schema,
+      diffOldChildren,
+      diffNewChildren
+    ).sort((a, b) => b.count - a.count);
+    const bestMatch = matchedNodes[0];
+    if (bestMatch) {
+      // console.log("==> bestMatch", bestMatch);
+      const { oldStartIndex, newStartIndex, oldEndIndex, newEndIndex } =
+        bestMatch;
+      const oldBeforeMatchChildren = diffOldChildren.slice(0, oldStartIndex);
+      const newBeforeMatchChildren = diffNewChildren.slice(0, newStartIndex);
+      // console.log(
+      //   "==> before match",
+      //   oldBeforeMatchChildren.length,
+      //   newBeforeMatchChildren.length,
+      //   oldBeforeMatchChildren,
+      //   newBeforeMatchChildren,
+      // );
+      finalLeftChildren.push(
+        ...patchRemainNodes(
+          schema,
+          oldBeforeMatchChildren,
+          newBeforeMatchChildren
+        )
+      );
+      finalLeftChildren.push(
+        ...diffOldChildren.slice(oldStartIndex, oldEndIndex)
+      );
+      // console.log("==> match", oldEndIndex - oldStartIndex);
+
+      const oldAfterMatchChildren = diffOldChildren.slice(oldEndIndex);
+      const newAfterMatchChildren = diffNewChildren.slice(newEndIndex);
+      // console.log(
+      //   "==> after match",
+      //   oldAfterMatchChildren.length,
+      //   newAfterMatchChildren.length,
+      // );
+      finalRightChildren.unshift(
+        ...patchRemainNodes(
+          schema,
+          oldAfterMatchChildren,
+          newAfterMatchChildren
+        )
+      );
+    } else {
+      // console.log("==> no best match found");
+      finalLeftChildren.push(
+        ...patchRemainNodes(schema, diffOldChildren, diffNewChildren)
+      );
+    }
+    // console.log("==> matchedNodes", matchedNodes);
+  } else {
+    finalLeftChildren.push(
+      ...patchRemainNodes(schema, diffOldChildren, diffNewChildren)
+    );
+  }
+
+  return createNewNode(oldNode, [...finalLeftChildren, ...finalRightChildren]);
+};
+
+const matchNodes = (schema, oldChildren, newChildren) => {
+  // console.log("==> matchNodes", oldChildren, newChildren);
+  const matches = [];
+  for (
+    let oldStartIndex = 0;
+    oldStartIndex < oldChildren.length;
+    oldStartIndex++
+  ) {
+    const oldStartNode = oldChildren[oldStartIndex];
+    const newStartIndex = findMatchNode(newChildren, oldStartNode);
+
+    if (newStartIndex !== -1) {
+      let oldEndIndex = oldStartIndex + 1;
+      let newEndIndex = newStartIndex + 1;
+      for (
+        ;
+        oldEndIndex < oldChildren.length && newEndIndex < newChildren.length;
+        oldEndIndex++, newEndIndex++
+      ) {
+        const oldEndNode = oldChildren[oldEndIndex];
+        if (!isNodeEqual(newChildren[newEndIndex], oldEndNode)) {
+          break;
+        }
+      }
+      // console.log(
+      //   "==> match",
+      //   oldStartIndex,
+      //   oldEndIndex,
+      //   newStartIndex,
+      //   newEndIndex,
+      // );
+      matches.push({
+        oldStartIndex,
+        newStartIndex,
+        oldEndIndex,
+        newEndIndex,
+        count: newEndIndex - newStartIndex,
+      });
+    }
+  }
+  return matches;
+};
+
+const findMatchNode = (children, node, startIndex = 0) => {
+  for (let i = startIndex; i < children.length; i++) {
+    if (isNodeEqual(children[i], node)) {
+      return i;
+    }
+  }
+  return -1;
+};
+
+const patchRemainNodes = (schema, oldChildren, newChildren) => {
+  const finalLeftChildren = [];
+  const finalRightChildren = [];
+  const oldChildLen = oldChildren.length;
+  const newChildLen = newChildren.length;
+  let left = 0;
+  let right = 0;
+  while (oldChildLen - left - right > 0 && newChildLen - left - right > 0) {
+    const leftOldNode = oldChildren[left];
+    const leftNewNode = newChildren[left];
+    const rightOldNode = oldChildren[oldChildLen - right - 1];
+    const rightNewNode = newChildren[newChildLen - right - 1];
+    let updateLeft =
+      !isTextNode(leftOldNode) && matchNodeType(leftOldNode, leftNewNode);
+    let updateRight =
+      !isTextNode(rightOldNode) && matchNodeType(rightOldNode, rightNewNode);
+    if (Array.isArray(leftOldNode) && Array.isArray(leftNewNode)) {
+      finalLeftChildren.push(
+        ...patchTextNodes(schema, leftOldNode, leftNewNode)
+      );
+      left += 1;
+      continue;
+    }
+
+    if (updateLeft && updateRight) {
+      const equalityLeft = computeChildEqualityFactor(leftOldNode, leftNewNode);
+      const equalityRight = computeChildEqualityFactor(
+        rightOldNode,
+        rightNewNode
+      );
+      if (equalityLeft < equalityRight) {
+        updateLeft = false;
+      } else {
+        updateRight = false;
+      }
+    }
+    if (updateLeft) {
+      finalLeftChildren.push(
+        patchDocumentNode(schema, leftOldNode, leftNewNode)
+      );
+      left += 1;
+    } else if (updateRight) {
+      finalRightChildren.unshift(
+        patchDocumentNode(schema, rightOldNode, rightNewNode)
+      );
+      right += 1;
+    } else {
+      // todo
+      finalLeftChildren.push(
+        createDiffNode(schema, leftOldNode, DiffType.Deleted)
+      );
+      finalLeftChildren.push(
+        createDiffNode(schema, leftNewNode, DiffType.Inserted)
+      );
+      left += 1;
+      // delete and insert
+    }
+  }
+
+  const deleteNodeLen = oldChildLen - left - right;
+  const insertNodeLen = newChildLen - left - right;
+  if (deleteNodeLen) {
+    finalLeftChildren.push(
+      ...oldChildren
+        .slice(left, left + deleteNodeLen)
+        .flat()
+        .map((node) => createDiffNode(schema, node, DiffType.Deleted))
+    );
+  }
+
+  if (insertNodeLen) {
+    finalRightChildren.unshift(
+      ...newChildren
+        .slice(left, left + insertNodeLen)
+        .flat()
+        .map((node) => createDiffNode(schema, node, DiffType.Inserted))
+    );
+  }
+
+  return [...finalLeftChildren, ...finalRightChildren];
+};
+
+export const patchTextNodes = (schema, oldNode, newNode) => {
+  const dmp = new diff_match_patch();
+  const oldText = oldNode.map((n) => getNodeText(n)).join('');
+  const newText = newNode.map((n) => getNodeText(n)).join('');
+  const diff = dmp.diff_main(oldText, newText);
+  let oldLen = 0;
+  let newLen = 0;
+  const res = diff
+    .map((d) => {
+      const [type, content] = [d[0], d[1]];
+      const node = createTextNode(
+        schema,
+        content,
+        type !== DiffType.Unchanged ? createDiffMark(schema, type) : []
+      );
+      const oldFrom = oldLen;
+      const oldTo = oldFrom + (type === DiffType.Inserted ? 0 : content.length);
+      const newFrom = newLen;
+      const newTo = newFrom + (type === DiffType.Deleted ? 0 : content.length);
+      oldLen = oldTo;
+      newLen = newTo;
+      return { node, type, oldFrom, oldTo, newFrom, newTo };
+    })
+    .map(({ node, type, oldFrom, oldTo, newFrom, newTo }) => {
+      if (type === DiffType.Deleted) {
+        const textItems = findTextNodes(oldNode, oldFrom, oldTo).filter(
+          (n) => Object.keys(n.node.attrs ?? {}).length || n.node.marks?.length
+        );
+        return applyTextNodeAttrsMarks(schema, node, oldFrom, textItems);
+      } else {
+        const textItems = findTextNodes(newNode, newFrom, newTo).filter(
+          (n) => Object.keys(n.node.attrs ?? {}).length || n.node.marks?.length
+        );
+        return applyTextNodeAttrsMarks(schema, node, newFrom, textItems);
+      }
+    });
+  return res.flat(Infinity);
+};
+
+const findTextNodes = (textNodes, from, to) => {
+  const result = [];
+  let start = 0;
+  for (let i = 0; i < textNodes.length && start < to; i++) {
+    const node = textNodes[i];
+    const text = getNodeText(node);
+    const end = start + text.length;
+    const intersect =
+      (start >= from && start < to) ||
+      (end > from && end <= to) ||
+      (start <= from && end >= to);
+    if (intersect) {
+      result.push({ node, from: start, to: end });
+    }
+    start += text.length;
+  }
+  return result;
+};
+
+const applyTextNodeAttrsMarks = (schema, node, base, textItems) => {
+  if (!textItems.length) {
+    return node;
+  }
+
+  const baseMarks = node.marks ?? [];
+  const firstItem = textItems[0];
+  const nodeText = getNodeText(node);
+  const nodeEnd = base + nodeText.length;
+  const result = [];
+  if (firstItem.from - base > 0) {
+    result.push(
+      createTextNode(
+        schema,
+        nodeText.slice(0, firstItem.from - base),
+        baseMarks
+      )
+    );
+  }
+  for (let i = 0; i < textItems.length; i++) {
+    const { from, node: textNode, to } = textItems[i];
+    result.push(
+      createTextNode(
+        schema,
+        nodeText.slice(Math.max(from, base) - base, to - base),
+        [...baseMarks, ...(textNode.marks ?? [])]
+      )
+    );
+    const nextFrom = i + 1 < textItems.length ? textItems[i + 1].from : nodeEnd;
+    if (nextFrom > to) {
+      result.push(
+        createTextNode(
+          schema,
+          nodeText.slice(to - base, nextFrom - base),
+          baseMarks
+        )
+      );
+    }
+  }
+  return result;
+};
+
+export const computeChildEqualityFactor = (node1, node2) => {
+  return 0;
+};
+
+export const assertNodeTypeEqual = (node1, node2) => {
+  if (getNodeProperty(node1, 'type') !== getNodeProperty(node2, 'type')) {
+    throw new Error(`node type not equal: ${node1.type} !== ${node2.type}`);
+  }
+};
+
+export const ensureArray = (value) => {
+  return Array.isArray(value) ? value : [value];
+};
+
+export const isNodeEqual = (node1, node2) => {
+  const isNode1Array = Array.isArray(node1);
+  const isNode2Array = Array.isArray(node2);
+  if (isNode1Array !== isNode2Array) {
+    return false;
+  }
+  if (isNode1Array) {
+    return (
+      node1.length === node2.length &&
+      node1.every((node, index) => isNodeEqual(node, node2[index]))
+    );
+  }
+
+  const type1 = getNodeProperty(node1, 'type');
+  const type2 = getNodeProperty(node2, 'type');
+  if (type1 !== type2) {
+    return false;
+  }
+  if (isTextNode(node1)) {
+    const text1 = getNodeProperty(node1, 'text');
+    const text2 = getNodeProperty(node2, 'text');
+    if (text1 !== text2) {
+      return false;
+    }
+  }
+  const attrs1 = getNodeAttributes(node1);
+  const attrs2 = getNodeAttributes(node2);
+  const attrs = [...new Set([...Object.keys(attrs1), ...Object.keys(attrs2)])];
+  for (const attr of attrs) {
+    if (attrs1[attr] !== attrs2[attr]) {
+      return false;
+    }
+  }
+  const marks1 = getNodeMarks(node1);
+  const marks2 = getNodeMarks(node2);
+  if (marks1.length !== marks2.length) {
+    return false;
+  }
+  for (let i = 0; i < marks1.length; i++) {
+    if (!isNodeEqual(marks1[i], marks2[i])) {
+      return false;
+    }
+  }
+  const children1 = getNodeChildren(node1);
+  const children2 = getNodeChildren(node2);
+  if (children1.length !== children2.length) {
+    return false;
+  }
+  for (let i = 0; i < children1.length; i++) {
+    if (!isNodeEqual(children1[i], children2[i])) {
+      return false;
+    }
+  }
+  return true;
+};
+
+export const normalizeNodeContent = (node) => {
+  const content = getNodeChildren(node) ?? [];
+  const res = [];
+  for (let i = 0; i < content.length; i++) {
+    const child = content[i];
+    if (isTextNode(child)) {
+      const textNodes = [];
+      for (
+        let textNode = content[i];
+        i < content.length && isTextNode(textNode);
+        textNode = content[++i]
+      ) {
+        textNodes.push(textNode);
+      }
+      i--;
+      res.push(textNodes);
+    } else {
+      res.push(child);
+    }
+  }
+  return res;
+};
+
+export const getNodeProperty = (node, property) => {
+  if (property === 'type') {
+    return node.type?.name;
+  }
+  return node[property];
+};
+export const getNodeAttribute = (node, attribute) =>
+  node.attrs ? node.attrs[attribute] : undefined;
+export const getNodeAttributes = (node) =>
+  node.attrs ? node.attrs : undefined;
+export const getNodeMarks = (node) => node.marks ?? [];
+export const getNodeChildren = (node) => node.content?.content ?? [];
+export const getNodeText = (node) => node.text;
+export const isTextNode = (node) => node.type?.name === 'text';
+export const matchNodeType = (node1, node2) =>
+  node1.type?.name === node2.type?.name ||
+  (Array.isArray(node1) && Array.isArray(node2));
+export const createNewNode = (oldNode, children) => {
+  if (!oldNode.type) {
+    throw new Error('oldNode.type is undefined');
+  }
+  return new Node(
+    oldNode.type,
+    oldNode.attrs,
+    Fragment.fromArray(children),
+    oldNode.marks
+  );
+};
+export const createDiffNode = (schema, node, type) => {
+  return mapDocumentNode(node, (node) => {
+    if (isTextNode(node)) {
+      return createTextNode(schema, getNodeText(node), [
+        ...(node.marks || []),
+        createDiffMark(schema, type),
+      ]);
+    }
+    return node;
+  });
+};
+function mapDocumentNode(node, mapper) {
+  const copy = node.copy(
+    Fragment.from(
+      node.content.content
+        .map((node) => mapDocumentNode(node, mapper))
+        .filter((n) => n)
+    )
+  );
+  return mapper(copy) || copy;
+}
+export const createDiffMark = (schema, type) => {
+  if (type === DiffType.Inserted) {
+    return schema.mark('diffMark', { type });
+  }
+  if (type === DiffType.Deleted) {
+    return schema.mark('diffMark', { type });
+  }
+  throw new Error('type is not valid');
+};
+export const createTextNode = (schema, content, marks = []) => {
+  return schema.text(content, marks);
+};
+
+export const diffEditor = (schema, oldDoc, newDoc) => {
+  const oldNode = Node.fromJSON(schema, oldDoc);
+  const newNode = Node.fromJSON(schema, newDoc);
+  return patchDocumentNode(schema, oldNode, newNode);
+};

lib/editor/index.js

@@ -0,0 +1,2 @@
+export * from "./diff.js";
+export * from "./DiffType";

lib/editor/react-renderer.tsx

@@ -0,0 +1,11 @@
+import { createRoot } from "react-dom/client";
+
+export class ReactRenderer {
+  static render(component: React.ReactElement, dom: HTMLElement) {
+    const root = createRoot(dom);
+    root.render(component);
+    return {
+      destroy: () => root.unmount(),
+    };
+  }
+}

lib/editor/suggestions.tsx

@@ -0,0 +1,118 @@
+import { Node } from 'prosemirror-model';
+import { PluginKey, Plugin } from 'prosemirror-state';
+import { DecorationSet, EditorView } from 'prosemirror-view';
+import { createRoot } from 'react-dom/client';
+
+import { Suggestion as PreviewSuggestion } from '@/components/custom/suggestion';
+import { Suggestion } from '@/db/schema';
+
+export interface UISuggestion extends Suggestion {
+  selectionStart: number;
+  selectionEnd: number;
+}
+
+interface Position {
+  start: number;
+  end: number;
+}
+
+function findPositionsInDoc(doc: Node, searchText: string): Position | null {
+  let positions: { start: number; end: number } | null = null;
+
+  doc.nodesBetween(0, doc.content.size, (node, pos) => {
+    if (node.isText && node.text) {
+      const index = node.text.indexOf(searchText);
+
+      if (index !== -1) {
+        positions = {
+          start: pos + index,
+          end: pos + index + searchText.length,
+        };
+
+        return false;
+      }
+    }
+
+    return true;
+  });
+
+  return positions;
+}
+
+export function projectWithHighlights(
+  doc: Node,
+  suggestions: Array<Suggestion>
+): Array<UISuggestion> {
+  return suggestions.map((suggestion) => {
+    const positions = findPositionsInDoc(doc, suggestion.originalText);
+
+    if (!positions) {
+      return {
+        ...suggestion,
+        selectionStart: 0,
+        selectionEnd: 0,
+      };
+    }
+
+    return {
+      ...suggestion,
+      selectionStart: positions.start,
+      selectionEnd: positions.end,
+    };
+  });
+}
+
+export function createSuggestionWidget(
+  suggestion: UISuggestion,
+  view: EditorView
+): { dom: HTMLElement; destroy: () => void } {
+  const dom = document.createElement('span');
+  const root = createRoot(dom);
+
+  const onApply = () => {
+    const { state, dispatch } = view;
+    const tr = state.tr.replaceWith(
+      suggestion.selectionStart,
+      suggestion.selectionEnd,
+      state.schema.text(suggestion.suggestedText)
+    );
+
+    dispatch(tr);
+  };
+
+  root.render(<PreviewSuggestion suggestion={suggestion} onApply={onApply} />);
+
+  return {
+    dom,
+    destroy: () => {
+      // Wrapping unmount in setTimeout to avoid synchronous unmounting during render
+      setTimeout(() => {
+        root.unmount();
+      }, 0);
+    },
+  };
+}
+
+export const suggestionsPluginKey = new PluginKey('suggestions');
+export const suggestionsPlugin = new Plugin({
+  key: suggestionsPluginKey,
+  state: {
+    init() {
+      return { decorations: DecorationSet.empty, selected: null };
+    },
+    apply(tr, state) {
+      const newDecorations = tr.getMeta(suggestionsPluginKey);
+      if (newDecorations) return newDecorations;
+
+      return {
+        decorations: state.decorations.map(tr.mapping, tr.doc),
+        selected: state.selected,
+      };
+    },
+  },
+  props: {
+    decorations(state) {
+      return this.getState(state)?.decorations ?? DecorationSet.empty;
+    },
+  },
+});

lib/model.ts

@@ -1,17 +0,0 @@
-// Define your models here.
-export const models = [
-  {
-    label: 'GPT 4o mini',
-    name: 'gpt-4o-mini',
-    description: 'Small model for fast, lightweight tasks',
-  },
-  {
-    label: 'GPT 4o',
-    name: 'gpt-4o',
-    description: 'For complex, multi-step tasks',
-  },
-] as const;
-
-export const DEFAULT_MODEL_NAME: Model['name'] = 'gpt-4o-mini';
-
-export type Model = (typeof models)[number];

lib/utils.ts

@@ -1,14 +1,15 @@
 import {
+  CoreAssistantMessage,
   CoreMessage,
   CoreToolMessage,
   generateId,
   Message,
   ToolInvocation,
-} from "ai";
+} from 'ai';
-import { clsx, type ClassValue } from "clsx";
+import { clsx, type ClassValue } from 'clsx';
-import { twMerge } from "tailwind-merge";
+import { twMerge } from 'tailwind-merge';
 
-import { Chat } from "@/db/schema";
+import { Chat } from '@/db/schema';
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
@@ -24,7 +25,7 @@ export const fetcher = async (url: string) => {
 
   if (!res.ok) {
     const error = new Error(
-      "An error occurred while fetching the data.",
+      'An error occurred while fetching the data.'
     ) as ApplicationError;
 
     error.info = await res.json();
@@ -37,16 +38,16 @@ export const fetcher = async (url: string) => {
 };
 
 export function getLocalStorage(key: string) {
-  if (typeof window !== "undefined") {
+  if (typeof window !== 'undefined') {
-    return JSON.parse(localStorage.getItem(key) || "[]");
+    return JSON.parse(localStorage.getItem(key) || '[]');
   }
   return [];
 }
 
 export function generateUUID(): string {
-  return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, function (c) {
+  return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function (c) {
     const r = (Math.random() * 16) | 0;
-    const v = c === "x" ? r : (r & 0x3) | 0x8;
+    const v = c === 'x' ? r : (r & 0x3) | 0x8;
     return v.toString(16);
   });
 }
@@ -64,13 +65,13 @@ function addToolMessageToChat({
         ...message,
         toolInvocations: message.toolInvocations.map((toolInvocation) => {
           const toolResult = toolMessage.content.find(
-            (tool) => tool.toolCallId === toolInvocation.toolCallId,
+            (tool) => tool.toolCallId === toolInvocation.toolCallId
           );
 
           if (toolResult) {
             return {
               ...toolInvocation,
-              state: "result",
+              state: 'result',
               result: toolResult.result,
             };
           }
@@ -85,28 +86,28 @@ function addToolMessageToChat({
 }
 
 export function convertToUIMessages(
-  messages: Array<CoreMessage>,
+  messages: Array<CoreMessage>
 ): Array<Message> {
   return messages.reduce((chatMessages: Array<Message>, message) => {
-    if (message.role === "tool") {
+    if (message.role === 'tool') {
       return addToolMessageToChat({
         toolMessage: message as CoreToolMessage,
         messages: chatMessages,
       });
     }
 
-    let textContent = "";
+    let textContent = '';
     let toolInvocations: Array<ToolInvocation> = [];
 
-    if (typeof message.content === "string") {
+    if (typeof message.content === 'string') {
       textContent = message.content;
     } else if (Array.isArray(message.content)) {
       for (const content of message.content) {
-        if (content.type === "text") {
+        if (content.type === 'text') {
           textContent += content.text;
-        } else if (content.type === "tool-call") {
+        } else if (content.type === 'tool-call') {
           toolInvocations.push({
-            state: "call",
+            state: 'call',
             toolCallId: content.toolCallId,
             toolName: content.toolName,
             args: content.args,
@@ -131,8 +132,92 @@ export function getTitleFromChat(chat: Chat) {
   const firstMessage = messages[0];
 
   if (!firstMessage) {
-    return "Untitled";
+    return 'Untitled';
   }
 
   return firstMessage.content;
 }
+
+const emptyAssistantMessage = [
+  {
+    role: 'assistant',
+    content: [
+      {
+        type: 'text',
+        text: '',
+      },
+    ],
+  },
+];
+
+export function sanitizeResponseMessages(
+  messages: Array<CoreToolMessage | CoreAssistantMessage>
+): Array<CoreToolMessage | CoreAssistantMessage> {
+  let toolResultIds: Array<string> = [];
+
+  for (const message of messages) {
+    if (message.role === 'tool') {
+      for (const content of message.content) {
+        if (content.type === 'tool-result') {
+          toolResultIds.push(content.toolCallId);
+        }
+      }
+    }
+  }
+
+  const messagesBySanitizedContent = messages.map((message) => {
+    if (message.role !== 'assistant') return message;
+
+    if (typeof message.content === 'string') return message;
+
+    const sanitizedContent = message.content.filter((content) =>
+      content.type === 'tool-call'
+        ? toolResultIds.includes(content.toolCallId)
+        : content.type === 'text'
+          ? content.text.length > 0
+          : true
+    );
+
+    return {
+      ...message,
+      content: sanitizedContent,
+    };
+  });
+
+  return messagesBySanitizedContent.filter(
+    (message) => message.content.length > 0
+  );
+}
+
+export function sanitizeUIMessages(messages: Array<Message>): Array<Message> {
+  const messagesBySanitizedToolInvocations = messages.map((message) => {
+    if (message.role !== 'assistant') return message;
+
+    if (!message.toolInvocations) return message;
+
+    let toolResultIds: Array<string> = [];
+
+    for (const toolInvocation of message.toolInvocations) {
+      if (toolInvocation.state === 'result') {
+        toolResultIds.push(toolInvocation.toolCallId);
+      }
+    }
+
+    const sanitizedToolInvocations = message.toolInvocations.filter(
+      (toolInvocation) =>
+        toolInvocation.state === 'result' ||
+        toolResultIds.includes(toolInvocation.toolCallId)
+    );
+
+    return {
+      ...message,
+      toolInvocations: sanitizedToolInvocations,
+    };
+  });
+
+  return messagesBySanitizedToolInvocations.filter(
+    (message) =>
+      message.content.length > 0 ||
+      (message.toolInvocations && message.toolInvocations.length > 0)
+  );
+}

package.json

@@ -23,12 +23,13 @@
     "@vercel/analytics": "^1.3.1",
     "@vercel/blob": "^0.24.1",
     "@vercel/postgres": "^0.10.0",
-    "ai": "3.4.9",
+    "ai": "3.4.27",
     "bcrypt-ts": "^5.0.2",
     "class-variance-authority": "^0.7.0",
     "classnames": "^2.5.1",
     "clsx": "^2.1.1",
     "date-fns": "^4.1.0",
+    "diff-match-patch": "^1.0.5",
     "dotenv": "^16.4.5",
     "drizzle-orm": "^0.34.0",
     "framer-motion": "^11.3.19",
@@ -37,7 +38,16 @@
     "next": "15.0.1",
     "next-auth": "5.0.0-beta.25",
     "next-themes": "^0.3.0",
+    "orderedmap": "^2.1.1",
     "postgres": "^3.4.4",
+    "prosemirror-example-setup": "^1.2.3",
+    "prosemirror-inputrules": "^1.4.0",
+    "prosemirror-markdown": "^1.13.1",
+    "prosemirror-model": "^1.23.0",
+    "prosemirror-schema-basic": "^1.2.3",
+    "prosemirror-schema-list": "^1.4.1",
+    "prosemirror-state": "^1.4.3",
+    "prosemirror-view": "^1.34.3",
     "react": "19.0.0-rc-45804af1-20241021",
     "react-dom": "19.0.0-rc-45804af1-20241021",
     "react-markdown": "^9.0.1",
@@ -51,6 +61,7 @@
     "zod": "^3.23.8"
   },
   "devDependencies": {
+    "@tailwindcss/typography": "^0.5.15",
     "@types/d3-scale": "^4.0.8",
     "@types/node": "^20",
     "@types/pdf-parse": "^1.1.4",

tailwind.config.ts

@@ -72,7 +72,7 @@ const config: Config = {
       },
     },
   },
-  plugins: [require('tailwindcss-animate')],
+  plugins: [require('tailwindcss-animate'), require('@tailwindcss/typography')],
   safelist: ['w-32', 'w-44', 'w-52'],
 };
 export default config;