Enable agents support in embed chat

src/components/ChatWindow/ChatContainer/index.jsx

@@ -1,8 +1,14 @@
 import React, { useState, useEffect } from "react";
 import ChatHistory from "./ChatHistory";
 import PromptInput from "./PromptInput";
-import handleChat from "@/utils/chat";
+import handleChat, { ABORT_STREAM_EVENT } from "@/utils/chat";
 import ChatService from "@/models/chatService";
+import handleSocketResponse, {
+  websocketURI,
+  AGENT_SESSION_END,
+  AGENT_SESSION_START,
+} from "@/utils/agent";
+import { v4 } from "uuid";
 export const SEND_TEXT_EVENT = "anythingllm-embed-send-prompt";
 
 export default function ChatContainer({
@@ -13,6 +19,8 @@ export default function ChatContainer({
   const [message, setMessage] = useState("");
   const [loadingResponse, setLoadingResponse] = useState(false);
   const [chatHistory, setChatHistory] = useState(knownHistory);
+  const [socketId, setSocketId] = useState(null);
+  const [websocket, setWebsocket] = useState(null);
 
   // Resync history if the ref to known history changes
   // eg: cleared.
@@ -93,6 +101,18 @@ export default function ChatContainer({
       const remHistory = chatHistory.length > 0 ? chatHistory.slice(0, -1) : [];
       var _chatHistory = [...remHistory];
 
+      // Override hook for new messages to now go to agents until the connection closes
+      if (!!websocket) {
+        if (!promptMessage || !promptMessage?.userMessage) return false;
+        websocket.send(
+          JSON.stringify({
+            type: "awaitingFeedback",
+            feedback: promptMessage?.userMessage,
+          })
+        );
+        return;
+      }
+
       if (!promptMessage || !promptMessage?.userMessage) {
         setLoadingResponse(false);
         return false;
@@ -108,14 +128,15 @@ export default function ChatContainer({
             setLoadingResponse,
             setChatHistory,
             remHistory,
-            _chatHistory
+            _chatHistory,
+            setSocketId
           )
       );
       return;
     }
 
     loadingResponse === true && fetchReply();
-  }, [loadingResponse, chatHistory]);
+  }, [loadingResponse, chatHistory, websocket]);
 
   const handleAutofillEvent = (event) => {
     if (!event.detail.command) return;
@@ -129,6 +150,65 @@ export default function ChatContainer({
     };
   }, []);
 
+  // Websocket connection management for agent sessions
+  useEffect(() => {
+    function handleWSS() {
+      try {
+        if (!socketId || !!websocket) return;
+        const socket = new WebSocket(
+          `${websocketURI(settings)}/api/agent-invocation/${socketId}`
+        );
+
+        window.addEventListener(ABORT_STREAM_EVENT, () => {
+          window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+          if (websocket) websocket.close();
+        });
+
+        socket.addEventListener("message", (event) => {
+          setLoadingResponse(true);
+          try {
+            handleSocketResponse(event, setChatHistory);
+          } catch (e) {
+            console.error("Failed to parse agent data:", e);
+            window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+            socket.close();
+          }
+          setLoadingResponse(false);
+        });
+
+        socket.addEventListener("close", (_event) => {
+          window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+          setLoadingResponse(false);
+          setWebsocket(null);
+          setSocketId(null);
+        });
+
+        setWebsocket(socket);
+        window.dispatchEvent(new CustomEvent(AGENT_SESSION_START));
+      } catch (e) {
+        setChatHistory((prev) => [
+          ...prev.filter((msg) => !!msg.content),
+          {
+            uuid: v4(),
+            type: "abort",
+            content: e.message,
+            role: "assistant",
+            sources: [],
+            closed: true,
+            error: e.message,
+            animate: false,
+            pending: false,
+            sentAt: Math.floor(Date.now() / 1000),
+          },
+        ]);
+        setLoadingResponse(false);
+        setWebsocket(null);
+        setSocketId(null);
+      }
+    }
+    handleWSS();
+  }, [socketId]);
+
   return (
     <div className="allm-h-full allm-w-full allm-flex allm-flex-col">
       <div className="allm-flex-1 allm-min-h-0 allm-mb-8">

src/components/ChatWindow/Header/index.jsx

@@ -49,9 +49,14 @@ export default function ChatWindowHeader({
       className="allm-flex allm-items-center allm-relative allm-rounded-t-2xl"
       id="anything-llm-header"
     >
-      <div className="allm-flex allm-justify-center allm-items-center allm-w-full allm-h-[76px]">
+      <div
+        className="allm-flex allm-justify-center allm-items-center allm-w-full allm-h-[76px]"
+        style={{
+            backgroundColor: settings.topHeaderBgColor
+        }}
+      >
         <img
-          style={{ maxWidth: 48, maxHeight: 48 }}
+          style={{ maxWidth: 48, maxHeight: 48}}
           src={iconUrl ?? AnythingLLMIcon}
           alt={iconUrl ? "Brand" : "AnythingLLM Logo"}
         />

src/hooks/useScriptAttributes.js

@@ -18,6 +18,7 @@ const DEFAULT_SETTINGS = {
   buttonColor: "#262626", // must be hex color code
   userBgColor: "#2C2F35", // user text bubble color
   assistantBgColor: "#2563eb", // assistant text bubble color
+  topHeaderBgColor: "#2563eb",
   noSponsor: null, // Shows sponsor in footer of chat
   sponsorText: "Powered by AnythingLLM", // default sponsor text
   sponsorLink: "https://anythingllm.com", // default sponsor link

src/utils/agent.js

@@ -0,0 +1,135 @@
+import { v4 } from "uuid";
+import { safeJsonParse } from "../utils/request";
+import { useState, useEffect } from "react";
+
+export const AGENT_SESSION_START = "agentSessionStart";
+export const AGENT_SESSION_END = "agentSessionEnd";
+
+const handledEvents = [
+  "statusResponse",
+  "fileDownload",
+  "awaitingFeedback",
+  "wssFailure",
+  "rechartVisualize",
+];
+
+export function websocketURI(embedSettings) {
+  const { baseApiUrl } = embedSettings;
+  const wsProtocol = baseApiUrl.startsWith("https://") ? "wss:" : "ws:";
+  return `${wsProtocol}//${new URL(baseApiUrl).host}`;
+}
+
+export default function handleSocketResponse(event, setChatHistory) {
+  const data = safeJsonParse(event.data, null);
+  if (data === null) return;
+  if (data.type === 'statusResponse') return;
+  // No message type is defined then this is a generic message
+  // that we need to print to the user as a system response
+  if (!data.hasOwnProperty("type")) {
+    return setChatHistory((prev) => {
+      return [
+        ...prev.filter((msg) => !!msg.content),
+        {
+          uuid: v4(),
+          content: data.content,
+          role: "assistant",
+          sources: [],
+          closed: true,
+          error: null,
+          animate: false,
+          pending: false,
+          sentAt: Math.floor(Date.now() / 1000),
+        },
+      ];
+    });
+  }
+
+  if (!handledEvents.includes(data.type) || !data.content) return;
+
+  if (data.type === "fileDownload") {
+    // File download functionality for embed
+    const blob = new Blob([atob(data.content.b64Content)], { type: 'application/octet-stream' });
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement('a');
+    a.href = url;
+    a.download = data.content.filename ?? 'unknown.txt';
+    document.body.appendChild(a);
+    a.click();
+    document.body.removeChild(a);
+    URL.revokeObjectURL(url);
+    return;
+  }
+
+  if (data.type === "rechartVisualize") {
+    return setChatHistory((prev) => {
+      return [
+        ...prev.filter((msg) => !!msg.content),
+        {
+          type: "rechartVisualize",
+          uuid: v4(),
+          content: data.content,
+          role: "assistant",
+          sources: [],
+          closed: true,
+          error: null,
+          animate: false,
+          pending: false,
+          sentAt: Math.floor(Date.now() / 1000),
+        },
+      ];
+    });
+  }
+
+  if (data.type === "wssFailure") {
+    return setChatHistory((prev) => {
+      return [
+        ...prev.filter((msg) => !!msg.content),
+        {
+          uuid: v4(),
+          content: data.content,
+          role: "assistant",
+          sources: [],
+          closed: true,
+          error: data.content,
+          animate: false,
+          pending: false,
+          sentAt: Math.floor(Date.now() / 1000),
+        },
+      ];
+    });
+  }
+
+  return setChatHistory((prev) => {
+    return [
+      ...prev.filter((msg) => !!msg.content),
+      {
+        uuid: v4(),
+        type: data.type,
+        content: data.content,
+        role: "assistant",
+        sources: [],
+        closed: true,
+        error: null,
+        animate: data?.animate || false,
+        pending: false,
+        sentAt: Math.floor(Date.now() / 1000),
+      },
+    ];
+  });
+}
+
+export function useIsAgentSessionActive() {
+  const [activeSession, setActiveSession] = useState(false);
+  useEffect(() => {
+    function listenForAgentSession() {
+      if (!window) return;
+      window.addEventListener(AGENT_SESSION_START, () =>
+        setActiveSession(true)
+      );
+      window.addEventListener(AGENT_SESSION_END, () => setActiveSession(false));
+    }
+    listenForAgentSession();
+  }, []);
+
+  return activeSession;
+}

src/utils/chat/index.js

@@ -1,10 +1,13 @@
+export const ABORT_STREAM_EVENT = "abort-chat-stream";
+
 // For handling of synchronous chats that are not utilizing streaming or chat requests.
 export default function handleChat(
   chatResult,
   setLoadingResponse,
   setChatHistory,
   remHistory,
-  _chatHistory
+  _chatHistory,
+  setSocketId = null
 ) {
   const {
     uuid,
@@ -14,6 +17,7 @@ export default function handleChat(
     error,
     close,
     errorMsg = null,
+    websocketUUID = null,
   } = chatResult;
 
   // Preserve the sentAt from the last message in the chat history
@@ -109,6 +113,10 @@ export default function handleChat(
       });
     }
     setChatHistory([..._chatHistory]);
+  } else if (type === "agentInitWebsocketConnection" && setSocketId) {
+    setSocketId(websocketUUID);
+  } else if (type === "statusResponse") {
+    setLoadingResponse(false);
   }
 }
 

src/utils/request.js

@@ -0,0 +1,8 @@
+export function safeJsonParse(jsonString, defaultValue) {
+  try {
+    return JSON.parse(jsonString);
+  } catch (error) {
+    console.error("Failed to parse JSON:", error);
+    return defaultValue;
+  }
+}