Improvements to chat UI (#7)

package.json

@@ -1,7 +1,7 @@
 {
   "name": "webgati-ai",
   "private": true,
-  "version": "0.1.6",
+  "version": "0.1.7",
   "scripts": {
     "dev": "vite",
     "build:dev": "vite build",

src/background/sw.ts

@@ -10,6 +10,7 @@ import {
   AppMessageUpdateModelId as SWMessageUpdateModel,
   AppMessageIndexWebpage,
   AppMessageTabStateInit,
+  QueryMode,
 } from "../utils/types";
 import { readAIModelConfig } from "../utils/storage";
 import {
@@ -188,6 +189,8 @@ function initTabState(tabId: number, url: string | null | undefined): TabState {
 }
 
 async function invokeBot(msg: AppMessageBotExecute, tabState: TabState) {
+  const queryMode = msg.payload.queryMode;
+
   try {
     postBotProcessing(tabState);
 
@@ -201,12 +204,12 @@ async function invokeBot(msg: AppMessageBotExecute, tabState: TabState) {
       msg,
       tabState.vectorStore,
       tabState.botAbortController,
-      (token) => postBotTokenResponse(tabState, token)
+      (token) => postBotTokenResponse(queryMode, tabState, token)
     );
   } catch (error: any) {
     if (error.message !== "AbortError") {
       console.log(error);
-      postBotError(tabState, error.message);
+      postBotError(queryMode, tabState, error.message);
     }
   } finally {
     postBotDone(tabState);
@@ -229,19 +232,25 @@ function postBotProcessing(tabState: TabState) {
   } as AppMessageBotProcessing);
 }
 
-function postBotTokenResponse(tabState: TabState, token: string) {
+function postBotTokenResponse(
+  queryMode: QueryMode,
+  tabState: TabState,
+  token: string
+) {
   tabState.port?.postMessage({
     type: "sw_bot-token-response",
     payload: {
+      queryMode,
       token,
     },
   } as AppMessageBotTokenResponse);
 }
 
-function postBotError(tabState: TabState, error: string) {
+function postBotError(queryMode: QueryMode, tabState: TabState, error: string) {
   tabState.port?.postMessage({
     type: "sw_bot-token-response",
     payload: {
+      queryMode,
       error,
     },
   } as AppMessageBotTokenResponse);

src/content-script/ContentScriptApp.tsx

@@ -5,6 +5,7 @@ import { PageSnipTool } from "./PageSnipTool";
 import { useContentScriptMessageListener } from "../utils/hooks/useContentScriptMessageListener";
 import { generatePageMarkdown } from "../utils/markdown";
 import {
+  AppMessageCheckSidePanelVisible,
   AppMessageGetWebpage,
   AppMessageImageCapture,
   AppMessageSelectionPrompt,
@@ -18,11 +19,26 @@ export function ContentScriptApp(): JSX.Element {
 
   const { selection } = useSelectionDialog(SELECTION_DEBOUNCE_DELAY_MS);
 
+  const checkAndInitSelectionDialog = useCallback(async () => {
+    try {
+      const result =
+        await chrome.runtime.sendMessage<AppMessageCheckSidePanelVisible>({
+          type: "cs_check-side-panel-visible",
+        });
+
+      if (result) {
+        setShowSelectionDialog(true);
+      }
+    } catch (error) {
+      // no-op
+    }
+  }, []);
+
   useEffect(() => {
     if (selection) {
-      setShowSelectionDialog(true);
+      checkAndInitSelectionDialog();
     }
-  }, [selection]);
+  }, [selection, checkAndInitSelectionDialog]);
 
   const handleSelectionDialogSubmit = useCallback((prompt: string) => {
     setShowSelectionDialog(false);

src/sidepanel/ChatUI.tsx

@@ -52,10 +52,11 @@ export const ChatUI = ({
 }: ChatUIProps): JSX.Element => {
   const formRef = useRef<HTMLFormElement>(null);
   const [showLoader, setShowLoader] = useState(false);
+  const [hasScrolled, setHasScrolled] = useState(false);
 
   const { scrollIntoView, scrollableRef, targetRef } =
     useScrollIntoView<HTMLDivElement>({
-      offset: 60,
+      offset: 100,
     });
 
   const form = useForm({
@@ -68,19 +69,6 @@ export const ChatUI = ({
     },
   });
 
-  useEffect(() => {
-    if (messages.length === 0) return;
-
-    setShowLoader(true);
-    scrollIntoView();
-  }, [messages, scrollIntoView]);
-
-  useEffect(() => {
-    if (showLoader) {
-      scrollIntoView();
-    }
-  }, [showLoader, scrollIntoView]);
-
   const handleEnterKey = (event: KeyboardEvent) => {
     event.preventDefault();
     form.validate();
@@ -89,6 +77,9 @@ export const ChatUI = ({
 
   useEffect(() => {
     setShowLoader(isLoading);
+    if (!isLoading) {
+      setHasScrolled(false);
+    }
   }, [isLoading]);
 
   useEffect(() => {
@@ -98,6 +89,7 @@ export const ChatUI = ({
   }, [error]);
 
   const handleFormSubmit = async (values: { question: string }) => {
+    setShowLoader(true);
     form.reset();
     await processUserPrompt(values.question.trim());
   };
@@ -107,6 +99,19 @@ export const ChatUI = ({
     clearImageData();
   };
 
+  const lastMessage = messages[messages.length - 1];
+
+  useEffect(() => {
+    if (
+      lastMessage &&
+      !hasScrolled &&
+      (lastMessage.role === "human" || lastMessage.queryMode === "summary")
+    ) {
+      scrollIntoView();
+      setHasScrolled(true);
+    }
+  }, [lastMessage, queryMode, hasScrolled, scrollIntoView]);
+
   return (
     <Box
       mt="8px"
@@ -125,7 +130,7 @@ export const ChatUI = ({
         sx={{ flex: 1, overflow: "scroll" }}
         ref={scrollableRef}
       >
-        {messages.map((message, index) => {
+        {messages.slice(0, -1).map((message, index) => {
           if (!message.content) {
             return null;
           }
@@ -137,12 +142,19 @@ export const ChatUI = ({
             />
           );
         })}
-        {showLoader && (
-          <Box sx={{ textAlign: "center" }} ref={targetRef}>
-            <Loader variant="dots" mb="8px" />
-          </Box>
+        {lastMessage && (
+          <Message
+            ref={targetRef}
+            role={lastMessage.role}
+            content={lastMessage.content}
+          />
         )}
       </Stack>
+      {showLoader && (
+        <Box sx={{ textAlign: "center" }}>
+          <Loader variant="dots" mb="8px" />
+        </Box>
+      )}
       <Box
         sx={{
           textAlign: "center",

src/sidepanel/SidePanel.tsx

@@ -63,38 +63,47 @@ export function SidePanel(): JSX.Element {
   const manifest = chrome.runtime.getManifest();
   const version = manifest.version;
 
-  const processToken = useCallback((token: string, isDone: boolean) => {
-    setMessages((messages) => {
-      const lastMessage = messages[messages.length - 1];
-      const prevMessages = messages.slice(0, messages.length - 1);
-
-      if (!lastMessage || lastMessage.isComplete || lastMessage.role !== "ai") {
+  const processToken = useCallback(
+    (queryMode: QueryMode | null, token: string, isDone: boolean) => {
+      setMessages((messages) => {
+        const lastMessage = messages[messages.length - 1];
+        const prevMessages = messages.slice(0, messages.length - 1);
+
+        if (
+          !lastMessage ||
+          lastMessage.isComplete ||
+          lastMessage.role !== "ai"
+        ) {
+          return [
+            ...messages,
+            {
+              role: "ai",
+              queryMode: queryMode!,
+              content: token,
+              isComplete: isDone,
+            },
+          ];
+        }
         return [
-          ...messages,
+          ...prevMessages,
           {
             role: "ai",
-            content: token,
+            queryMode: lastMessage.queryMode,
+            content: lastMessage.content + token,
             isComplete: isDone,
           },
         ];
-      }
-      return [
-        ...prevMessages,
-        {
-          role: "ai",
-          content: lastMessage.content + token,
-          isComplete: isDone,
-        },
-      ];
-    });
-  }, []);
+      });
+    },
+    []
+  );
 
   const handleBotMessagePayload = useCallback(
     (payload: AppMessageBotTokenResponse["payload"], isDone: boolean) => {
       if (payload.error) {
         setError(payload.error);
       } else {
-        processToken(payload.token, isDone);
+        processToken(payload.queryMode, payload.token, isDone);
       }
     },
     [processToken]
@@ -159,7 +168,7 @@ export function SidePanel(): JSX.Element {
       const prevMessages: ChatMessage[] = [];
       setMessages((messages) => {
         prevMessages.push(...messages);
-        return [...messages, { role: "human", content: prompt }];
+        return [...messages, { role: "human", queryMode, content: prompt }];
       });
 
       if (queryMode === "webpage-text-qa") {

src/utils/hooks/useChatMessaging.ts

@@ -28,7 +28,7 @@ export function useChatMessaging(
           onMessage(msg.payload, false);
           break;
         case "sw_bot-done":
-          onMessage({ token: "" }, true);
+          onMessage({ queryMode: null, token: "" }, true);
           setIsBotProcessing(false);
           break;
         default:

src/utils/hooks/useSidePanelMessageListener.ts

@@ -38,6 +38,9 @@ export function useSidePanelMessageListener(
               onImageCapture(message.payload.imageData);
               sendResponse("OK");
               break;
+            case "cs_check-side-panel-visible":
+              sendResponse(true);
+              break;
           }
         }
       }

src/utils/types.ts

@@ -35,6 +35,7 @@ export type IndexedData = {
 
 export type ChatMessage = {
   role: "human" | "ai";
+  queryMode: QueryMode;
   content: string;
   isComplete?: boolean;
 };
@@ -82,6 +83,10 @@ export type AppMessageSelectionPrompt = {
   };
 };
 
+export type AppMessageCheckSidePanelVisible = {
+  type: "cs_check-side-panel-visible";
+};
+
 export type AppMessageImageCapture = {
   type: "cs_image-capture";
   payload: {
@@ -109,6 +114,7 @@ export type AppMessageBotProcessing = {
 export type AppMessageBotTokenResponse = {
   type: "sw_bot-token-response";
   payload: {
+    queryMode: QueryMode | null;
     token: string;
     error?: string;
   };
@@ -183,7 +189,8 @@ export type AppMessage =
   | AppMessageSidePanelInit
   | AppMessageStartPageSnipTool
   | AppMessageSelectionPrompt
-  | AppMessageImageCapture;
+  | AppMessageImageCapture
+  | AppMessageCheckSidePanelVisible;
 
 export type TabState = {
   tabId: number;