feat: add polling for pending notifications (#1152)

* feat: add notification controller

* feat: add polling logic on pending notifications

* feat: refecth notifications on Feed

backend/main.py

@@ -19,6 +19,7 @@ from routes.chat_routes import chat_router
 from routes.crawl_routes import crawl_router
 from routes.explore_routes import explore_router
 from routes.misc_routes import misc_router
+from routes.notification_routes import notification_router
 from routes.prompt_routes import prompt_router
 from routes.subscription_routes import subscription_router
 from routes.upload_routes import upload_router
@@ -54,6 +55,7 @@ app.include_router(user_router)
 app.include_router(api_key_router)
 app.include_router(subscription_router)
 app.include_router(prompt_router)
+app.include_router(notification_router)
 
 
 @app.exception_handler(HTTPException)

backend/routes/notification_routes.py

@@ -0,0 +1,24 @@
+from uuid import UUID
+
+from auth import AuthBearer
+from fastapi import APIRouter, Depends
+from repository.notification.get_chat_notifications import (
+    get_chat_notifications,
+)
+
+notification_router = APIRouter()
+
+
+@notification_router.get(
+    "/notifications/{chat_id}",
+    dependencies=[Depends(AuthBearer())],
+    tags=["Notification"],
+)
+async def get_notifications(
+    chat_id: UUID,
+):
+    """
+    Get notifications by chat_id
+    """
+
+    return get_chat_notifications(chat_id)

frontend/app/chat/[chatId]/components/ActionsBar/ActionsBar.tsx

@@ -10,9 +10,12 @@ export const ActionsBar = (): JSX.Element => {
     shouldDisplayUploadCard,
     setShouldDisplayUploadCard,
     hasPendingRequests,
+    setHasPendingRequests,
   } = useActionBar();
   const { addContent, contents, feedBrain, removeContent } =
-    useKnowledgeUploader();
+    useKnowledgeUploader({
+      setHasPendingRequests,
+    });
 
   const { t } = useTranslation(["chat"]);
 

frontend/app/chat/[chatId]/components/ActionsBar/hooks/useActionBar.ts

@@ -19,5 +19,6 @@ export const useActionBar = () => {
     shouldDisplayUploadCard,
     setShouldDisplayUploadCard,
     hasPendingRequests,
+    setHasPendingRequests,
   };
 };

frontend/app/chat/[chatId]/components/ActionsBar/hooks/useKnowledgeUploader.ts

@@ -1,31 +1,40 @@
 /* eslint-disable max-lines */
 import axios from "axios";
 import { UUID } from "crypto";
-import { useParams } from "next/navigation";
+import { useParams, useRouter } from "next/navigation";
 import { useCallback, useState } from "react";
 import { useTranslation } from "react-i18next";
 
 import { useChatApi } from "@/lib/api/chat/useChatApi";
 import { useCrawlApi } from "@/lib/api/crawl/useCrawlApi";
+import { useNotificationApi } from "@/lib/api/notification/useNotificationApi";
 import { useUploadApi } from "@/lib/api/upload/useUploadApi";
+import { useChatContext } from "@/lib/context";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
 import { useToast } from "@/lib/hooks";
 
 import { FeedItemCrawlType, FeedItemType, FeedItemUploadType } from "../types";
 
