import { ActionIcon, Box, Group, HoverCard, JsonInput, List, ScrollArea, Stack, Tabs, Text, Textarea, TextInput, Transition, useMantineTheme, } from "@mantine/core"; import { memo, useEffect, useState } from "react"; import { defaultParameters, defaultSystemPrompt, useStore, } from "../../../state"; import { usePageContext } from "vike-react/usePageContext"; import { useData } from "vike-react/useData"; import type { Data } from "./+data"; import type { CommittedMessage, DraftMessage, OtherParameters, SendMessageStatus, SendMessageStatusUI, } from "@/types"; import Markdown from "react-markdown"; import { IconTrash, IconEdit, IconCheck, IconX, IconLoaderQuarter, } from "@tabler/icons-react"; import { useTRPC, useTRPCClient } from "../../trpc-client"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import type { Conversation } from "@/database/common"; export default function ChatPage() { const pageContext = usePageContext(); const conversationId = pageContext.routeParams.id; const { conversation, // messages: initialMessages, // facts: initialFacts, // factTriggers: initialFactTriggers, } = useData(); const conversationTitle = conversation?.title; const message = useStore((state) => state.message); const systemPrompt = useStore((state) => state.systemPrompt); const parameters = useStore((state) => state.parameters); const loading = useStore((state) => state.loading); const setMessage = useStore((state) => state.setMessage); const setSystemPrompt = useStore((state) => state.setSystemPrompt); const setParameters = useStore((state) => state.setParameters); const setLoading = useStore((state) => state.setLoading); const trpc = useTRPC(); const trpcClient = useTRPCClient(); const queryClient = useQueryClient(); const messagesResult = useQuery( trpc.chat.messages.fetchByConversationId.queryOptions({ conversationId, }) ); const messages: Array | undefined = messagesResult.data?.map((m) => ({ ...m, parts: m.parts.filter((p) => p.type === "text"), })) || []; const facts = useQuery( trpc.chat.facts.fetchByConversationId.queryOptions({ conversationId, }) ); const factTriggers = useQuery( trpc.chat.factTriggers.fetchByConversationId.queryOptions({ conversationId, }) ); const deleteFact = useMutation( trpc.chat.facts.deleteOne.mutationOptions({ onMutate: async ({ factId: factIdToDelete }) => { /** Cancel affected queries that may be in-flight: */ await queryClient.cancelQueries({ queryKey: trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), }); /** Optimistically update the affected queries in react-query's cache: */ const previousFacts = await queryClient.getQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }) ); if (!previousFacts) { return { previousFacts: [], newFacts: [], }; } const newFacts = previousFacts.filter((f) => f.id !== factIdToDelete); queryClient.setQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), newFacts ); return { previousFacts, newFacts }; }, onSettled: async (data, variables, context) => { await queryClient.invalidateQueries({ queryKey: trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), }); }, onError: async (error, variables, context) => { console.error(error); if (!context) return; queryClient.setQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), context.previousFacts ); }, }) ); const updateFact = useMutation( trpc.chat.facts.update.mutationOptions({ onMutate: async ({ factId, content }) => { /** Cancel affected queries that may be in-flight: */ await queryClient.cancelQueries({ queryKey: trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), }); /** Optimistically update the affected queries in react-query's cache: */ const previousFacts = await queryClient.getQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }) ); if (!previousFacts) { return { previousFacts: [], newFacts: [], }; } const newFacts = previousFacts.map((f) => f.id === factId ? { ...f, content } : f ); queryClient.setQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), newFacts ); return { previousFacts, newFacts }; }, onSettled: async (data, variables, context) => { await queryClient.invalidateQueries({ queryKey: trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), }); }, onError: async (error, variables, context) => { console.error(error); if (!context) return; queryClient.setQueryData( trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), context.previousFacts ); }, }) ); const deleteFactTrigger = useMutation( trpc.chat.factTriggers.deleteOne.mutationOptions({ onMutate: async ({ factTriggerId: factTriggerIdToDelete }) => { /** Cancel affected queries that may be in-flight: */ await queryClient.cancelQueries({ queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), }); /** Optimistically update the affected queries in react-query's cache: */ const previousFactTriggers = await queryClient.getQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }) ); if (!previousFactTriggers) { return { previousFactTriggers: [], newFactTriggers: [], }; } const newFactTriggers = previousFactTriggers.filter( (ft) => ft.id !== factTriggerIdToDelete ); queryClient.setQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), newFactTriggers ); return { previousFactTriggers, newFactTriggers }; }, onSettled: async (data, variables, context) => { await queryClient.invalidateQueries({ queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), }); }, onError: async (error, variables, context) => { console.error(error); if (!context) return; queryClient.setQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), context.previousFactTriggers ); }, }) ); const updateFactTrigger = useMutation( trpc.chat.factTriggers.update.mutationOptions({ onMutate: async ({ factTriggerId, content }) => { /** Cancel affected queries that may be in-flight: */ await queryClient.cancelQueries({ queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), }); /** Optimistically update the affected queries in react-query's cache: */ const previousFactTriggers = await queryClient.getQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }) ); if (!previousFactTriggers) { return { previousFactTriggers: [], newFactTriggers: [], }; } const newFactTriggers = previousFactTriggers.map((ft) => ft.id === factTriggerId ? { ...ft, content } : ft ); queryClient.setQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), newFactTriggers ); return { previousFactTriggers, newFactTriggers }; }, onSettled: async (data, variables, context) => { await queryClient.invalidateQueries({ queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), }); }, onError: async (error, variables, context) => { console.error(error); if (!context) return; queryClient.setQueryData( trpc.chat.factTriggers.fetchByConversationId.queryKey({ conversationId, }), context.previousFactTriggers ); }, }) ); const updateConversationTitle = useMutation( trpc.chat.conversations.updateTitle.mutationOptions({ onMutate: async ({ id, title }) => { /** Cancel affected queries that may be in-flight: */ await queryClient.cancelQueries({ queryKey: trpc.chat.conversations.fetchAll.queryKey(), }); /** Optimistically update the affected queries in react-query's cache: */ const previousConversations = await queryClient.getQueryData( trpc.chat.conversations.fetchAll.queryKey() ); if (!previousConversations) { return { previousConversations: [], newConversations: null, }; } const newConversations: Array = [ ...previousConversations, { ...conversation, title, } as Conversation, ]; queryClient.setQueryData( trpc.chat.conversations.fetchAll.queryKey(), newConversations ); return { previousConversations, newConversations }; }, onSettled: async (data, variables, context) => { await queryClient.invalidateQueries({ queryKey: trpc.chat.conversations.fetchOne.queryKey({ id: conversationId, }), }); await queryClient.invalidateQueries({ queryKey: trpc.chat.conversations.fetchAll.queryKey(), }); }, onError: async (error, variables, context) => { console.error(error); if (!context) return; queryClient.setQueryData( trpc.chat.conversations.fetchAll.queryKey(), context.previousConversations ); }, }) ); // Get state from Zustand store const sendMessageStatus = useStore((state) => state.sendMessageStatus); const isSendingMessage = useStore((state) => state.isSendingMessage); const setSendMessageStatus = useStore((state) => state.setSendMessageStatus); const setIsSendingMessage = useStore((state) => state.setIsSendingMessage); const theme = useMantineTheme(); // Function to send message using subscription const sendSubscriptionMessage = async ({ conversationId, messages, systemPrompt, parameters, }: { conversationId: string; messages: Array; systemPrompt: string; parameters: OtherParameters; }) => { setIsSendingMessage(true); setSendMessageStatus(null); try { // Create an abort controller for the subscription const abortController = new AbortController(); // Start the subscription const subscription = trpcClient.chat.sendMessage.subscribe( { conversationId, messages, systemPrompt, parameters, }, { signal: abortController.signal, onData: (data) => { setSendMessageStatus(data as SendMessageStatus); // If we've completed, update the UI and invalidate queries if (data.status === "completed") { setIsSendingMessage(false); // Invalidate queries to refresh the data queryClient.invalidateQueries({ queryKey: trpc.chat.messages.fetchByConversationId.queryKey({ conversationId, }), }); queryClient.invalidateQueries({ queryKey: trpc.chat.facts.fetchByConversationId.queryKey({ conversationId, }), }); queryClient.invalidateQueries({ queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey( { conversationId, } ), }); } else { setSendMessageStatus(data); } }, onError: (error) => { console.error("Subscription error:", error); setIsSendingMessage(false); setSendMessageStatus({ status: "error" as const, message: "An error occurred while sending the message", }); }, } ); // Return a function to unsubscribe if needed return () => { abortController.abort(); subscription.unsubscribe(); }; } catch (error) { console.error("Failed to start subscription:", error); setIsSendingMessage(false); setSendMessageStatus({ status: "error" as const, message: "Failed to start message sending process", }); } }; // State for editing facts const [editingFactId, setEditingFactId] = useState(null); const [editingFactContent, setEditingFactContent] = useState(""); // State for editing fact triggers const [editingFactTriggerId, setEditingFactTriggerId] = useState< string | null >(null); const [editingFactTriggerContent, setEditingFactTriggerContent] = useState(""); // Handle clicking outside to cancel editing useEffect(() => { function handleClickOutside(event: MouseEvent) { if (editingFactId && event.target instanceof Element) { const editingElement = event.target.closest(".editing-fact"); if (!editingElement) { setEditingFactId(null); } } if (editingFactTriggerId && event.target instanceof Element) { const editingElement = event.target.closest(".editing-fact-trigger"); if (!editingElement) { setEditingFactTriggerId(null); } } } document.addEventListener("mousedown", handleClickOutside); return () => { document.removeEventListener("mousedown", handleClickOutside); }; }, [editingFactId, editingFactTriggerId]); async function handleDeleteFact(factId: string) { await deleteFact.mutateAsync({ factId }); } async function handleUpdateFact(factId: string, content: string) { await updateFact.mutateAsync({ factId, content }); } async function handleDeleteFactTrigger(factTriggerId: string) { await deleteFactTrigger.mutateAsync({ factTriggerId }); } async function handleUpdateFactTrigger( factTriggerId: string, content: string ) { await updateFactTrigger.mutateAsync({ factTriggerId, content }); } return ( { // setConversationTitle(e.target.value); // }} onBlur={(e) => { updateConversationTitle.mutateAsync({ id: conversationId, title: e.target.value, }); }} variant="unstyled" styles={{ input: { // backgroundColor: "transparent", // border: "none", // padding: 0, // margin: 0, fontFamily: theme.headings.fontFamily, fontSize: theme.fontSizes.lg, lineHeight: theme.lineHeights["4xl"], }, wrapper: { marginTop: 0, }, }} /> {isSendingMessage && } {sendMessageStatus && ( )} Message System Prompt Parameters Facts Fact Triggers