feat: Add start stream message function

tjtanjin · Nov 22, 2024 · e1b24ba · e1b24ba
1 parent ee6071e
commit e1b24ba
Show file tree

Hide file tree

Showing 10 changed files with 81 additions and 25 deletions.
diff --git a/__tests__/hooks/internal/useMessagesInternal.test.ts b/__tests__/hooks/internal/useMessagesInternal.test.ts
@@ -54,6 +54,7 @@ describe("useMessagesInternal", () => {
   it("should return expected functions and values", () => {
     const { result } = renderHook(() => useMessagesInternal());
 
+    expect(result.current).toHaveProperty("startStreamMessage");
     expect(result.current).toHaveProperty("endStreamMessage");
     expect(result.current).toHaveProperty("injectMessage");
     expect(result.current).toHaveProperty("removeMessage");
@@ -107,6 +108,18 @@ describe("useMessagesInternal", () => {
     expect(mockStreamMessageMap.current.has("BOT")).toBeTruthy();
   });
 
+  it("should start stream message correctly", async () => {
+    mockStreamMessageMap.current.set("BOT", "test-id");
+    const { result } = renderHook(() => useMessagesInternal());
+
+    await act(async () => {
+      const success = await result.current.startStreamMessage("BOT");
+      expect(success).toBeTruthy();
+    });
+
+    expect(mockStreamMessageMap.current.has("BOT")).toBeTruthy();
+  });
+
   it("should end stream message correctly", async () => {
     mockStreamMessageMap.current.set("BOT", "test-id");
     const { result } = renderHook(() => useMessagesInternal());

diff --git a/__tests__/services/BlockService/BlockService.test.tsx b/__tests__/services/BlockService/BlockService.test.tsx
@@ -59,6 +59,7 @@ describe("BlockService", () => {
 		setTextAreaValue: jest.fn(),
 		streamMessage: jest.fn(),
 		removeMessage: jest.fn(),
+		startStreamMessage: jest.fn(),
 		endStreamMessage: jest.fn(),
 		showToast: jest.fn(),
 		dismissToast: jest.fn(),

diff --git a/__tests__/services/BlockService/PathProcessor.test.ts b/__tests__/services/BlockService/PathProcessor.test.ts
@@ -15,6 +15,7 @@ const params: Params = {
 	injectMessage: jest.fn(),
 	streamMessage: jest.fn(),
 	removeMessage: jest.fn(),
+	startStreamMessage: jest.fn(),
 	endStreamMessage: jest.fn(),
 	showToast: jest.fn(),
 	dismissToast: jest.fn(),

diff --git a/__tests__/services/FunctionProcessor.test.ts b/__tests__/services/FunctionProcessor.test.ts
@@ -17,6 +17,7 @@ describe('processFunction', () => {
             injectMessage: jest.fn(),
             streamMessage: jest.fn(),
             removeMessage: jest.fn(),
+            startStreamMessage: jest.fn(),
             endStreamMessage: jest.fn(),
             showToast: jest.fn(),
             dismissToast: jest.fn(),

diff --git a/src/components/Buttons/FileAttachmentButton/FileAttachmentButton.tsx b/src/components/Buttons/FileAttachmentButton/FileAttachmentButton.tsx
@@ -28,7 +28,13 @@ const FileAttachmentButton = () => {
 	const { styles } = useStylesContext();
 
 	// handles messages
-	const { injectMessage, streamMessage, removeMessage, endStreamMessage } = useMessagesInternal();
+	const {
+		injectMessage,
+		streamMessage,
+		removeMessage,
+		endStreamMessage,
+		startStreamMessage,
+	} = useMessagesInternal();
 
 	// handles paths and blocks
 	const { getCurrPath, getPrevPath, goToPath, blockAllowsAttachment } = usePathsInternal();
@@ -124,7 +130,7 @@ const FileAttachmentButton = () => {
 			await handleSubmitText("📄 " + fileNames.join(", "), settings.fileAttachment?.sendFileName);
 			await fileHandler({userInput: inputRef.current?.value as string, prevPath: getPrevPath(),
 				currPath: getCurrPath(), goToPath, setTextAreaValue, injectMessage, streamMessage,
-				removeMessage, endStreamMessage, openChat, showToast, dismissToast, files
+				removeMessage, startStreamMessage, endStreamMessage, openChat, showToast, dismissToast, files
 			});
 		}
 	};

diff --git a/src/hooks/internal/useBotEffectsInternal.tsx b/src/hooks/internal/useBotEffectsInternal.tsx
@@ -37,6 +37,7 @@ export const useBotEffectsInternal = () => {
 
 	// handles messages
 	const {
+		startStreamMessage,
 		endStreamMessage,
 		injectMessage,
 		removeMessage,
@@ -234,8 +235,8 @@ export const useBotEffectsInternal = () => {
 		}
 
 		const params = {prevPath: getPrevPath(), currPath: getCurrPath(), goToPath, setTextAreaValue,
-			userInput: paramsInputRef.current, endStreamMessage, injectMessage, removeMessage, streamMessage,
-			openChat, showToast, dismissToast
+			userInput: paramsInputRef.current, startStreamMessage, endStreamMessage, injectMessage,
+			removeMessage, streamMessage, openChat, showToast, dismissToast
 		};
 
 		// calls the new block for preprocessing upon change to path.

diff --git a/src/hooks/internal/useMessagesInternal.ts b/src/hooks/internal/useMessagesInternal.ts
@@ -209,24 +209,8 @@ export const useMessagesInternal = () => {
 		sender = sender.toUpperCase();
 
 		if (!streamMessageMap.current.has(sender)) {
-			const message = createMessage(content, sender);
-			// handles start stream message event
-			if (settings.event?.rcbStartStreamMessage) {
-				const event = await callRcbEvent(RcbEvent.START_STREAM_MESSAGE, {message});
-				if (event.defaultPrevented) {
-					return null;
-				}
-			}
-
-			setIsBotTyping(false);
-			setMessages((prevMessages) => {
-				const updatedMessages = [...prevMessages, message];
-				handlePostMessagesUpdate(updatedMessages);
-				return [...prevMessages, message];
-			});
-			setUnreadCount(prev => prev + 1);
-			streamMessageMap.current.set(sender, message.id);
-			return message.id;
+			await startStreamMessage(sender);
+			return streamMessageMap.current.get(sender) ?? null;
 		}
 
 		const message = {...createMessage(content, sender), id: streamMessageMap.current.get(sender) as string};
@@ -254,7 +238,45 @@ export const useMessagesInternal = () => {
 	},[callRcbEvent, settings.event?.rcbChunkStreamMessage, settings.event?.rcbStartStreamMessage, streamMessageMap]);
 
 	/**
-	 * Sets the streaming mode of the chatbot.
+	 * Starts the stream for specified sender.
+	 * 
+	 * @param sender sender whose stream is being started
+	 * 
+	 * Note: This is not strictly needed, since there are auto-detections to start stream
+	 * of messages in place. However, it is recommended to specify this for code clarity and
+	 * for future-proofing.
+	 */
+	const startStreamMessage = useCallback(async (sender = "BOT"): Promise<boolean> => {
+		// always convert to uppercase for checks
+		sender = sender.toUpperCase();
+
+		// nothing to start if already streaming
+		if (streamMessageMap.current.has(sender)) {
+			return true;
+		}
+
+		// handles start stream message event
+		const message = createMessage("", sender);
+		if (settings.event?.rcbStartStreamMessage) {
+			const event = await callRcbEvent(RcbEvent.START_STREAM_MESSAGE, {message});
+			if (event.defaultPrevented) {
+				return false;
+			}
+		}
+
+		setIsBotTyping(false);
+		setMessages((prevMessages) => {
+			const updatedMessages = [...prevMessages, message];
+			handlePostMessagesUpdate(updatedMessages);
+			return [...prevMessages, message];
+		});
+		setUnreadCount(prev => prev + 1);
+		streamMessageMap.current.set(sender, message.id);
+		return true;
+	}, [callRcbEvent, messages, settings.event?.rcbStartStreamMessage, streamMessageMap])
+
+	/**
+	 * Ends the stream for specified sender.
 	 * 
 	 * @param sender sender whose stream is being ended
 	 * 
@@ -346,6 +368,7 @@ export const useMessagesInternal = () => {
 	}, [handlePostMessagesUpdate])
 
 	return {
+		startStreamMessage,
 		endStreamMessage,
 		injectMessage,
 		removeMessage,

diff --git a/src/hooks/internal/useSubmitInputInternal.ts b/src/hooks/internal/useSubmitInputInternal.ts
@@ -23,7 +23,13 @@ export const useSubmitInputInternal = () => {
 	const { settings } = useSettingsContext();
 
 	// handles messages
-	const { endStreamMessage, injectMessage, removeMessage, streamMessage } = useMessagesInternal();
+	const {
+		startStreamMessage,
+		endStreamMessage,
+		injectMessage,
+		removeMessage,
+		streamMessage
+	} = useMessagesInternal();
 
 	// handles paths
 	const { getCurrPath, getPrevPath, goToPath } = usePathsInternal();
@@ -143,7 +149,8 @@ export const useSubmitInputInternal = () => {
 
 		setTimeout(async () => {
 			const params = {prevPath: getPrevPath(), currPath: getCurrPath(), goToPath, setTextAreaValue, userInput, 
-				injectMessage, streamMessage, removeMessage, endStreamMessage, openChat, showToast, dismissToast
+				injectMessage, streamMessage, removeMessage, startStreamMessage, endStreamMessage, openChat,
+				showToast, dismissToast
 			};
 			const hasNextPath = await postProcessBlock(flowRef.current as Flow, path, params, goToPath);
 			if (!hasNextPath) {

diff --git a/src/hooks/useMessages.ts b/src/hooks/useMessages.ts
@@ -6,6 +6,7 @@ import { useMessagesInternal } from "./internal/useMessagesInternal";
 export const useMessages = () => {
 	// handles messages
 	const {
+		startStreamMessage,
 		endStreamMessage,
 		injectMessage,
 		removeMessage,
@@ -15,6 +16,7 @@ export const useMessages = () => {
 	} = useMessagesInternal();
 
 	return {
+		startStreamMessage,
 		endStreamMessage,
 		injectMessage,
 		removeMessage,

diff --git a/src/types/Params.ts b/src/types/Params.ts
@@ -13,6 +13,7 @@ export type Params = {
 	injectMessage: (content: string | JSX.Element, sender?: string) => Promise<string | null>;
 	streamMessage: (content: string | JSX.Element, sender?: string) => Promise<string | null>;
 	removeMessage: (id: string) => Promise<string | null>;
+	startStreamMessage: (sender: string) => Promise<boolean>;
 	endStreamMessage: (sender: string) => Promise<boolean>;
 	showToast: (content: string | JSX.Element, timeout?: number) => Promise<string | null>;
 	dismissToast: (id: string) => Promise<string | null>;