Merge pull request #10 from code-kern-ai/websocket-change

Prep for rework including useWebsocket hook

Author: JWittmeyer

src/components/models-download/ModelsDownload.tsx

@@ -11,13 +11,12 @@ import LoadingIcon from "../shared/loading/LoadingIcon";
 import { openModal, setModalStates } from "@/src/reduxStore/states/modal";
 import { ModalEnum } from "@/src/types/shared/modal";
 import { selectIsManaged } from "@/src/reduxStore/states/general";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { CurrentPage } from "@/src/types/shared/general";
 import { timer } from "rxjs";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
 import { TOOLTIPS_DICT } from "@/src/util/tooltip-constants";
 import AddModelDownloadModal from "./AddModelDownloadModal";
 import DeleteModelDownloadModal from "./DeleteModelDownloadModal";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 export default function ModelsDownload() {
     const router = useRouter();
@@ -28,14 +27,8 @@ export default function ModelsDownload() {
 
     const [refetchModelsDownload] = useLazyQuery(GET_MODEL_PROVIDER_INFO, { fetchPolicy: 'network-only', nextFetchPolicy: 'cache-first' });
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.MODELS_DOWNLOAD]), []);
-
     useEffect(() => {
         refetchModels();
-        WebSocketsService.subscribeToNotification(CurrentPage.MODELS_DOWNLOAD, {
-            whitelist: ['model_provider_download'],
-            func: handleWebsocketNotification
-        });
     }, []);
 
     function refetchModels() {
@@ -53,9 +46,7 @@ export default function ModelsDownload() {
         }
     }, []);
 
-    useEffect(() => {
-        WebSocketsService.updateFunctionPointer(null, CurrentPage.MODELS_DOWNLOAD, handleWebsocketNotification)
-    }, [handleWebsocketNotification]);
+    useWebsocket(CurrentPage.MODELS_DOWNLOAD, handleWebsocketNotification);
 
     return (<div className="p-4 bg-gray-100 flex-1 flex flex-col h-[calc(100vh-4rem)] overflow-y-auto">
         <div className="flex flex-row items-center">

src/components/projects/ButtonsContainer.tsx

@@ -9,9 +9,8 @@ import { useCallback, useEffect, useState } from "react";
 import { useDispatch, useSelector } from "react-redux"
 import ModalUpload from "../shared/upload/ModalUpload";
 import SampleProjectsDropdown from "./SampleProjectsDropdown";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
-import { CurrentPage } from "@/src/types/shared/general";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
+import { CurrentPage, CurrentPageSubKey } from "@/src/types/shared/general";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 const BASE_OPTIONS = { reloadOnFinish: false, deleteProjectOnFail: true, closeModalOnClick: true, isModal: true, navigateToProject: true, showBadPasswordMsg: null };
 
@@ -24,15 +23,6 @@ export default function ButtonsContainer() {
     const [uploadOptions, setUploadOptions] = useState<UploadOptions>(BASE_OPTIONS);
     const [showBadPasswordMsg, setShowBadPasswordMsg] = useState(false);
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.PROJECTS]), []);
-
-    useEffect(() => {
-        WebSocketsService.subscribeToNotification(CurrentPage.PROJECTS, {
-            whitelist: ['bad_password'],
-            func: handleWebsocketNotification
-        });
-    }, []);
-
     useEffect(() => {
         setUploadOptions({ ...BASE_OPTIONS, showBadPasswordMsg: showBadPasswordMsg });
     }, [showBadPasswordMsg]);
@@ -43,9 +33,7 @@ export default function ButtonsContainer() {
         }
     }, []);
 
