Make straming answers work more reliably.

components/Chat/components/Answer/SourceDocuments.tsx

@@ -1,10 +1,9 @@
-import AnswerCard from "@/components/Chat/components/Answer/Card";
 import React from "react";
-import { SourceDocument } from "@/components/Chat/types/chat";
+import { Work } from "@nulib/dcapi-types";
 import { styled } from "@/stitches.config";
 
 interface SourceDocumentsProps {
-  source_documents: SourceDocument[];
+  source_documents: Work[];
 }
 
 const SourceDocuments: React.FC<SourceDocumentsProps> = ({
@@ -13,7 +12,11 @@ const SourceDocuments: React.FC<SourceDocumentsProps> = ({
   return (
     <Sources>
       {source_documents.map((document, idx) => (
-        <AnswerCard {...document} key={idx} />
+        <div key={document.id}>
+          <strong>{document.title}</strong>
+          <img src={document.thumbnail} />
+        </div>
+        // <AnswerCard {...document} key={idx} />
       ))}
     </Sources>
   );

components/Chat/components/Wrapper.tsx

@@ -1,38 +1,16 @@
-import { useEffect, useState } from "react";
-
 import Chat from "@/components/Chat";
-import { ChatConfig } from "@/components/Chat/types/chat";
-import axios from "axios";
-
-const WS_ENDPOINTS = {
-  production: "https://api.dc.library.northwestern.edu/api/v2/chat-endpoint",
-  staging:
-    "https://dcapi.rdc-staging.library.northwestern.edu/api/v2/chat-endpoint",
-  weaviateEndpoint: `https://dcapi-prototype.rdc-staging.library.northwestern.edu/api/v2/chat-endpoint`,
-};
+import useChatSocket from "../../../hooks/useChatSocket";
+import useQueryParams from "@/hooks/useQueryParams";
 
 const ChatWrapper = () => {
-  const [chatConfig, setChatConfig] = useState<ChatConfig>();
-
-  useEffect(() => {
-    axios({
-      method: "GET",
-      url: WS_ENDPOINTS.production,
-      withCredentials: true,
-    })
-      .then((response) => {
-        setChatConfig(response.data);
-      })
-      .catch((error) => {
-        console.error(error);
-      });
-  }, []);
+  const { searchTerm: question } = useQueryParams();
+  const { authToken, chatSocket } = useChatSocket();
 
-  if (!chatConfig) return null;
+  if (!authToken || !chatSocket || !question) return null;
 
   return (
     <div style={{ background: "#f0f0f0", padding: "2rem" }}>
-      <Chat chatConfig={chatConfig} />
+      <Chat authToken={authToken} chatSocket={chatSocket} question={question} />
     </div>
   );
 };

components/Chat/index.tsx

@@ -1,92 +1,81 @@
-import { Answer, QuestionRendered, StreamingMessage } from "./types/chat";
-import React, { useCallback, useEffect, useState } from "react";
-import {
-  StyledActions,
-  StyledAnswerHeader,
-  StyledAnswerItem,
-  StyledRemoveButton,
-} from "./components/Answer/Answer.styled";
+import React, { useEffect, useState } from "react";
 
-import AnswerInformation from "./components/Answer/Information";
-import AnswerLoader from "./components/Answer/Loader";
-import { ChatConfig } from "@/components/Chat/types/chat";
-import Icon from "@/components/Shared/Icon";
-import { IconClear } from "@/components/Shared/SVG/Icons";
-import QuestionInput from "./components/Question/Input";
 import SourceDocuments from "./components/Answer/SourceDocuments";
 import StreamingAnswer from "./components/Answer/StreamingAnswer";
-import useQueryParams from "@/hooks/useQueryParams";
-import useStreamingAnswers from "./hooks/useStreamingAnswers";
-
-const Chat = ({ chatConfig }: { chatConfig: ChatConfig }) => {
-  console.log("\n RE RENDERING");
-  const { auth: authToken, endpoint } = chatConfig;
-  const { prepareQuestion, updateStreamAnswers } = useStreamingAnswers();
-
-  const [chatSocket, setChatSocket] = React.useState<WebSocket>();
-  const [readyState, setReadyState] = React.useState<WebSocket["readyState"]>();
-
+import { StreamingMessage } from "./types/chat";
+import { Work } from "@nulib/dcapi-types";
+import { prepareQuestion } from "../../lib/chat-helpers";
+
+const Chat = ({
+  authToken,
+  chatSocket,
+  question,
+}: {
+  authToken: string;
+  chatSocket?: WebSocket;
+  question?: string;
+}) => {
+  const [isReadyStateOpen, setIsReadyStateOpen] = useState(false);
+  const [isStreamingComplete, setIsStreamingComplete] = useState(false);
+  const [sourceDocuments, setSourceDocuments] = useState<Work[]>([]);
   const [streamedAnswer, setStreamedAnswer] = useState("");
 
-  const { searchTerm: question } = useQueryParams();
-
-  const handleReadyStateChange = (event: Event) => {
-    const target = event.target as WebSocket;
-    console.log("target.readyState", target.readyState);
-    setReadyState(target.readyState);
+  const handleReadyStateChange = () => {
+    setIsReadyStateOpen(chatSocket?.readyState === 1);
   };
 
   // Handle web socket stream updates
   const handleMessageUpdate = (event: MessageEvent) => {
     const data: StreamingMessage = JSON.parse(event.data);
-    console.log("handleMessageUpdate", data);
+    // console.log("handleMessageUpdate", data);
 
-    if (data.token) {
+    if (data.source_documents) {
+      setSourceDocuments(data.source_documents);
+    } else if (data.token) {
       setStreamedAnswer((prev) => {
         return prev + data.token;
       });
     } else if (data.answer) {
       setStreamedAnswer(data.answer);
+      setIsStreamingComplete(true);
     }
   };
 
   useEffect(() => {
-    if (question && chatSocket?.readyState === 1) {
+    if (question && isReadyStateOpen && chatSocket) {
       const preparedQuestion = prepareQuestion(question, authToken);
-      console.log("preparedQuestion", preparedQuestion, chatSocket);
       chatSocket?.send(JSON.stringify(preparedQuestion));
     }
-  }, [authToken, chatSocket, question, prepareQuestion]);
+  }, [chatSocket, isReadyStateOpen, prepareQuestion]);
 
   useEffect(() => {
-    if (!authToken || !endpoint) return;
-
-    const socket = new WebSocket(endpoint);
-
-    socket.addEventListener("open", handleReadyStateChange);
-    socket.addEventListener("close", handleReadyStateChange);
-    socket.addEventListener("error", handleReadyStateChange);
-    socket.addEventListener("message", handleMessageUpdate);
-
-    setChatSocket(socket);
+    if (chatSocket) {
+      chatSocket.addEventListener("message", handleMessageUpdate);
+      chatSocket.addEventListener("open", handleReadyStateChange);
+      chatSocket.addEventListener("close", handleReadyStateChange);
+      chatSocket.addEventListener("error", handleReadyStateChange);
+    }
 
     return () => {
-      if (socket) {
-        socket.close();
-        socket.removeEventListener("open", handleReadyStateChange);
-        socket.removeEventListener("close", handleReadyStateChange);
-        socket.removeEventListener("error", handleReadyStateChange);
-        socket.removeEventListener("message", handleMessageUpdate);
+      if (chatSocket) {
+        chatSocket.removeEventListener("message", handleMessageUpdate);
+        chatSocket.removeEventListener("open", handleReadyStateChange);
+        chatSocket.removeEventListener("close", handleReadyStateChange);
+        chatSocket.removeEventListener("error", handleReadyStateChange);
       }
     };
-  }, [authToken, endpoint]);
+  }, [chatSocket, chatSocket?.url]);
 
   return (
     <>
+      <h3>{question}</h3>
       {question && (
         <>
-          {/* <SourceDocuments source_documents={answer.source_documents} /> */}
-          <StreamingAnswer answer={streamedAnswer} isComplete={false} />
+          <SourceDocuments source_documents={sourceDocuments} />
+          <StreamingAnswer
+            answer={streamedAnswer}
+            isComplete={isStreamingComplete}
+          />
         </>
       )}
     </>

components/Chat/types/chat.ts

@@ -1,3 +1,5 @@
+import { Work } from "@nulib/dcapi-types";
+
 export type QuestionRendered = {
   question: string;
   ref: string;
@@ -10,26 +12,19 @@ export type Question = {
   ref: string;
 };
 
-export type SourceDocument = {
-  page_content: string;
-  metadata: {
-    [key: string]: any;
-  };
-};
-
 export type Answer = {
   answer: string;
   isComplete: boolean;
-  question?: string; // revisit this
+  question?: string;
   ref: string;
-  source_documents: Array<SourceDocument>;
+  source_documents: Array<Work>;
 };
 
 export type StreamingMessage = {
   answer?: string;
   question?: string;
   ref: string;
-  source_documents?: Array<SourceDocument>;
+  source_documents?: Array<Work>;
   token?: string;
 };
 

hooks/useChatSocket.ts

@@ -0,0 +1,38 @@
+import { useEffect, useState } from "react";
+
+import { DCAPI_CHAT_ENDPOINT } from "@/lib/constants/endpoints";
+import axios from "axios";
+
+const useChatSocket = () => {
+  const [chatSocket, setChatSocket] = useState<WebSocket | null>(null);
+  const [authToken, setAuthToken] = useState<string | null>(null);
+
+  useEffect(() => {
+    axios({
+      method: "GET",
+      url: DCAPI_CHAT_ENDPOINT,
+      withCredentials: true,
+    })
+      .then((response) => {
+        const { auth: authToken, endpoint } = response.data;
+
+        if (!authToken || !endpoint) return;
+
+        const socket = new WebSocket(endpoint);
+
+        setAuthToken(authToken);
+        setChatSocket(socket);
+
+        return () => {
+          if (socket) socket.close();
+        };
+      })
+      .catch((error) => {
+        console.error(error);
+      });
+  }, []);
+
+  return { authToken, chatSocket };
+};
+
+export default useChatSocket;