FEAT: Unified Notifications

backend/account/authentication_controller.py

@@ -39,6 +39,7 @@
 from django.middleware import csrf
 from django.shortcuts import redirect
 from django_tenants.utils import tenant_context
+from logs_helper.log_service import LogService
 from rest_framework import status
 from rest_framework.request import Request
 from rest_framework.response import Response
@@ -259,6 +260,8 @@ def make_user_organization_display_name(self, user_name: str) -> str:
         return self.auth_service.make_user_organization_display_name(user_name)
 
     def user_logout(self, request: Request) -> Response:
+        session_id: str = request.COOKIES.get("sessionid")
+        LogService.remove_logs_on_logout(session_id=session_id)
         response = self.auth_service.user_logout(request=request)
         organization_id = UserSessionUtils.get_organization_id(request)
         user_id = UserSessionUtils.get_user_id(request)

backend/logs_helper/log_service.py

@@ -0,0 +1,17 @@
+from django_redis import get_redis_connection
+
+
+class LogService:
+    @staticmethod
+    def remove_logs_on_logout(session_id):
+
+        if session_id:
+            # Get the Redis connection
+            r = get_redis_connection("default")
+
+            key_pattern = f"logs:{session_id}*"
+
+            # Retrieve keys matching the pattern and delete them
+            keys = r.keys(key_pattern)
+            if keys:
+                r.delete(*keys)

backend/logs_helper/urls.py

@@ -1,9 +1,9 @@
 from django.urls import path
 from rest_framework.urlpatterns import format_suffix_patterns
 
-from .views import LogsHelperView
+from .views import LogsHelperViewSet
 