-    useEffect(() => {
-        WebSocketsService.updateFunctionPointer(null, CurrentPage.PROJECTS, handleWebsocketNotification)
-    }, [handleWebsocketNotification]);
+    useWebsocket(CurrentPage.PROJECTS, handleWebsocketNotification, null, CurrentPageSubKey.BUTTONS_CONTAINER);
 
     return (
         user && user.role === UserRole.ENGINEER ? (<div>

src/components/projects/ProjectsList.tsx

@@ -1,6 +1,5 @@
 import { selectInactiveOrganization, selectIsDemo, selectIsManaged, selectUser, setComments } from "@/src/reduxStore/states/general"
 import { selectAllProjects, setAllProjects } from "@/src/reduxStore/states/project";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { GET_OVERVIEW_STATS, GET_PROJECT_LIST } from "@/src/services/gql/queries/projects";
 import { Project, ProjectStatistics } from "@/src/types/components/projects/projects-list";
 import { CurrentPage } from "@/src/types/shared/general";
@@ -15,14 +14,12 @@ import { GET_CAN_CREATE_LOCAL_ORG } from "@/src/services/gql/queries/organizatio
 import { ADD_USER_TO_ORGANIZATION, CREATE_ORGANIZATION } from "@/src/services/gql/mutations/organizations";
 import style from "@/src/styles/components/projects/projects-list.module.css";
 import { useRouter } from "next/router";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
 import AdminDeleteProjectModal from "./AdminDeleteProjectModal";
 import { setAllAttributes, setAllEmbeddings, setLabelingTasksAll } from "@/src/reduxStore/states/pages/settings";
 import { setOverviewFilters } from "@/src/reduxStore/states/tmp";
 import { setDataSlices, setFullSearchStore, setSearchGroupsStore } from "@/src/reduxStore/states/pages/data-browser";
-import { closeModal } from "@/src/reduxStore/states/modal";
-import { ModalEnum } from "@/src/types/shared/modal";
 import { SearchGroup } from "@/submodules/javascript-functions/enums/enums";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 export default function ProjectsList() {
     const router = useRouter();
@@ -44,13 +41,7 @@ export default function ProjectsList() {
     const [createOrgMut] = useMutation(CREATE_ORGANIZATION);
     const [addUserToOrgMut] = useMutation(ADD_USER_TO_ORGANIZATION);
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.PROJECTS]), []);
-
     useEffect(() => {
-        WebSocketsService.subscribeToNotification(CurrentPage.PROJECTS, {
-            whitelist: ['project_created', 'project_deleted', 'project_update', 'file_upload'],
-            func: handleWebsocketNotification
-        });
         dispatch(setLabelingTasksAll(null));
         dispatch(setOverviewFilters(null));
         dispatch(setSearchGroupsStore(null));
@@ -119,9 +110,7 @@ export default function ProjectsList() {
         }
     }, []);
 
-    useEffect(() => {
-        WebSocketsService.updateFunctionPointer(null, CurrentPage.PROJECTS, handleWebsocketNotification)
-    }, [handleWebsocketNotification]);
+    useWebsocket(CurrentPage.PROJECTS, handleWebsocketNotification);
 
     return (
         <div>

src/components/projects/projectId/admin/ProjectAdmin.tsx

@@ -13,10 +13,9 @@ import { useDispatch, useSelector } from "react-redux";
 import NewPersonalToken from "./NewPersonalTokenModal";
 import DeletePersonalToken from "./DeletePersonalTokenModal";
 import { CurrentPage } from "@/src/types/shared/general";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { useRouter } from "next/router";
 import { selectIsAdmin } from "@/src/reduxStore/states/general";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 export default function ProjectAdmin() {
     const router = useRouter();
@@ -29,8 +28,6 @@ export default function ProjectAdmin() {
 
     const [refetchAccessTokens] = useLazyQuery(GET_ALL_PERSONAL_ACCESS_TOKENS, { fetchPolicy: "network-only" });
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.ADMIN_PAGE]), []);
-
     useEffect(() => {
         if (!projectId) return;
         refetchAccessTokensAndProcess();
@@ -41,11 +38,6 @@ export default function ProjectAdmin() {
             return;
         }
 
-        WebSocketsService.subscribeToNotification(CurrentPage.ADMIN_PAGE, {
-            projectId: projectId,
-            whitelist: ['pat'],
-            func: handleWebsocketNotification
-        });
     }, [projectId]);
 
     function refetchAccessTokensAndProcess() {
@@ -64,10 +56,7 @@ export default function ProjectAdmin() {
         }
     }, [accessTokens]);
 
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.updateFunctionPointer(projectId, CurrentPage.ADMIN_PAGE, handleWebsocketNotification)
-    }, [handleWebsocketNotification, projectId]);
+    useWebsocket(CurrentPage.ADMIN_PAGE, handleWebsocketNotification, projectId);
 
     return (<>
         {accessTokens && <div className="p-4 bg-gray-100 h-full overflow-y-auto flex-1 flex flex-col">

src/components/projects/projectId/attributes/attributeId/AttributeCalculations.tsx

@@ -23,17 +23,16 @@ import DangerZone from "@/src/components/shared/danger-zone/DangerZone";
 import { DangerZoneEnum } from "@/src/types/shared/danger-zone";
 import ContainerLogs from "@/src/components/shared/logs/ContainerLogs";
 import LoadingIcon from "@/src/components/shared/loading/LoadingIcon";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { debounceTime, distinctUntilChanged, fromEvent, timer } from "rxjs";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
 import { TOOLTIPS_DICT } from "@/src/util/tooltip-constants";
-import { selectAllUsers, setBricksIntegrator, setComments } from "@/src/reduxStore/states/general";
+import { selectAllUsers, setComments } from "@/src/reduxStore/states/general";
 import { REQUEST_COMMENTS } from "@/src/services/gql/queries/projects";
 import { CommentDataManager } from "@/src/util/classes/comments";
 import { CommentType } from "@/src/types/shared/comments";
 import BricksIntegrator from "@/src/components/shared/bricks-integrator/BricksIntegrator";
 import { AttributeCodeLookup } from "@/src/util/classes/attribute-calculation";
 import Dropdown2 from "@/submodules/react-components/components/Dropdown2";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 import { postProcessLabelingTasks, postProcessLabelingTasksSchema } from "@/src/util/components/projects/projectId/settings/labeling-tasks-helper";
 
 const EDITOR_OPTIONS = { theme: 'vs-light', language: 'python', readOnly: false };
@@ -68,8 +67,6 @@ export default function AttributeCalculation() {
     const [refetchComments] = useLazyQuery(REQUEST_COMMENTS, { fetchPolicy: "no-cache" });
     const [refetchLabelingTasksByProjectId] = useLazyQuery(GET_LABELING_TASKS_BY_PROJECT_ID, { fetchPolicy: "network-only" });
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.ATTRIBUTE_CALCULATION]), []);
-
     useEffect(() => {
         if (!currentAttribute) return;
         if (isInitial == null) setIsInitial(AttributeCodeLookup.isCodeStillTemplate(currentAttribute.sourceCode, currentAttribute.dataType))
@@ -92,11 +89,6 @@ export default function AttributeCalculation() {
         }
         refetchLabelingTasksAndProcess();
         checkProjectTokenization();
-        WebSocketsService.subscribeToNotification(CurrentPage.ATTRIBUTE_CALCULATION, {
-            projectId: projectId,
-            whitelist: ['attributes_updated', 'calculate_attribute', 'tokenization', 'knowledge_base_updated', 'knowledge_base_deleted', 'knowledge_base_created'],
-            func: handleWebsocketNotification
-        });
     }, [projectId, attributes, currentAttribute]);
 
     useEffect(() => {
@@ -298,10 +290,7 @@ export default function AttributeCalculation() {
         }
     }, [projectId, currentAttribute]);
 
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.updateFunctionPointer(projectId, CurrentPage.ATTRIBUTE_CALCULATION, handleWebsocketNotification)
-    }, [handleWebsocketNotification, projectId]);
+    useWebsocket(CurrentPage.ATTRIBUTE_CALCULATION, handleWebsocketNotification, projectId);
 
     return (projectId && <div className="bg-white p-4 pb-16 overflow-y-auto h-screen" style={{ width: 'calc(100vw - 75px)' }} onScroll={(e: any) => onScrollEvent(e)}>
         {currentAttribute && <div>

src/components/projects/projectId/data-browser/DataBrowser.tsx

@@ -13,19 +13,16 @@ import { GET_ORGANIZATION_USERS_WITH_COUNT } from "@/src/services/gql/queries/or
 import { selectAllUsers, selectUser, setComments } from "@/src/reduxStore/states/general";
 import DataBrowserRecords from "./DataBrowserRecords";
 import { postProcessingEmbeddings } from "@/src/util/components/projects/projectId/settings/embeddings-helper";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { CurrentPage } from "@/src/types/shared/general";
 import { CommentType } from "@/src/types/shared/comments";
 import { CommentDataManager } from "@/src/util/classes/comments";
 import { REQUEST_COMMENTS } from "@/src/services/gql/queries/projects";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
-import { useRouter } from "next/router";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 const SEARCH_REQUEST = { offset: 0, limit: 20 };
 
 export default function DataBrowser() {
     const dispatch = useDispatch();
-    const router = useRouter();
 
     const projectId = useSelector(selectProjectId);
     const users = useSelector(selectAllUsers);
@@ -45,8 +42,6 @@ export default function DataBrowser() {
     const [refetchComments] = useLazyQuery(REQUEST_COMMENTS, { fetchPolicy: "no-cache" });
     const [refetchUniqueValues] = useLazyQuery(GET_UNIQUE_VALUES_BY_ATTRIBUTES, { fetchPolicy: "no-cache" });
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.DATA_BROWSER]), []);
-
     useEffect(() => {
         if (!projectId) return;
         if (!users || !user) return;
@@ -56,11 +51,6 @@ export default function DataBrowser() {
         refetchUsersCountAndProcess();
         refetchEmbeddingsAndPostProcess();
         refetchUniqueValuesAndProcess();
-        WebSocketsService.subscribeToNotification(CurrentPage.DATA_BROWSER, {
-            projectId: projectId,
-            whitelist: ['data_slice_created', 'data_slice_updated', 'data_slice_deleted', 'label_created', 'label_deleted', 'labeling_task_deleted', 'labeling_task_updated', 'labeling_task_created', 'information_source_created', 'information_source_updated', 'information_source_deleted', 'attributes_updated', 'calculate_attribute', 'embedding', 'embedding_deleted'],
-            func: handleWebsocketNotification
-        });
     }, [projectId, users, user]);
 
     useEffect(() => {
@@ -181,10 +171,7 @@ export default function DataBrowser() {
         }
     }, [projectId]);
 
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.updateFunctionPointer(projectId, CurrentPage.DATA_BROWSER, handleWebsocketNotification)
-    }, [handleWebsocketNotification, projectId]);
+    useWebsocket(CurrentPage.DATA_BROWSER, handleWebsocketNotification, projectId);
 
     return (<>
         {projectId && <div className="flex flex-row h-full">

src/components/projects/projectId/edit-records/EditRecords.tsx

@@ -10,11 +10,10 @@ import { IconAlertCircle, IconAlertTriangle, IconAlertTriangleFilled, IconBallpe
 import { Tooltip } from "@nextui-org/react";
 import { TOOLTIPS_DICT } from "@/src/util/tooltip-constants";
 import EditField from "./EditField";
-import { WebSocketsService } from "@/src/services/base/web-sockets/WebSocketsService";
 import { useRouter } from "next/router";
-import { unsubscribeWSOnDestroy } from "@/src/services/base/web-sockets/web-sockets-helper";
 import { ModalEnum } from "@/src/types/shared/modal";
 import { openModal } from "@/src/reduxStore/states/modal";
+import { useWebsocket } from "@/src/services/base/web-sockets/useWebsocket";
 
 export default function EditRecords() {
     const dispatch = useDispatch();
@@ -26,8 +25,6 @@ export default function EditRecords() {
     const [erdData, setErdData] = useState(createDefaultEditRecordComponentData());
     const [alertLastVisible, setAlertLastVisible] = useState<number>(null);
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.EDIT_RECORDS]), []);
-
     useEffect(() => {
         if (!projectId) return;
         if (!sessionData) {
@@ -44,11 +41,6 @@ export default function EditRecords() {
         if (!erdDataCopy.modals.hideExplainModal) {
             dispatch(openModal(ModalEnum.EXPLAIN_EDIT_RECORDS));
         }
-        WebSocketsService.subscribeToNotification(CurrentPage.EDIT_RECORDS, {
-            projectId: projectId,
-            whitelist: ['calculate_attribute'],
-            func: handleWebsocketNotification
-        });
     }, [sessionData, projectId]);
 
     useEffect(() => {
@@ -77,10 +69,7 @@ export default function EditRecords() {
         setAlertLastVisible(Date.now());
     }, [alertLastVisible]);
 
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.updateFunctionPointer(projectId, CurrentPage.EDIT_RECORDS, handleWebsocketNotification)
-    }, [handleWebsocketNotification, projectId]);
+    useWebsocket(CurrentPage.EDIT_RECORDS, handleWebsocketNotification, projectId);
 
     return (<>{projectId && <div className="h-screen bg-white flex flex-col">
         <NavBarTopEditRecords erdData={erdData} setErdData={(erdData) => setErdData(erdData)} />

src/components/projects/projectId/heuristics/HeuristicsHeader.tsx

@@ -4,8 +4,6 @@ import { openModal, selectModal } from '@/src/reduxStore/states/modal';
 import { selectHeuristicsAll, setHeuristicType } from '@/src/reduxStore/states/pages/heuristics';
 import { selectLabelingTasksAll } from '@/src/reduxStore/states/pages/settings';
 import { selectProjectId } from '@/src/reduxStore/states/project';
-import { WebSocketsService } from '@/src/services/base/web-sockets/WebSocketsService';
-import { unsubscribeWSOnDestroy } from '@/src/services/base/web-sockets/web-sockets-helper';
 import { CREATE_INFORMATION_SOURCE_PAYLOAD, RUN_ZERO_SHOT_PROJECT, SET_ALL_HEURISTICS, START_WEAK_SUPERVISIONS } from '@/src/services/gql/mutations/heuristics';
 import { GET_CURRENT_WEAK_SUPERVISION_RUN } from '@/src/services/gql/queries/heuristics';
 import style from '@/src/styles/components/projects/projectId/heuristics/heuristics.module.css';
@@ -26,6 +24,7 @@ import { useDispatch, useSelector } from 'react-redux';
 import LastWeakSupervisionModal from './modals/LastWeakSupervisionModal';
 import DeleteHeuristicsModal from './DeleteHeuristicsModal';
 import Dropdown2 from '@/submodules/react-components/components/Dropdown2';
+import { useWebsocket } from '@/src/services/base/web-sockets/useWebsocket';
 
 
 export default function HeuristicsHeader(props: HeuristicsHeaderProps) {
@@ -52,17 +51,6 @@ export default function HeuristicsHeader(props: HeuristicsHeaderProps) {
     const [createTaskMut] = useMutation(CREATE_INFORMATION_SOURCE_PAYLOAD);
     const [runZeroShotMut] = useMutation(RUN_ZERO_SHOT_PROJECT);
 
-    useEffect(unsubscribeWSOnDestroy(router, [CurrentPage.HEURISTICS]), []);
-
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.subscribeToNotification(CurrentPage.HEURISTICS, {
-            projectId: projectId,
-            whitelist: ['weak_supervision_started', 'weak_supervision_finished'],
-            func: handleWebsocketNotification
-        });
-    }, [projectId]);
-
     useEffect(() => {
         if (!heuristics) return;
         setAreHeuristicsSelected(checkSelectedHeuristics(heuristics, false));
@@ -173,10 +161,7 @@ export default function HeuristicsHeader(props: HeuristicsHeaderProps) {
         else if (msgParts[1] == 'weak_supervision_finished') setLoadingIconWS(false);
     }, []);
 
-    useEffect(() => {
-        if (!projectId) return;
-        WebSocketsService.updateFunctionPointer(projectId, CurrentPage.HEURISTICS, handleWebsocketNotification)
-    }, [handleWebsocketNotification, projectId]);
+    useWebsocket(CurrentPage.HEURISTICS, handleWebsocketNotification, projectId);
 
     return (
         <div className="flex-shrink-0 block xl:flex justify-between items-center">