diff --git a/app/(chat)/actions.ts b/app/(chat)/actions.ts index d88d8a7..e0f8771 100644 --- a/app/(chat)/actions.ts +++ b/app/(chat)/actions.ts @@ -1,6 +1,6 @@ 'use server'; -import { type CoreUserMessage, generateText } from 'ai'; +import { type CoreUserMessage, generateText, Message } from 'ai'; import { cookies } from 'next/headers'; import { customModel } from '@/lib/ai'; @@ -19,7 +19,7 @@ export async function saveModelId(model: string) { export async function generateTitleFromUserMessage({ message, }: { - message: CoreUserMessage; + message: Message; }) { const { text: title } = await generateText({ model: customModel('gpt-4o-mini'), diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts index 14eef57..dbf4cb4 100644 --- a/app/(chat)/api/chat/route.ts +++ b/app/(chat)/api/chat/route.ts @@ -1,6 +1,5 @@ import { type Message, - convertToCoreMessages, createDataStreamResponse, smoothStream, streamText, @@ -65,8 +64,7 @@ export async function POST(request: Request) { return new Response('Model not found', { status: 404 }); } - const coreMessages = convertToCoreMessages(messages); - const userMessage = getMostRecentUserMessage(coreMessages); + const userMessage = getMostRecentUserMessage(messages); if (!userMessage) { return new Response('No user message found', { status: 400 }); @@ -79,28 +77,20 @@ export async function POST(request: Request) { await saveChat({ id, userId: session.user.id, title }); } - const userMessageId = generateUUID(); - await saveMessages({ - messages: [ - { ...userMessage, id: userMessageId, createdAt: new Date(), chatId: id }, - ], + messages: [{ ...userMessage, createdAt: new Date(), chatId: id }], }); return createDataStreamResponse({ execute: (dataStream) => { - dataStream.writeData({ - type: 'user-message-id', - content: userMessageId, - }); - const result = streamText({ model: customModel(model.apiIdentifier), system: systemPrompt, - messages: coreMessages, + messages, maxSteps: 5, experimental_activeTools: allTools, experimental_transform: smoothStream({ chunking: 'word' }), + experimental_generateMessageId: generateUUID, tools: { getWeather, createDocument: createDocument({ session, dataStream, model }), @@ -120,16 +110,8 @@ export async function POST(request: Request) { await saveMessages({ messages: responseMessagesWithoutIncompleteToolCalls.map( (message) => { - const messageId = generateUUID(); - - if (message.role === 'assistant') { - dataStream.writeMessageAnnotation({ - messageIdFromServer: messageId, - }); - } - return { - id: messageId, + id: message.id, chatId: id, role: message.role, content: message.content, diff --git a/components/chat.tsx b/components/chat.tsx index a2ea8f5..fb9b2a3 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -7,7 +7,7 @@ import useSWR, { useSWRConfig } from 'swr'; import { ChatHeader } from '@/components/chat-header'; import type { Vote } from '@/lib/db/schema'; -import { fetcher } from '@/lib/utils'; +import { fetcher, generateUUID } from '@/lib/utils'; import { Block } from './block'; import { MultimodalInput } from './multimodal-input'; @@ -45,6 +45,8 @@ export function Chat({ body: { id, modelId: selectedModelId }, initialMessages, experimental_throttle: 100, + sendExtraMessageFields: true, + generateId: generateUUID, onFinish: () => { mutate('/api/history'); }, diff --git a/components/data-stream-handler.tsx b/components/data-stream-handler.tsx index 02a66ee..1328a36 100644 --- a/components/data-stream-handler.tsx +++ b/components/data-stream-handler.tsx @@ -5,7 +5,6 @@ import { useEffect, useRef } from 'react'; import { blockDefinitions, BlockKind } from './block'; import { Suggestion } from '@/lib/db/schema'; import { initialBlockData, useBlock } from '@/hooks/use-block'; -import { useUserMessageId } from '@/hooks/use-user-message-id'; export type DataStreamDelta = { type: @@ -17,14 +16,12 @@ export type DataStreamDelta = { | 'suggestion' | 'clear' | 'finish' - | 'user-message-id' | 'kind'; content: string | Suggestion; }; export function DataStreamHandler({ id }: { id: string }) { const { data: dataStream } = useChat({ id }); - const { setUserMessageIdFromServer } = useUserMessageId(); const { block, setBlock, setMetadata } = useBlock(); const lastProcessedIndex = useRef(-1); @@ -35,11 +32,6 @@ export function DataStreamHandler({ id }: { id: string }) { lastProcessedIndex.current = dataStream.length - 1; (newDeltas as DataStreamDelta[]).forEach((delta: DataStreamDelta) => { - if (delta.type === 'user-message-id') { - setUserMessageIdFromServer(delta.content as string); - return; - } - const blockDefinition = blockDefinitions.find( (blockDefinition) => blockDefinition.kind === block.kind, ); @@ -97,7 +89,7 @@ export function DataStreamHandler({ id }: { id: string }) { } }); }); - }, [dataStream, setBlock, setUserMessageIdFromServer, setMetadata, block]); + }, [dataStream, setBlock, setMetadata, block]); return null; } diff --git a/components/message-actions.tsx b/components/message-actions.tsx index fc29e08..dc67dc5 100644 --- a/components/message-actions.tsx +++ b/components/message-actions.tsx @@ -4,7 +4,6 @@ import { useSWRConfig } from 'swr'; import { useCopyToClipboard } from 'usehooks-ts'; import type { Vote } from '@/lib/db/schema'; -import { getMessageIdFromAnnotations } from '@/lib/utils'; import { CopyIcon, ThumbDownIcon, ThumbUpIcon } from './icons'; import { Button } from './ui/button'; @@ -62,13 +61,11 @@ export function PureMessageActions({ disabled={vote?.isUpvoted} variant="outline" onClick={async () => { - const messageId = getMessageIdFromAnnotations(message); - const upvote = fetch('/api/vote', { method: 'PATCH', body: JSON.stringify({ chatId, - messageId, + messageId: message.id, type: 'up', }), }); @@ -116,13 +113,11 @@ export function PureMessageActions({ variant="outline" disabled={vote && !vote.isUpvoted} onClick={async () => { - const messageId = getMessageIdFromAnnotations(message); - const downvote = fetch('/api/vote', { method: 'PATCH', body: JSON.stringify({ chatId, - messageId, + messageId: message.id, type: 'down', }), }); diff --git a/components/message-editor.tsx b/components/message-editor.tsx index 35fbcf9..9e845d6 100644 --- a/components/message-editor.tsx +++ b/components/message-editor.tsx @@ -6,7 +6,6 @@ import { Dispatch, SetStateAction, useEffect, useRef, useState } from 'react'; import { Textarea } from './ui/textarea'; import { deleteTrailingMessages } from '@/app/(chat)/actions'; import { toast } from 'sonner'; -import { useUserMessageId } from '@/hooks/use-user-message-id'; export type MessageEditorProps = { message: Message; @@ -25,7 +24,6 @@ export function MessageEditor({ setMessages, reload, }: MessageEditorProps) { - const { userMessageIdFromServer } = useUserMessageId(); const [isSubmitting, setIsSubmitting] = useState(false); const [draftContent, setDraftContent] = useState(message.content); @@ -74,16 +72,9 @@ export function MessageEditor({ disabled={isSubmitting} onClick={async () => { setIsSubmitting(true); - const messageId = userMessageIdFromServer ?? message.id; - - if (!messageId) { - toast.error('Something went wrong, please try again!'); - setIsSubmitting(false); - return; - } await deleteTrailingMessages({ - id: messageId, + id: message.id, }); setMessages((messages) => { diff --git a/hooks/use-user-message-id.ts b/hooks/use-user-message-id.ts deleted file mode 100644 index 63a3262..0000000 --- a/hooks/use-user-message-id.ts +++ /dev/null @@ -1,10 +0,0 @@ -'use client'; - -import useSWR from 'swr'; - -export function useUserMessageId() { - const { data: userMessageIdFromServer, mutate: setUserMessageIdFromServer } = - useSWR('userMessageIdFromServer', null); - - return { userMessageIdFromServer, setUserMessageIdFromServer }; -} diff --git a/lib/utils.ts b/lib/utils.ts index 60b0299..fa19776 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -126,9 +126,12 @@ export function convertToUIMessages( }, []); } +type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage; +type ResponseMessage = ResponseMessageWithoutId & { id: string }; + export function sanitizeResponseMessages( - messages: Array, -): Array { + messages: Array, +): Array { const toolResultIds: Array = []; for (const message of messages) { @@ -198,7 +201,7 @@ export function sanitizeUIMessages(messages: Array): Array { ); } -export function getMostRecentUserMessage(messages: Array) { +export function getMostRecentUserMessage(messages: Array) { const userMessages = messages.filter((message) => message.role === 'user'); return userMessages.at(-1); } @@ -212,13 +215,3 @@ export function getDocumentTimestampByIndex( return documents[index].createdAt; } - -export function getMessageIdFromAnnotations(message: Message) { - if (!message.annotations) return message.id; - - const [annotation] = message.annotations; - if (!annotation) return message.id; - - // @ts-expect-error messageIdFromServer is not defined in MessageAnnotation - return annotation.messageIdFromServer; -}