From 7d4007b86ef4f5a7ba91c9f9ec403843895bea0e Mon Sep 17 00:00:00 2001 From: Mat Jordan Date: Fri, 1 Mar 2024 13:55:58 -0500 Subject: [PATCH] Make straming answers work more reliably. --- .../components/Answer/SourceDocuments.tsx | 11 ++- components/Chat/components/Wrapper.tsx | 34 ++----- components/Chat/hooks/useStreamingAnswers.tsx | 90 ----------------- components/Chat/index.tsx | 99 +++++++++---------- components/Chat/types/chat.ts | 15 +-- hooks/useChatSocket.ts | 38 +++++++ lib/chat-helpers.ts | 31 ++++++ lib/constants/endpoints.ts | 3 + 8 files changed, 134 insertions(+), 187 deletions(-) delete mode 100644 components/Chat/hooks/useStreamingAnswers.tsx create mode 100644 hooks/useChatSocket.ts create mode 100644 lib/chat-helpers.ts diff --git a/components/Chat/components/Answer/SourceDocuments.tsx b/components/Chat/components/Answer/SourceDocuments.tsx index 79fc8b74..8fb7a027 100644 --- a/components/Chat/components/Answer/SourceDocuments.tsx +++ b/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 = ({ @@ -13,7 +12,11 @@ const SourceDocuments: React.FC = ({ return ( {source_documents.map((document, idx) => ( - +
+ {document.title} + +
+ // ))} ); diff --git a/components/Chat/components/Wrapper.tsx b/components/Chat/components/Wrapper.tsx index 20978bf3..6b2c81e0 100644 --- a/components/Chat/components/Wrapper.tsx +++ b/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(); - - 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 (
- +
); }; diff --git a/components/Chat/hooks/useStreamingAnswers.tsx b/components/Chat/hooks/useStreamingAnswers.tsx deleted file mode 100644 index 1b77a99c..00000000 --- a/components/Chat/hooks/useStreamingAnswers.tsx +++ /dev/null @@ -1,90 +0,0 @@ -import { Answer, Question, StreamingMessage } from "../types/chat"; - -const updateStreamAnswers = ( - data: StreamingMessage, - updatedStreamAnswers: Answer[] -) => { - // Check if the answer with the given 'ref' already exists in the state - const answerIndex = updatedStreamAnswers.findIndex( - (answer) => answer.ref === data.ref - ); - const existingAnswer = updatedStreamAnswers[answerIndex]; - - let updatedAnswer: Answer; - if (existingAnswer) { - // Create a shallow copy of the existing answer to modify - updatedAnswer = { ...existingAnswer }; - } else { - // Initialize a new answer - updatedAnswer = { - answer: "", - isComplete: false, - question: data.question, - ref: data.ref, - source_documents: [], - }; - } - - // Update the properties of the answer based on the incoming data - if (data.token) { - updatedAnswer.answer += data.token; - } - if (data.source_documents) { - updatedAnswer.source_documents = data.source_documents; - } - if (data.answer) { - updatedAnswer.answer = data.answer; - updatedAnswer.isComplete = true; - } - - // Replace or append the answer in the state array - if (existingAnswer) { - updatedStreamAnswers[answerIndex] = updatedAnswer; - } else { - /** - * save the question in local storage - */ - - // questions.unshift({ id: ref, question: questionString, timestamp }); - // saveQuestions(questions); - - // Update the state with the modified array - updatedStreamAnswers.push(updatedAnswer); - } - - return updatedStreamAnswers; -}; - -const prepareQuestion = (questionString: string, authToken: string) => { - const date = new Date(); - const timestamp = date.getTime(); - - /** - * hackily generate unique id from string and timestamp - */ - const uniqueString = `${questionString}${timestamp}`; - - // Refactor the following as a SHA1[0..4] - const ref = uniqueString - .split("") - .reduce((a, b) => { - a = (a << 5) - a + b.charCodeAt(0); - return a & a; - }, 0) - .toString(); - - const question: Question = { - auth: authToken, - message: "chat", - question: questionString, - ref, - }; - - return question; -}; - -const useStreamingAnswers = () => { - return { prepareQuestion, updateStreamAnswers }; -}; - -export default useStreamingAnswers; diff --git a/components/Chat/index.tsx b/components/Chat/index.tsx index 3ef78f1c..52a6d7c2 100644 --- a/components/Chat/index.tsx +++ b/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(); - const [readyState, setReadyState] = React.useState(); - +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([]); 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 ( <> +

{question}

{question && ( <> - {/* */} - + + )} diff --git a/components/Chat/types/chat.ts b/components/Chat/types/chat.ts index 0d3c6a59..eafc26dd 100644 --- a/components/Chat/types/chat.ts +++ b/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; + source_documents: Array; }; export type StreamingMessage = { answer?: string; question?: string; ref: string; - source_documents?: Array; + source_documents?: Array; token?: string; }; diff --git a/hooks/useChatSocket.ts b/hooks/useChatSocket.ts new file mode 100644 index 00000000..25184f6d --- /dev/null +++ b/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(null); + const [authToken, setAuthToken] = useState(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; diff --git a/lib/chat-helpers.ts b/lib/chat-helpers.ts new file mode 100644 index 00000000..0c6382c1 --- /dev/null +++ b/lib/chat-helpers.ts @@ -0,0 +1,31 @@ +import { Question } from "../components/Chat/types/chat"; + +const prepareQuestion = (questionString: string, authToken: string) => { + const date = new Date(); + const timestamp = date.getTime(); + + /** + * hackily generate unique id from string and timestamp + */ + const uniqueString = `${questionString}${timestamp}`; + + // Refactor the following as a SHA1[0..4] + const ref = uniqueString + .split("") + .reduce((a, b) => { + a = (a << 5) - a + b.charCodeAt(0); + return a & a; + }, 0) + .toString(); + + const question: Question = { + auth: authToken, + message: "chat", + question: questionString, + ref, + }; + + return question; +}; + +export { prepareQuestion }; diff --git a/lib/constants/endpoints.ts b/lib/constants/endpoints.ts index b8422866..aebbc562 100644 --- a/lib/constants/endpoints.ts +++ b/lib/constants/endpoints.ts @@ -1,3 +1,5 @@ +const DCAPI_CHAT_ENDPOINT = + "https://api.dc.library.northwestern.edu/api/v2/chat-endpoint"; const DCAPI_ENDPOINT = process.env.NEXT_PUBLIC_DCAPI_ENDPOINT; const DCAPI_PRODUCTION_ENDPOINT = "https://api.dc.library.northwestern.edu/api/v2"; @@ -9,6 +11,7 @@ const PRODUCTION_URL = "https://digitalcollections.library.northwestern.edu"; export { DC_URL, + DCAPI_CHAT_ENDPOINT, DCAPI_ENDPOINT, DCAPI_PRODUCTION_ENDPOINT, DC_API_SEARCH_URL,