-logs_helper_get = LogsHelperView.as_view({"get": "get_logs", "post": "store_log"})
+logs_helper_get = LogsHelperViewSet.as_view({"get": "get_logs", "post": "store_log"})
 
 urlpatterns = format_suffix_patterns(
     [

backend/logs_helper/views.py

@@ -3,12 +3,10 @@
 import os
 from datetime import datetime, timezone
 
-import redis
+from django_redis import get_redis_connection
 from logs_helper.constants import LogsHelperKeys
-from logs_helper.exceptions import InvalidValueError, MissingFieldsKeyError
 from rest_framework import status, viewsets
 from rest_framework.decorators import action
-from rest_framework.exceptions import APIException
 from rest_framework.response import Response
 from utils.local_context import StateStore
 
@@ -17,77 +15,52 @@
 logger = logging.getLogger(__name__)
 
 
-class LogsHelperView(viewsets.ModelViewSet):
+class LogsHelperViewSet(viewsets.ModelViewSet):
     """Viewset to handle all Tool Studio prompt related API logics."""
 
-    r = redis.Redis(
-        host=os.environ.get("REDIS_HOST", "http://localhost"),
-        port=os.environ.get("REDIS_PORT", "6379"),
-        username=os.environ.get("REDIS_USER", ""),
-        password=os.environ.get("REDIS_PASSWORD", ""),
-    )
+    r = get_redis_connection("default")
 
     @action(detail=False, methods=["get"])
     def get_logs(self, request):
-        try:
-            # Extract the session ID
-            session_id: str = StateStore.get(LogsHelperKeys.LOG_EVENTS_ID)
-
-            # Construct the Redis key pattern to match keys
-            # associated with the session ID
-            redis_key = f"logs:{session_id}*"
-
-            # Retrieve keys matching the pattern
-            keys = self.r.keys(redis_key)
-
-            # Retrieve values corresponding to the keys and sort them by timestamp
-            logs = []
-            for key in keys:
-                log_data = self.r.get(key).decode()
-                log_entry = json.loads(log_data)
-                logs.append(log_entry)
-
-            # Sort logs based on timestamp
-            sorted_logs = sorted(logs, key=lambda x: x["timestamp"])
-
-            return Response({"data": sorted_logs}, status=status.HTTP_200_OK)
-        except Exception as e:
-            # Handle other exceptions
-            error_msg = "An unexpected error occurred while retrieving logs"
-            logger.error(f"{error_msg}: {e}")
-            raise APIException(error_msg)
+        # Extract the session ID
+        session_id: str = StateStore.get(LogsHelperKeys.LOG_EVENTS_ID)
+
+        # Construct the Redis key pattern to match keys
+        # associated with the session ID
+        redis_key = f"logs:{session_id}*"
+
+        # Retrieve keys matching the pattern
+        keys = self.r.keys(redis_key)
+
+        # Retrieve values corresponding to the keys and sort them by timestamp
+        logs = []
+        for key in keys:
+            log_data = self.r.get(key).decode()
+            log_entry = json.loads(log_data)
+            logs.append(log_entry)
+
+        # Sort logs based on timestamp
+        sorted_logs = sorted(logs, key=lambda x: x["timestamp"])
+
+        return Response({"data": sorted_logs}, status=status.HTTP_200_OK)
 
     @action(detail=False, methods=["post"])
     def store_log(self, request):
         """Store log message in Redis."""
-        try:
-            # Extract the session ID
-            logs_expiry = int(os.environ.get("LOGS_EXPIRATION_TIME_IN_SECOND", 3600))
-            session_id: str = StateStore.get(LogsHelperKeys.LOG_EVENTS_ID)
-
-            serializer = StoreLogMessagesSerializer(data=request.data)
-            serializer.is_valid(raise_exception=True)
-
-            # Extract the log message from the validated data
-            log: str = serializer.validated_data.get("log")
-
-            timestamp = datetime.now(timezone.utc).timestamp()
-
-            redis_key = f"logs:{session_id}:{timestamp}"
-
-            self.r.setex(redis_key, logs_expiry, log)
-
-            return Response({"message": "Successfully stored the message in redis"})
-        except KeyError as e:
-            # Handle KeyError
-            logger.error(f"Log is missing: {e}")
-            raise MissingFieldsKeyError()
-        except ValueError as e:
-            # Handle ValueError
-            logger.error(f"Invalid value: {e}")
-            raise InvalidValueError()
-        except Exception as e:
-            # Handle other exceptions
-            error_msg = "An unexpected error occurred while store the log message"
-            logger.error(f"{error_msg}: {e}")
-            raise APIException(error_msg)
+        # Extract the session ID
+        logs_expiry = int(os.environ.get("LOGS_EXPIRATION_TIME_IN_SECOND", 3600))
+        session_id: str = StateStore.get(LogsHelperKeys.LOG_EVENTS_ID)
+
+        serializer = StoreLogMessagesSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+
+        # Extract the log message from the validated data
+        log: str = serializer.validated_data.get("log")
+
+        timestamp = datetime.now(timezone.utc).timestamp()
+
+        redis_key = f"logs:{session_id}:{timestamp}"
+
+        self.r.setex(redis_key, logs_expiry, log)
+
+        return Response({"message": "Successfully stored the message in redis"})

frontend/src/components/helpers/custom-tools/CustomToolsHelper.js

@@ -92,6 +92,7 @@ function CustomToolsHelper() {
   useEffect(() => {
     return () => {
       setDefaultCustomTool();
+      resetTokenUsage();
     };
   }, []);
 

frontend/src/layouts/page-layout/PageLayout.jsx

@@ -41,7 +41,7 @@ function PageLayout() {
       const response = await axios(requestOptions);
       return response.data;
     } catch (error) {
-      return;
+      return { data: [] };
     }
   };
 

unstract/core/src/unstract/core/pubsub_helper.py

@@ -4,16 +4,11 @@
 from datetime import datetime, timezone
 from typing import Any, Optional
 
-import redis
+from django_redis import get_redis_connection
 
 
 class LogPublisher:
-    r = redis.Redis(
-        host=os.environ.get("REDIS_HOST", "http://localhost"),
-        port=os.environ.get("REDIS_PORT", "6379"),
-        username=os.environ.get("REDIS_USER", ""),
-        password=os.environ.get("REDIS_PASSWORD", ""),
-    )
+    r = get_redis_connection("default")
 
     @staticmethod
     def log_usage(