Merge branch 'master' of github.com:notedsource/Perplexica into hristo/deploy-on-gcp-gke

ui/components/ChatWindow.tsx

@@ -1,11 +1,13 @@
 'use client';
 
-import { useEffect, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { Document } from '@langchain/core/documents';
 import Navbar from './Navbar';
 import Chat from './Chat';
 import EmptyChat from './EmptyChat';
 import { toast } from 'sonner';
+import { useSearchParams } from 'next/navigation';
+import { getSuggestions } from '@/lib/actions';
 import { clientFetch } from '@/lib/utils';
 import { getAccessKey } from '@/lib/config';
 
@@ -14,10 +16,11 @@ export type Message = {
   createdAt: Date;
   content: string;
   role: 'user' | 'assistant';
+  suggestions?: string[];
   sources?: Document[];
 };
 
-const useSocket = (url: string) => {
+const useSocket = (url: string, setIsReady: (ready: boolean) => void) => {
   const [ws, setWs] = useState<WebSocket | null>(null);
 
   useEffect(() => {
@@ -107,9 +110,17 @@ const useSocket = (url: string) => {
 
         ws.onopen = () => {
           console.log('[DEBUG] open');
-          setWs(ws);
         };
 
+        const stateCheckInterval = setInterval(() => {
+          if (ws.readyState === 1) {
+            setIsReady(true);
+            clearInterval(stateCheckInterval);
+          }
+        }, 100);
+
+        setWs(ws);
+
         ws.onmessage = (e) => {
           const parsedData = JSON.parse(e.data);
           if (parsedData.type === 'error') {
@@ -128,19 +139,29 @@ const useSocket = (url: string) => {
       ws?.close();
       console.log('[DEBUG] closed');
     };
-  }, [ws, url]);
+  }, [ws, url, setIsReady]);
 
   return ws;
 };
 
 const ChatWindow = () => {
-  const ws = useSocket(process.env.NEXT_PUBLIC_WS_URL!);
+  const searchParams = useSearchParams();
+  const initialMessage = searchParams.get('q');
+
+  const [isReady, setIsReady] = useState(false);
+  const ws = useSocket(process.env.NEXT_PUBLIC_WS_URL!, setIsReady);
+
   const [chatHistory, setChatHistory] = useState<[string, string][]>([]);
   const [messages, setMessages] = useState<Message[]>([]);
+  const messagesRef = useRef<Message[]>([]);
   const [loading, setLoading] = useState(false);
   const [messageAppeared, setMessageAppeared] = useState(false);
   const [focusMode, setFocusMode] = useState('webSearch');
 
+  useEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
+
   const sendMessage = async (message: string) => {
     if (loading) return;
     setLoading(true);
@@ -169,7 +190,7 @@ const ChatWindow = () => {
       },
     ]);
 
-    const messageHandler = (e: MessageEvent) => {
+    const messageHandler = async (e: MessageEvent) => {
       const data = JSON.parse(e.data);
 
       if (data.type === 'error') {
@@ -231,8 +252,28 @@ const ChatWindow = () => {
           ['human', message],
           ['assistant', recievedMessage],
         ]);
+
         ws?.removeEventListener('message', messageHandler);
         setLoading(false);
+
+        const lastMsg = messagesRef.current[messagesRef.current.length - 1];
+
+        if (
+          lastMsg.role === 'assistant' &&
+          lastMsg.sources &&
+          lastMsg.sources.length > 0 &&
+          !lastMsg.suggestions
+        ) {
+          const suggestions = await getSuggestions(messagesRef.current);
+          setMessages((prev) =>
+            prev.map((msg) => {
+              if (msg.id === lastMsg.id) {
+                return { ...msg, suggestions: suggestions };
+              }
+              return msg;
+            }),
+          );
+        }
       }
     };
 
@@ -256,7 +297,14 @@ const ChatWindow = () => {
     sendMessage(message.content);
   };
 
-  return ws ? (
+  useEffect(() => {
+    if (isReady && initialMessage) {
+      sendMessage(initialMessage);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [isReady, initialMessage]);
+
+  return isReady ? (
     <div>
       {messages.length > 0 ? (
         <>