+type UseKnowledgeUploaderProps = {
+  setHasPendingRequests: (hasPendingRequests: boolean) => void;
+};
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
-export const useKnowledgeUploader = () => {
+export const useKnowledgeUploader = ({
+  setHasPendingRequests,
+}: UseKnowledgeUploaderProps) => {
   const [contents, setContents] = useState<FeedItemType[]>([]);
   const { publish } = useToast();
   const { uploadFile } = useUploadApi();
   const { t } = useTranslation(["upload"]);
   const { crawlWebsiteUrl } = useCrawlApi();
   const { createChat } = useChatApi();
-
+  const { currentBrainId } = useBrainContext();
+  const { setNotifications } = useChatContext();
+  const { getChatNotifications } = useNotificationApi();
+  const router = useRouter();
   const params = useParams();
   const chatId = params?.chatId as UUID | undefined;
 
-  const { currentBrainId } = useBrainContext();
   const addContent = (content: FeedItemType) => {
     setContents((prevContents) => [...prevContents, content]);
   };
@@ -33,6 +42,11 @@ export const useKnowledgeUploader = () => {
     setContents((prevContents) => prevContents.filter((_, i) => i !== index));
   };
 
+  const fetchNotifications = async (currentChatId: UUID): Promise<void> => {
+    const fetchedNotifications = await getChatNotifications(currentChatId);
+    setNotifications(fetchedNotifications);
+  };
+
   const crawlWebsiteHandler = useCallback(
     async (url: string, brainId: UUID, chat_id: UUID) => {
       // Configure parameters
@@ -50,6 +64,7 @@ export const useKnowledgeUploader = () => {
           config,
           chat_id,
         });
+        await fetchNotifications(chat_id);
       } catch (error: unknown) {
         publish({
           variant: "danger",
@@ -114,6 +129,7 @@ export const useKnowledgeUploader = () => {
     }
 
     try {
+      setHasPendingRequests(true);
       const currentChatId = chatId ?? (await createChat("New Chat")).chat_id;
       const uploadPromises = files.map((file) =>
         uploadFileHandler(file, currentBrainId, currentChatId)
@@ -126,6 +142,12 @@ export const useKnowledgeUploader = () => {
 
       setContents([]);
 
+      if (chatId === undefined) {
+        void router.push(`/chat/${currentChatId}`);
+      } else {
+        await fetchNotifications(currentChatId);
+      }
+
       publish({
         variant: "success",
         text: t("knowledgeUploaded"),
@@ -135,6 +157,8 @@ export const useKnowledgeUploader = () => {
         variant: "danger",
         text: JSON.stringify(e),
       });
+    } finally {
+      setHasPendingRequests(false);
     }
   };
 

frontend/app/chat/[chatId]/hooks/useSelectedChatPage.ts

@@ -1,20 +1,59 @@
+import { useQuery } from "@tanstack/react-query";
 import { useParams } from "next/navigation";
 import { useEffect } from "react";
 
 import { useChatApi } from "@/lib/api/chat/useChatApi";
+import { useNotificationApi } from "@/lib/api/notification/useNotificationApi";
 import { useChatContext } from "@/lib/context";
 
+import { getChatNotificationsQueryKey } from "../utils/getChatNotificationsQueryKey";
 import { getMessagesFromChatItems } from "../utils/getMessagesFromChatItems";
 import { getNotificationsFromChatItems } from "../utils/getNotificationsFromChatItems";
 
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const useSelectedChatPage = () => {
-  const { setMessages, setNotifications } = useChatContext();
+  const { setMessages, setNotifications, notifications } = useChatContext();
   const { getChatItems } = useChatApi();
-
+  const { getChatNotifications } = useNotificationApi();
   const params = useParams();
+
   const chatId = params?.chatId as string | undefined;
 
+  const chatNotificationsQueryKey = getChatNotificationsQueryKey(chatId ?? "");
+  const { data: fetchedNotifications = [] } = useQuery({
+    queryKey: [chatNotificationsQueryKey],
+    enabled: notifications.length > 0,
+    queryFn: () => {
+      if (chatId === undefined) {
+        return [];
+      }
+
+      return getChatNotifications(chatId);
+    },
+    refetchInterval: () => {
+      if (notifications.length === 0) {
+        return false;
+      }
+      const hasAPendingNotification = notifications.find(
+        (item) => item.status === "Pending"
+      );
+
+      if (hasAPendingNotification) {
+        //30 seconds
+        return 30_000;
+      }
+
+      return false;
+    },
+  });
+
+  useEffect(() => {
+    if (fetchedNotifications.length === 0) {
+      return;
+    }
+    setNotifications(fetchedNotifications);
+  }, [fetchedNotifications]);
+
   useEffect(() => {
     const fetchHistory = async () => {
       if (chatId === undefined) {
@@ -30,5 +69,5 @@ export const useSelectedChatPage = () => {
       setNotifications(getNotificationsFromChatItems(chatItems));
     };
     void fetchHistory();
-  }, [chatId, setMessages]);
+  }, [chatId]);
 };

frontend/app/chat/[chatId]/utils/getChatNotificationsQueryKey.ts

@@ -0,0 +1,2 @@
+export const getChatNotificationsQueryKey = (chatId: string): string =>
+  `notifications-${chatId}`;

frontend/app/chat/components/ChatsList/__tests__/ChatsList.test.tsx

@@ -1,4 +1,5 @@
 /* eslint-disable max-lines */
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { act, fireEvent, render, screen } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 
@@ -11,6 +12,7 @@ import * as useChatsListModule from "../hooks/useChatsList";
 import { ChatsList } from "../index";
 
 const getChatsMock = vi.fn(() => []);
+const queryClient = new QueryClient();
 
 const setOpenMock = vi.fn();
 
@@ -56,9 +58,11 @@ describe("ChatsList", () => {
 
   it("should render correctly", () => {
     const { getByTestId } = render(
-      <ChatProviderMock>
-        <ChatsList />
-      </ChatProviderMock>
+      <QueryClientProvider client={queryClient}>
+        <ChatProviderMock>
+          <ChatsList />
+        </ChatProviderMock>
+      </QueryClientProvider>
     );
     const chatsList = getByTestId("chats-list");
     expect(chatsList).toBeDefined();
@@ -72,9 +76,11 @@ describe("ChatsList", () => {
 
   it("renders the chats list with correct number of items", () => {
     render(
-      <ChatProviderMock>
-        <ChatsList />
-      </ChatProviderMock>
+      <QueryClientProvider client={queryClient}>
+        <ChatProviderMock>
+          <ChatsList />
+        </ChatProviderMock>
+      </QueryClientProvider>
     );
     const chatItems = screen.getAllByTestId("chats-list-item");
     expect(chatItems).toHaveLength(2);
@@ -88,9 +94,11 @@ describe("ChatsList", () => {
 
     await act(() =>
       render(
-        <ChatProviderMock>
-          (<ChatsList />)
-        </ChatProviderMock>
+        <QueryClientProvider client={queryClient}>
+          <ChatProviderMock>
+            (<ChatsList />)
+          </ChatProviderMock>
+        </QueryClientProvider>
       )
     );
 
@@ -109,9 +117,11 @@ describe("ChatsList", () => {
     }));
     await act(() =>
       render(
-        <ChatProviderMock>
-          <ChatsList />
-        </ChatProviderMock>
+        <QueryClientProvider client={queryClient}>
+          <ChatProviderMock>
+            <ChatsList />
+          </ChatProviderMock>
+        </QueryClientProvider>
       )
     );
 

frontend/lib/api/notification/__tests__/useNotificationApi.test.ts

@@ -0,0 +1,28 @@
+import { renderHook } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+
+import { useNotificationApi } from "../useNotificationApi";
+
+const axiosGetMock = vi.fn(() => ({}));
+
+vi.mock("@/lib/hooks", () => ({
+  useAxios: () => ({
+    axiosInstance: {
+      get: axiosGetMock,
+    },
+  }),
+}));
+
+describe("useNotificationApi", () => {
+  it("should call getChatNotifications with the correct parameters", async () => {
+    const chatId = "test-chat-id";
+    const {
+      result: {
+        current: { getChatNotifications },
+      },
+    } = renderHook(() => useNotificationApi());
+    await getChatNotifications(chatId);
+    expect(axiosGetMock).toHaveBeenCalledTimes(1);
+    expect(axiosGetMock).toHaveBeenCalledWith(`/notifications/${chatId}`);
+  });
+});

frontend/lib/api/notification/notification.ts

@@ -0,0 +1,11 @@
+import { AxiosInstance } from "axios";
+
+import { Notification } from "@/app/chat/[chatId]/types";
+
+export const getChatNotifications = async (
+  chatId: string,
+  axiosInstance: AxiosInstance
+): Promise<Notification[]> => {
+  return (await axiosInstance.get<Notification[]>(`/notifications/${chatId}`))
+    .data;
+};

frontend/lib/api/notification/useNotificationApi.ts

@@ -0,0 +1,13 @@
+import { useAxios } from "@/lib/hooks";
+
+import { getChatNotifications } from "./notification";
+
+// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
+export const useNotificationApi = () => {
+  const { axiosInstance } = useAxios();
+
+  return {
+    getChatNotifications: async (chatId: string) =>
+      await getChatNotifications(chatId, axiosInstance),
+  };
+};