Handle chat response 'end' message

components/Chat/Chat.test.tsx

@@ -21,6 +21,7 @@ jest.mock("@/context/search-context", () => {
         chat: {
           answer: "",
           documents: [],
+          end: "stop",
           question: "",
         },
         searchFixed: false,
@@ -45,7 +46,7 @@ jest.mock("@/hooks/useChatSocket");
 (useChatSocket as jest.Mock).mockImplementation(() => ({
   authToken: "fake-token-1",
   isConnected: false,
-  message: { answer: "fake-answer-1" },
+  message: { answer: "fake-answer-1", end: "stop" },
   sendMessage: mockSendMessage,
 }));
 
@@ -120,4 +121,54 @@ describe("Chat component", () => {
 
     expect(mockSendMessage).not.toHaveBeenCalled();
   });
+
+  it("displays an error message when the response hits the LLM token limit", () => {
+    (useChatSocket as jest.Mock).mockImplementation(() => ({
+      authToken: "fake",
+      isConnected: true,
+      message: {
+        end: {
+          reason: "length",
+          ref: "fake",
+        },
+      },
+      sendMessage: mockSendMessage,
+    }));
+
+    mockRouter.setCurrentUrl("/search?q=boats");
+
+    render(
+      <SearchProvider>
+        <Chat />
+      </SearchProvider>,
+    );
+
+    const error = screen.getByText("The response has hit the LLM token limit.");
+    expect(error).toBeInTheDocument();
+  });
+
+  it("displays an error message when the response times out", () => {
+    (useChatSocket as jest.Mock).mockImplementation(() => ({
+      authToken: "fake",
+      isConnected: true,
+      message: {
+        end: {
+          reason: "timeout",
+          ref: "fake",
+        },
+      },
+      sendMessage: mockSendMessage,
+    }));
+
+    mockRouter.setCurrentUrl("/search?q=boats");
+
+    render(
+      <SearchProvider>
+        <Chat />
+      </SearchProvider>,
+    );
+
+    const error = screen.getByText("The response has timed out.");
+    expect(error).toBeInTheDocument();
+  });
 });

components/Chat/Chat.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useState } from "react";
 
+import Announcement from "@/components/Shared/Announcement";
 import { Button } from "@nulib/design-system";
 import ChatFeedback from "@/components/Chat/Feedback/Feedback";
 import ChatResponse from "@/components/Chat/Response/Response";
@@ -19,6 +20,8 @@ const Chat = ({ totalResults }: { totalResults?: number }) => {
   const { chat } = searchState;
   const { answer, documents = [], question } = chat;
 
+  const [streamingError, setStreamingError] = useState("");
+
   const sameQuestionExists = !!question && searchTerm === question;
 
   const [sourceDocuments, setSourceDocuments] = useState<Work[]>([]);
@@ -63,6 +66,22 @@ const Chat = ({ totalResults }: { totalResults?: number }) => {
       return;
     }
 
+    if (message.end) {
+      switch (message.end.reason) {
+        case "length":
+          setStreamingError("The response has hit the LLM token limit.");
+          break;
+        case "timeout":
+          setStreamingError("The response has timed out.");
+          break;
+        case "eos_token":
+          setStreamingError("This should never happen.");
+          break;
+        default:
+          break;
+      }
+    }
+
     if (message.answer) {
       updateChat();
     }
@@ -94,6 +113,13 @@ const Chat = ({ totalResults }: { totalResults?: number }) => {
         sourceDocuments={sameQuestionExists ? documents : sourceDocuments}
         streamedAnswer={sameQuestionExists ? answer : streamedAnswer}
       />
+      {streamingError && (
+        <Container>
+          <Announcement css={{ marginTop: "1rem" }}>
+            {streamingError}
+          </Announcement>
+        </Container>
+      )}
       {answer && (
         <>
           <Container>

types/components/chat.ts

@@ -22,6 +22,10 @@ export type Answer = {
 
 export type StreamingMessage = {
   answer?: string;
+  end?: {
+    reason: "stop" | "length" | "timeout" | "eos_token";
+    ref: string;
+  };
   question?: string;
   ref: string;
   source_documents?: Array<Work>;