diff --git a/src/components/ChatWindow/ChatContainer/index.jsx b/src/components/ChatWindow/ChatContainer/index.jsx index b74ac15..3fb3862 100644 --- a/src/components/ChatWindow/ChatContainer/index.jsx +++ b/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 (
diff --git a/src/components/ChatWindow/Header/index.jsx b/src/components/ChatWindow/Header/index.jsx index 2cc00fe..b23b745 100644 --- a/src/components/ChatWindow/Header/index.jsx +++ b/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" > -
+
{iconUrl diff --git a/src/hooks/useScriptAttributes.js b/src/hooks/useScriptAttributes.js index f9964de..3a16c11 100644 --- a/src/hooks/useScriptAttributes.js +++ b/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 diff --git a/src/utils/agent.js b/src/utils/agent.js new file mode 100644 index 0000000..16cccb0 --- /dev/null +++ b/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; +} \ No newline at end of file diff --git a/src/utils/chat/index.js b/src/utils/chat/index.js index 62d25cb..ffdba6f 100644 --- a/src/utils/chat/index.js +++ b/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); } } diff --git a/src/utils/request.js b/src/utils/request.js new file mode 100644 index 0000000..b4c48ec --- /dev/null +++ b/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; + } +} \ No newline at end of file