Agent support for @agent default agent inside workspace chat (#1093)

V1 of agent support via built-in `@agent` that can be invoked alongside normal workspace RAG chat.

frontend/src/components/WorkspaceChat/ChatContainer/index.jsx

@@ -2,16 +2,24 @@ import { useState, useEffect } from "react";
 import ChatHistory from "./ChatHistory";
 import PromptInput from "./PromptInput";
 import Workspace from "@/models/workspace";
-import handleChat from "@/utils/chat";
+import handleChat, { ABORT_STREAM_EVENT } from "@/utils/chat";
 import { isMobile } from "react-device-detect";
 import { SidebarMobileHeader } from "../../Sidebar";
 import { useParams } from "react-router-dom";
+import { v4 } from "uuid";
+import handleSocketResponse, {
+  websocketURI,
+  AGENT_SESSION_END,
+  AGENT_SESSION_START,
+} from "@/utils/chat/agent";
 
 export default function ChatContainer({ workspace, knownHistory = [] }) {
   const { threadSlug = null } = useParams();
   const [message, setMessage] = useState("");
   const [loadingResponse, setLoadingResponse] = useState(false);
   const [chatHistory, setChatHistory] = useState(knownHistory);
+  const [socketId, setSocketId] = useState(null);
+  const [websocket, setWebsocket] = useState(null);
   const handleMessageChange = (event) => {
     setMessage(event.target.value);
   };
@@ -68,6 +76,19 @@ export default function ChatContainer({ workspace, knownHistory = [] }) {
       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;
+      }
+
+      // TODO: Simplify this
       if (!promptMessage || !promptMessage?.userMessage) return false;
       if (!!threadSlug) {
         await Workspace.threads.streamChat(
@@ -79,7 +100,8 @@ export default function ChatContainer({ workspace, knownHistory = [] }) {
               setLoadingResponse,
               setChatHistory,
               remHistory,
-              _chatHistory
+              _chatHistory,
+              setSocketId
             )
         );
       } else {
@@ -92,7 +114,8 @@ export default function ChatContainer({ workspace, knownHistory = [] }) {
               setLoadingResponse,
               setChatHistory,
               remHistory,
-              _chatHistory
+              _chatHistory,
+              setSocketId
             )
         );
       }
@@ -101,6 +124,77 @@ export default function ChatContainer({ workspace, knownHistory = [] }) {
     loadingResponse === true && fetchReply();
   }, [loadingResponse, chatHistory, workspace]);
 
+  // TODO: Simplify this WSS stuff
+  useEffect(() => {
+    function handleWSS() {
+      try {
+        if (!socketId || !!websocket) return;
+        const socket = new WebSocket(
+          `${websocketURI()}/api/agent-invocation/${socketId}`
+        );
+
+        window.addEventListener(ABORT_STREAM_EVENT, () => {
+          window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+          websocket.close();
+        });
+
+        socket.addEventListener("message", (event) => {
+          setLoadingResponse(true);
+          try {
+            handleSocketResponse(event, setChatHistory);
+          } catch (e) {
+            console.error("Failed to parse data");
+            window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+            socket.close();
+          }
+          setLoadingResponse(false);
+        });
+
+        socket.addEventListener("close", (_event) => {
+          window.dispatchEvent(new CustomEvent(AGENT_SESSION_END));
+          setChatHistory((prev) => [
+            ...prev.filter((msg) => !!msg.content),
+            {
+              uuid: v4(),
+              type: "statusResponse",
+              content: "Agent session complete.",
+              role: "assistant",
+              sources: [],
+              closed: true,
+              error: null,
+              animate: false,
+              pending: false,
+            },
+          ]);
+          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,
+          },
+        ]);
+        setLoadingResponse(false);
+        setWebsocket(null);
+        setSocketId(null);
+      }
+    }
+    handleWSS();
+  }, [socketId]);
+
   return (
     <div
       style={{ height: isMobile ? "100%" : "calc(100% - 32px)" }}