diff --git a/apps/web/app/(app)/[emailAccountId]/simple/Summary.tsx b/apps/web/app/(app)/[emailAccountId]/simple/Summary.tsx index 074494fd2c..ae51dc1e3e 100644 --- a/apps/web/app/(app)/[emailAccountId]/simple/Summary.tsx +++ b/apps/web/app/(app)/[emailAccountId]/simple/Summary.tsx @@ -1,6 +1,6 @@ "use client"; -import { useCompletion } from "ai/react"; +import { useCompletion } from "@ai-sdk/react"; import { useEffect } from "react"; import { ButtonLoader } from "@/components/Loading"; import { ViewMoreButton } from "@/app/(app)/[emailAccountId]/simple/ViewMoreButton"; diff --git a/apps/web/app/api/chat/route.ts b/apps/web/app/api/chat/route.ts index 2e31a23131..1d19d0ef9c 100644 --- a/apps/web/app/api/chat/route.ts +++ b/apps/web/app/api/chat/route.ts @@ -1,4 +1,4 @@ -import { appendClientMessage, appendResponseMessages } from "ai"; +import { convertToModelMessages, type UIMessage } from "ai"; import { z } from "zod"; import { withEmailAccount } from "@/utils/middleware"; import { getEmailAccountWithAi } from "@/utils/user/get"; @@ -6,7 +6,9 @@ import { NextResponse } from "next/server"; import { aiProcessAssistantChat } from "@/utils/ai/assistant/chat"; import { createScopedLogger } from "@/utils/logger"; import prisma from "@/utils/prisma"; -import { Prisma, type ChatMessage } from "@prisma/client"; +import type { Prisma } from "@prisma/client"; +import { convertToUIMessages } from "@/components/assistant-chat/helpers"; +import { captureException } from "@/utils/error"; export const maxDuration = 120; @@ -21,19 +23,8 @@ const assistantInputSchema = z.object({ id: z.string(), message: z.object({ id: z.string(), - createdAt: z.coerce.date(), role: z.enum(["user"]), - content: z.string().min(1).max(3000), parts: z.array(textPartSchema), - // experimental_attachments: z - // .array( - // z.object({ - // url: z.string().url(), - // name: z.string().min(1).max(100), - // contentType: z.enum(["image/png", "image/jpg", "image/jpeg"]), - // }), - // ) - // .optional(), }), }); @@ -49,7 +40,6 @@ export const POST = withEmailAccount(async (request) => { if (error) return NextResponse.json({ error: error.errors }, { status: 400 }); - // create chat if it doesn't exist const chat = (await getChatById(data.id)) || (await createNewChat({ emailAccountId, chatId: data.id })); @@ -69,63 +59,27 @@ export const POST = withEmailAccount(async (request) => { } const { message } = data; - const mappedDbMessages = chat.messages.map((dbMsg: ChatMessage) => { - return { - ...dbMsg, - role: convertDbRoleToSdkRole(dbMsg.role), - content: "", - parts: dbMsg.parts as any, - }; - }); - - const messages = appendClientMessage({ - messages: mappedDbMessages, - message, - }); + const uiMessages = [...convertToUIMessages(chat), message]; await saveChatMessage({ chat: { connect: { id: chat.id } }, id: message.id, role: "user", parts: message.parts, - // attachments: message.experimental_attachments ?? [], }); try { const result = await aiProcessAssistantChat({ - messages, + messages: convertToModelMessages(uiMessages), emailAccountId, user, - onFinish: async ({ response }) => { - const assistantMessages = response.messages.filter( - (message) => message.role === "assistant", - ); - const assistantId = getTrailingMessageId(assistantMessages); - - if (!assistantId) { - logger.error("No assistant message found!", { response }); - throw new Error("No assistant message found!"); - } - - // handles all tool calls - const [, assistantMessage] = appendResponseMessages({ - messages: [message], - responseMessages: response.messages, - }); - - await saveChatMessage({ - id: assistantId, - chat: { connect: { id: chat.id } }, - role: assistantMessage.role, - parts: assistantMessage.parts - ? (assistantMessage.parts as Prisma.InputJsonValue) - : Prisma.JsonNull, - // attachments: assistantMessage.experimental_attachments ?? [], - }); - }, }); - return result.toDataStreamResponse(); + return result.toUIMessageStreamResponse({ + onFinish: async ({ messages }) => { + await saveChatMessages(messages, chat.id); + }, + }); } catch (error) { logger.error("Error in assistant chat", { error }); return NextResponse.json( @@ -155,10 +109,6 @@ async function createNewChat({ } } -async function saveChatMessage(message: Prisma.ChatMessageCreateInput) { - return prisma.chatMessage.create({ data: message }); -} - async function getChatById(chatId: string) { const chat = await prisma.chat.findUnique({ where: { id: chatId }, @@ -167,29 +117,23 @@ async function getChatById(chatId: string) { return chat; } -function convertDbRoleToSdkRole( - role: string, -): "user" | "assistant" | "system" | "data" { - switch (role) { - case "user": - return "user"; - case "assistant": - return "assistant"; - case "system": - return "system"; - case "data": - return "data"; - default: - return "assistant"; - } +async function saveChatMessage(message: Prisma.ChatMessageCreateInput) { + return prisma.chatMessage.create({ data: message }); } -function getTrailingMessageId( - messages: Array, -): string | null { - const trailingMessage = messages.at(-1); - - if (!trailingMessage) return null; - - return trailingMessage.id; +async function saveChatMessages(messages: UIMessage[], chatId: string) { + try { + return prisma.chatMessage.createMany({ + data: messages.map((message) => ({ + id: message.id, + chatId, + role: message.role, + parts: message.parts as Prisma.InputJsonValue, + })), + }); + } catch (error) { + logger.error("Failed to save chat messages", { error, chatId }); + captureException(error, { extra: { chatId } }); + throw error; + } } diff --git a/apps/web/components/assistant-chat/ChatContext.tsx b/apps/web/components/assistant-chat/ChatContext.tsx index 4560124faa..400b3626a9 100644 --- a/apps/web/components/assistant-chat/ChatContext.tsx +++ b/apps/web/components/assistant-chat/ChatContext.tsx @@ -24,9 +24,7 @@ export function ChatProvider({ export function useChat(): ChatContextType { const context = useContext(ChatContext); if (context === undefined) { - // TODO: throw error once this feature is live - // throw new Error("useChat must be used within a ChatProvider"); - return { setInput: null }; + throw new Error("useChat must be used within a ChatProvider"); } return context; } diff --git a/apps/web/components/assistant-chat/chat.tsx b/apps/web/components/assistant-chat/chat.tsx index 81f0e35fce..9ec70423cb 100644 --- a/apps/web/components/assistant-chat/chat.tsx +++ b/apps/web/components/assistant-chat/chat.tsx @@ -1,9 +1,8 @@ "use client"; -import type React from "react"; import { useEffect, useMemo, useState } from "react"; -import { type ScopedMutator, SWRConfig, useSWRConfig } from "swr"; -import type { UIMessage } from "ai"; +import { useSWRConfig } from "swr"; +import { DefaultChatTransport } from "ai"; import { useChat } from "@ai-sdk/react"; import { ArrowLeftToLineIcon, @@ -20,7 +19,6 @@ import { ResizablePanelGroup } from "@/components/ui/resizable"; import { ResizablePanel } from "@/components/ui/resizable"; import { AssistantTabs } from "@/app/(app)/[emailAccountId]/assistant/AssistantTabs"; import { ChatProvider } from "./ChatContext"; -import { SWRProvider } from "@/providers/SWRProvider"; import { Button } from "@/components/ui/button"; import { DropdownMenu, @@ -31,18 +29,12 @@ import { import { useChats } from "@/hooks/useChats"; import { LoadingContent } from "@/components/LoadingContent"; import { useChatMessages } from "@/hooks/useChatMessages"; -import type { GetChatResponse } from "@/app/api/chats/[chatId]/route"; import { useIsMobile } from "@/hooks/use-mobile"; import { ExamplesDialog } from "@/components/assistant-chat/examples-dialog"; import { Tooltip } from "@/components/Tooltip"; import { toastError } from "@/components/Toast"; - -// Some mega hacky code used here to workaround AI SDK's use of SWR -// AI SDK uses SWR too and this messes with the global SWR config -// Wrapping in SWRConfig to disable global fetcher for this component -// https://github.com/vercel/ai/issues/3214#issuecomment-2675872030 -// We then re-enable the regular SWRProvider in the AssistantTabs component -// AI SDK v5 won't use SWR anymore so we can remove this workaround +import type { ChatMessage } from "@/components/assistant-chat/types"; +import { convertToUIMessages } from "./helpers"; const MAX_MESSAGES = 20; @@ -65,9 +57,6 @@ export function Chat(props: ChatProps) { } function ChatWithEmptySWR(props: ChatProps & { chatId: string }) { - // Use parent SWR config for mutate - const { mutate } = useSWRConfig(); - const { data } = useChatMessages(props.chatId); const [{ input, tab }] = useQueryStates({ @@ -81,20 +70,13 @@ function ChatWithEmptySWR(props: ChatProps & { chatId: string }) { }, [input]); return ( - - - + ); } @@ -102,31 +84,38 @@ function ChatInner({ chatId, initialMessages, emailAccountId, - mutate, initialInput, tab, }: ChatProps & { chatId: string; - initialMessages: Array; - mutate: ScopedMutator; + initialMessages: ChatMessage[]; initialInput?: string; tab?: string; }) { - const chat = useChat({ + const { mutate } = useSWRConfig(); + + const [input, setInput] = useState(""); + + const chat = useChat({ id: chatId, - api: "/api/chat", - experimental_prepareRequestBody: (body) => ({ - id: chatId, - message: body.messages.at(-1), + transport: new DefaultChatTransport({ + api: "/api/chat", + headers: { + [EMAIL_ACCOUNT_HEADER]: emailAccountId, + }, + prepareSendMessagesRequest({ messages, id, body }) { + return { + body: { + id, + message: messages.at(-1), + ...body, + }, + }; + }, }), - initialMessages, - initialInput, + messages: initialMessages, experimental_throttle: 100, - sendExtraMessageFields: true, generateId: generateUUID, - headers: { - [EMAIL_ACCOUNT_HEADER]: emailAccountId, - }, onFinish: () => { mutate("/api/user/rules"); }, @@ -139,22 +128,40 @@ function ChatInner({ }, }); + const handleSubmit = () => { + chat.sendMessage({ + role: "user", + parts: [ + { + type: "text", + text: input, + }, + ], + }); + + setInput(""); + }; + const isMobile = useIsMobile(); - const chatPanel = ; + const chatPanel = ( + + ); return ( - + {tab ? ( - {/* re-enable the regular SWRProvider */} - - - + @@ -168,47 +175,44 @@ function ChatInner({ ); } -function ChatUI({ chat }: { chat: ReturnType }) { - const { - messages, - setMessages, - handleSubmit, - input, - setInput, - status, - stop, - reload, - } = chat; - +function ChatUI({ + chat: { messages, setMessages, status, stop, regenerate }, + input, + setInput, + handleSubmit, +}: { + chat: ReturnType>; + input: string; + setInput: (input: string) => void; + handleSubmit: () => void; +}) { return (
- -
- {messages.length > MAX_MESSAGES ? ( -
- The chat is too long. Please start a new conversation. -
- ) : ( -
- )} - -
- - - - +
+ {messages.length > MAX_MESSAGES ? ( +
+ The chat is too long. Please start a new conversation.
+ ) : ( +
+ )} + +
+ + + +
- - - +
+ +
{ - return ( - chat?.messages.map((message) => ({ - id: message.id, - parts: message.parts as UIMessage["parts"], - role: message.role as UIMessage["role"], - // Note: content will soon be deprecated in @ai-sdk/react - content: "", - createdAt: message.createdAt, - // experimental_attachments: (message.attachments as Array) ?? [], - })) || [] - ); -} diff --git a/apps/web/components/assistant-chat/data-stream-handler.tsx b/apps/web/components/assistant-chat/data-stream-handler.tsx deleted file mode 100644 index b58a016c13..0000000000 --- a/apps/web/components/assistant-chat/data-stream-handler.tsx +++ /dev/null @@ -1,31 +0,0 @@ -"use client"; - -import { useChat } from "@ai-sdk/react"; -import { useEffect, useRef } from "react"; - -export type DataStreamDelta = { - type: - | "text-delta" - | "code-delta" - | "sheet-delta" - | "image-delta" - | "title" - | "id" - | "suggestion" - | "clear" - | "finish" - | "kind"; - content: string; -}; - -export function DataStreamHandler({ id }: { id: string }) { - const { data: dataStream } = useChat({ id }); - const lastProcessedIndex = useRef(-1); - - useEffect(() => { - if (!dataStream?.length) return; - lastProcessedIndex.current = dataStream.length - 1; - }, [dataStream]); - - return null; -} diff --git a/apps/web/components/assistant-chat/helpers.ts b/apps/web/components/assistant-chat/helpers.ts new file mode 100644 index 0000000000..070d74c2e3 --- /dev/null +++ b/apps/web/components/assistant-chat/helpers.ts @@ -0,0 +1,20 @@ +import type { UIMessage, UIMessagePart } from "ai"; +import type { GetChatResponse } from "@/app/api/chats/[chatId]/route"; +import type { + ChatMessage, + ChatTools, + CustomUIDataTypes, +} from "@/components/assistant-chat/types"; + +export function convertToUIMessages(chat: GetChatResponse): ChatMessage[] { + return ( + chat?.messages.map((message) => ({ + id: message.id, + role: message.role as UIMessage["role"], + parts: message.parts as UIMessagePart[], + // metadata: { + // createdAt: formatISO(message.createdAt), + // }, + })) || [] + ); +} diff --git a/apps/web/components/assistant-chat/message-editor.tsx b/apps/web/components/assistant-chat/message-editor.tsx index e18bc61daf..af88aa50be 100644 --- a/apps/web/components/assistant-chat/message-editor.tsx +++ b/apps/web/components/assistant-chat/message-editor.tsx @@ -1,6 +1,6 @@ "use client"; -import type { Message } from "ai"; +import type { UIMessage } from "ai"; import { type Dispatch, type SetStateAction, @@ -11,24 +11,27 @@ import { import type { UseChatHelpers } from "@ai-sdk/react"; import { Button } from "@/components/ui/button"; import { Textarea } from "@/components/ui/textarea"; +import type { ChatMessage } from "@/components/assistant-chat/types"; // import { deleteTrailingMessages } from "@/app/(app)/assistant/chat/actions"; -export type MessageEditorProps = { - message: Message; +type MessageEditorProps = { + message: ChatMessage; setMode: Dispatch>; - setMessages: UseChatHelpers["setMessages"]; - reload: UseChatHelpers["reload"]; + setMessages: UseChatHelpers["setMessages"]; + regenerate: UseChatHelpers["regenerate"]; }; export function MessageEditor({ message, setMode, setMessages, - reload, + regenerate, }: MessageEditorProps) { const [isSubmitting, setIsSubmitting] = useState(false); - const [draftContent, setDraftContent] = useState(message.content); + const [draftContent, setDraftContent] = useState( + getTextFromMessage(message), + ); const textareaRef = useRef(null); useEffect(() => { @@ -88,7 +91,6 @@ export function MessageEditor({ if (index !== -1) { const updatedMessage = { ...message, - content: draftContent, parts: [{ type: "text", text: draftContent }], }; @@ -99,7 +101,7 @@ export function MessageEditor({ }); setMode("view"); - reload(); + regenerate(); }} > {isSubmitting ? "Sending..." : "Send"} @@ -108,3 +110,10 @@ export function MessageEditor({
); } + +function getTextFromMessage(message: ChatMessage): string { + return message.parts + .filter((part) => part.type === "text") + .map((part) => part.text) + .join(""); +} diff --git a/apps/web/components/assistant-chat/message.tsx b/apps/web/components/assistant-chat/message.tsx index 61721ba110..81522df85c 100644 --- a/apps/web/components/assistant-chat/message.tsx +++ b/apps/web/components/assistant-chat/message.tsx @@ -1,7 +1,6 @@ "use client"; import { memo, useState } from "react"; -import type { UIMessage } from "ai"; import type { UseChatHelpers } from "@ai-sdk/react"; import { AnimatePresence, motion } from "framer-motion"; import { SparklesIcon } from "lucide-react"; @@ -10,19 +9,27 @@ import { Markdown } from "./markdown"; import { cn } from "@/utils"; import { MessageEditor } from "./message-editor"; import { MessageReasoning } from "./message-reasoning"; -import { ToolCard } from "@/components/assistant-chat/tools"; -import { Skeleton } from "@/components/ui/skeleton"; +import { + AddToKnowledgeBase, + BasicToolInfo, + CreatedRuleToolCard, + UpdateAbout, + UpdatedLearnedPatterns, + UpdatedRuleActions, + UpdatedRuleConditions, +} from "@/components/assistant-chat/tools"; +import type { ChatMessage } from "@/components/assistant-chat/types"; const PurePreviewMessage = ({ message, isLoading, setMessages, - reload, + regenerate, }: { - message: UIMessage; + message: ChatMessage; isLoading: boolean; - setMessages: UseChatHelpers["setMessages"]; - reload: UseChatHelpers["reload"]; + setMessages: UseChatHelpers["setMessages"]; + regenerate: UseChatHelpers["regenerate"]; }) => { const [mode, setMode] = useState<"view" | "edit">("view"); @@ -62,7 +69,7 @@ const PurePreviewMessage = ({ ); } @@ -94,33 +101,236 @@ const PurePreviewMessage = ({ message={message} setMode={setMode} setMessages={setMessages} - reload={reload} + regenerate={regenerate} />
); } } - if (type === "tool-invocation") { - const { toolInvocation } = part; - const { toolName, toolCallId, state } = toolInvocation; + if (type === "tool-getUserRulesAndSettings") { + const { toolCallId, state } = part; - if (state === "call") { - return ; + if (state === "input-available") { + return ; } - if (state === "result") { + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ; + } + } + + if (type === "tool-getLearnedPatterns") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ; + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ; + } + } + + if (type === "tool-createRule") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ( + + ); + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ( + + ); + } + } + + if (type === "tool-updateRuleConditions") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ( + + ); + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ( + + ); + } + } + + if (type === "tool-updateRuleActions") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ( + + ); + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ( + + ); + } + } + + if (type === "tool-updateLearnedPatterns") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ( + + ); + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + return ( - ); } } + + if (type === "tool-updateAbout") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ; + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ; + } + } + + if (type === "tool-addToKnowledgeBase") { + const { toolCallId, state } = part; + + if (state === "input-available") { + return ; + } + + if (state === "output-available") { + const { output } = part; + + if ("error" in output) { + return ( + + ); + } + + return ; + } + } })}
@@ -129,6 +339,10 @@ const PurePreviewMessage = ({ ); }; +function ErrorToolCard({ error }: { error: string }) { + return
Error: {error}
; +} + export const PreviewMessage = memo( PurePreviewMessage, (prevProps, nextProps) => { diff --git a/apps/web/components/assistant-chat/messages.tsx b/apps/web/components/assistant-chat/messages.tsx index e6569c4ff0..1e2a87578e 100644 --- a/apps/web/components/assistant-chat/messages.tsx +++ b/apps/web/components/assistant-chat/messages.tsx @@ -5,21 +5,22 @@ import { Overview } from "./overview"; import { memo } from "react"; import equal from "fast-deep-equal"; import type { UseChatHelpers } from "@ai-sdk/react"; +import type { ChatMessage } from "@/components/assistant-chat/types"; interface MessagesProps { - status: UseChatHelpers["status"]; - messages: Array; - setMessages: UseChatHelpers["setMessages"]; - reload: UseChatHelpers["reload"]; + status: UseChatHelpers["status"]; + messages: Array; + setMessages: UseChatHelpers["setMessages"]; + regenerate: UseChatHelpers["regenerate"]; isArtifactVisible: boolean; - setInput: UseChatHelpers["setInput"]; + setInput: (input: string) => void; } function PureMessages({ status, messages, setMessages, - reload, + regenerate, setInput, }: MessagesProps) { const [messagesContainerRef, messagesEndRef] = @@ -38,7 +39,7 @@ function PureMessages({ message={message} isLoading={status === "streaming" && messages.length - 1 === index} setMessages={setMessages} - reload={reload} + regenerate={regenerate} /> ))} diff --git a/apps/web/components/assistant-chat/multimodal-input.tsx b/apps/web/components/assistant-chat/multimodal-input.tsx index 7da355e7fa..f5089315df 100644 --- a/apps/web/components/assistant-chat/multimodal-input.tsx +++ b/apps/web/components/assistant-chat/multimodal-input.tsx @@ -11,6 +11,7 @@ import { Button } from "@/components/ui/button"; import { Textarea } from "@/components/ui/textarea"; // import { SuggestedActions } from "./suggested-actions"; import { cn } from "@/utils"; +import type { ChatMessage } from "@/components/assistant-chat/types"; function PureMultimodalInput({ // chatId, @@ -25,16 +26,16 @@ function PureMultimodalInput({ className, }: { // chatId?: string; - input: UseChatHelpers["input"]; - setInput: UseChatHelpers["setInput"]; - status: UseChatHelpers["status"]; + input: UseChatHelpers["input"]; + setInput: UseChatHelpers["setInput"]; + status: UseChatHelpers["status"]; stop: () => void; // attachments: Array; // setAttachments: Dispatch>>; // messages: Array; - setMessages: UseChatHelpers["setMessages"]; + setMessages: UseChatHelpers["setMessages"]; // append: UseChatHelpers["append"]; - handleSubmit: UseChatHelpers["handleSubmit"]; + handleSubmit: UseChatHelpers["handleSubmit"]; className?: string; }) { const textareaRef = useRef(null); @@ -168,7 +169,7 @@ function PureStopButton({ setMessages, }: { stop: () => void; - setMessages: UseChatHelpers["setMessages"]; + setMessages: UseChatHelpers["setMessages"]; }) { return (