streaming response

pages/chat/@id/+Page.tsx

@@ -19,10 +19,20 @@ import {
 import { usePageContext } from "vike-react/usePageContext";
 import { useData } from "vike-react/useData";
 import type { Data } from "./+data";
-import type { CommittedMessage, DraftMessage } from "../../../types";
+import type {
+  CommittedMessage,
+  DraftMessage,
+  OtherParameters,
+} from "../../../types";
 import Markdown from "react-markdown";
-import { IconTrash, IconEdit, IconCheck, IconX } from "@tabler/icons-react";
-import { useTRPC } from "../../../trpc/client";
+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 { nanoid } from "nanoid";
 import type { Conversation } from "../../../database/common";
@@ -49,6 +59,7 @@ export default function ChatPage() {
   const setParameters = useStore((state) => state.setParameters);
   const setLoading = useStore((state) => state.setLoading);
   const trpc = useTRPC();
+  const trpcClient = useTRPCClient();
   const queryClient = useQueryClient();
 
   const messagesResult = useQuery(
@@ -334,83 +345,94 @@ export default function ChatPage() {
     })
   );
 
-  const sendMessage = useMutation(
-    trpc.chat.sendMessage.mutationOptions({
-      onMutate: async ({
-        conversationId,
-        messages,
-        systemPrompt,
-        parameters,
-      }) => {
-        /** Cancel affected queries that may be in-flight: */
-        await queryClient.cancelQueries({
-          queryKey: trpc.chat.messages.fetchByConversationId.queryKey({
-            conversationId,
-          }),
-        });
-        /** Optimistically update the affected queries in react-query's cache: */
-        const previousMessages: Array<CommittedMessage> | undefined =
-          await queryClient.getQueryData(
-            trpc.chat.messages.fetchByConversationId.queryKey({
-              conversationId,
-            })
-          );
-        if (!previousMessages) {
-          return {
-            previousMessages: [],
-            newMessages: [],
-          };
+  // 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);
+
+  // Function to send message using subscription
+  const sendSubscriptionMessage = async ({
+    conversationId,
+    messages,
+    systemPrompt,
+    parameters,
+  }: {
+    conversationId: string;
+    messages: Array<DraftMessage | CommittedMessage>;
+    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);
+
+            // 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",
+              message: "An error occurred while sending the message",
+            });
+          },
         }
-        const newMessages: Array<CommittedMessage> = [
-          ...previousMessages,
-          {
-            /** placeholder id; will be overwritten when we get the true id from the backend */
-            id: nanoid(),
-            conversationId,
-            // content: messages[messages.length - 1].content,
-            // role: "user" as const,
-            ...messages[messages.length - 1],
-            index: previousMessages.length,
-            createdAt: new Date().toISOString(),
-          } as CommittedMessage,
-        ];
-        queryClient.setQueryData(
-          trpc.chat.messages.fetchByConversationId.queryKey({
-            conversationId,
-          }),
-          newMessages
-        );
-        return { previousMessages, newMessages };
-      },
-      onSettled: async (data, variables, context) => {
-        await queryClient.invalidateQueries({
-          queryKey: trpc.chat.messages.fetchByConversationId.queryKey({
-            conversationId,
-          }),
-        });
-        await queryClient.invalidateQueries({
-          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
-            conversationId,
-          }),
-        });
-        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.messages.fetchByConversationId.queryKey({
-            conversationId,
-          }),
-          context.previousMessages
-        );
-      },
-    })
-  );
+      );
+
+      // 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",
+        message: "Failed to start message sending process",
+      });
+    }
+  };
 
   // State for editing facts
   const [editingFactId, setEditingFactId] = useState<string | null>(null);
@@ -483,6 +505,8 @@ export default function ChatPage() {
             });
           }}
         />
+        {isSendingMessage && <IconLoaderQuarter size={16} stroke={1.5} />}
+        {sendMessageStatus && <span>{sendMessageStatus.message}</span>}
       </div>
       <Tabs defaultValue="message">
         <Tabs.List>
@@ -504,7 +528,7 @@ export default function ChatPage() {
               if (e.key === "Enter") {
                 e.preventDefault();
                 setLoading(true);
-                await sendMessage.mutateAsync({
+                await sendSubscriptionMessage({
                   conversationId,
                   messages: [
                     ...(messages || []),