Dependabot fixes + mongodb vector search update + improvement on websocket connection on UI (#43)

* working version pre dependency bot break

* update websocket connection feedback and pom updates

* version bumps in search function and mongodb vector search update

* bump memory allocation for lambdas

Author: stephenheg

frontend/package-lock.json

@@ -1,6 +1,6 @@
 import { WS_URL } from "astro:env/client";
 import useWebSocket, { ReadyState } from "react-use-websocket";
-import React, { useState, useCallback } from "react";
+import React, { useState, useCallback, useEffect } from "react";
 import ChatInput from "./ChatInput.tsx";
 import ChatMessages from "./ChatMessages.tsx";
 import type { ChatMessage } from "./ChatMessage.tsx";
@@ -21,7 +21,9 @@ export default function Chat({ username }: { username: string }) {
 
   const [input, setInput] = useState("");
   const [messageHistory, setMessageHistory] = useState<ChatMessage[]>([]);
-  const [isLoading, setIsLoading] = useState(false); // Add this line
+  const [isLoading, setIsLoading] = useState(false);
+  const [connectionStatus, setConnectionStatus] = useState<string>("Connecting...");
+  const [isReconnecting, setIsReconnecting] = useState(false);
 
   const handleMessage = useCallback((event: WebSocketEventMap["message"]) => {
     console.log(event);
@@ -37,25 +39,92 @@ export default function Chat({ username }: { username: string }) {
     } catch (e) {
       console.debug("Invalid JSON message received:", event.data);
     }
-    setIsLoading(false); // Add this line
+    setIsLoading(false);
   }, []);
 
   const { sendJsonMessage, readyState } = useWebSocket<ChatMessage>(WS_URL, {
-    onOpen: () => console.log("WebSocket connection opened!"),
-    onClose: (event) => console.log("WebSocket connection closed!: ", event),
-    onError: (event) => console.error("WebSocket error:", event),
+    onOpen: () => {
+      console.log("WebSocket connection opened!");
+      setConnectionStatus("Connected");
+      setIsReconnecting(false);
+    },
+    onClose: (event) => {
+      console.log("WebSocket connection closed!: ", event);
+      // Only show disconnected if we're not actively reconnecting
+      if (!isReconnecting) {
+        setConnectionStatus("Disconnected");
+      }
+    },
+    onError: (event) => {
+      console.error("WebSocket error:", event);
+      setConnectionStatus("Error - Retrying...");
+      setIsReconnecting(true);
+    },
     onMessage: handleMessage,
+    onReconnectStop: () => {
+      console.log("Reconnection attempts stopped");
+      setConnectionStatus("Disconnected");
+      setIsReconnecting(false);
+    },
+    // Add retry configuration
+    shouldReconnect: (closeEvent) => {
+      console.log("Should reconnect:", closeEvent);
+      setIsReconnecting(true);
+      setConnectionStatus("Reconnecting...");
+      return true; // Always try to reconnect
+    },
+    reconnectInterval: (attemptNumber) => {
+      // Exponential backoff: 1s, 2s, 4s, 8s, max 30s
+      return Math.min(1000 * Math.pow(2, attemptNumber), 30000);
+    },
+    reconnectAttempts: 10,
+    // Connection timeout
+    heartbeat: {
+      message: JSON.stringify({ type: "ping" }),
+      returnMessage: JSON.stringify({ type: "pong" }),
+      timeout: 60000, // 60 seconds
+      interval: 60000, // 60 seconds
+    },
   });
 
+  // Update connection status based on readyState
+  useEffect(() => {
+    switch (readyState) {
+      case ReadyState.CONNECTING:
+        if (!isReconnecting) {
+          setConnectionStatus("Connecting...");
+        }
+        break;
+      case ReadyState.OPEN:
+        setConnectionStatus("Connected");
+        setIsReconnecting(false);
+        break;
+      case ReadyState.CLOSING:
+        setConnectionStatus("Disconnecting...");
+        break;
+      case ReadyState.CLOSED:
+        if (!isReconnecting) {
+          setConnectionStatus("Disconnected");
+        }
+        break;
+      case ReadyState.UNINSTANTIATED:
+        setConnectionStatus("Not Connected");
+        break;
+    }
+  }, [readyState, isReconnecting]);
+
   const handleInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     setInput(event.target.value);
   };
 
   const handleSubmit = (event: React.FormEvent<HTMLFormElement>) => {
     event.preventDefault();
     if (input.trim() === "") return;
+    
+    // Only prevent submission if truly disconnected (not reconnecting)
+    if (readyState !== ReadyState.OPEN && !isReconnecting) return;
 
-    setIsLoading(true); // Add this line
+    setIsLoading(true);
 
     const chatMessage: ChatMessage = {
       role: "client",
@@ -79,8 +148,29 @@ export default function Chat({ username }: { username: string }) {
 
     console.log(clientMessage);
 
-    sendJsonMessage(clientMessage);
-    setInput("");
+    try {
+      sendJsonMessage(clientMessage);
+      setInput("");
+    } catch (error) {
+      console.error("Failed to send message:", error);
+      setIsLoading(false);
+      // Optionally show error message to user
+    }
+  };
+
+  const getConnectionStatusColor = () => {
+    if (isReconnecting) {
+      return "text-yellow-500"; // Show yellow during reconnection attempts
+    }
+    
+    switch (readyState) {
+      case ReadyState.OPEN:
+        return "text-green-500";
+      case ReadyState.CONNECTING:
+        return "text-yellow-500";
+      default:
+        return "text-red-500";
+    }
   };
 
   return (
@@ -95,17 +185,22 @@ export default function Chat({ username }: { username: string }) {
         <div className="fixed bottom-32 right-4 bg-gray-800 text-white rounded-lg shadow-lg flex flex-col w-11/12 md:w-1/3">
           <div className="flex justify-between items-center p-4 bg-blue-500 rounded-t-lg">
             <h1 className="text-xl font-bold">Big Friendly Bank</h1>
-            <button onClick={toggleChat} className="text-white">
-              X
-            </button>
+            <div className="flex items-center gap-2">
+              <span className={`text-sm ${getConnectionStatusColor()}`}>
+                {connectionStatus}
+              </span>
+              <button onClick={toggleChat} className="text-white">
+                X
+              </button>
+            </div>
           </div>
           <div className="flex flex-col flex-grow w-full max-w-screen-lg rounded-lg h-full overflow-y-auto">
             <ChatMessages messages={messageHistory} isLoading={isLoading} />
             <ChatInput
               input={input}
               handleSubmit={handleSubmit}
               handleInputChange={handleInputChange}
-              isDisabled={readyState !== ReadyState.OPEN}
+              isDisabled={readyState !== ReadyState.OPEN && !isReconnecting}
             />
           </div>
         </div>

frontend/src/components/Chat.tsx

@@ -1,4 +1,7 @@
 boto3==1.38.18
 confluent-kafka==2.10.0
 jsonschema==4.23.0
-requests==2.32.4
+requests==2.32.4
+httpx>=0.28.0
+cachetools>=6.1.0
+authlib==1.6.0 

infrastructure/modules/backend/functions/requirements.txt

@@ -117,6 +117,7 @@ module "websocket_chat_api_lambda" {
   handler       = "websocket_lambda.lambda_handler"
   runtime       = "python3.12"
   architectures = [var.system_architecture]
+  memory_size   = 512
 
   environment_variables = {
     SCHEMA_REGISTRY_URL        = var.schema_registry_url
@@ -281,6 +282,7 @@ module "lambda_trigger_vector_search" {
   handler                           = "io.confluent.pie.search.SearchHandler::handleRequest"
   runtime                           = "java17"
   architectures                     = [var.system_architecture]
+  memory_size                       = 1024
   timeout                           = 600
   snap_start                        = false
   provisioned_concurrent_executions = 1

infrastructure/modules/backend/main.tf

@@ -30,13 +30,13 @@
         <dependency>
             <groupId>io.confluent</groupId>
             <artifactId>kafka-json-schema-serializer</artifactId>
-            <version>7.9.1-180</version>
+            <version>7.7.1</version>
         </dependency>
 
         <dependency>
             <groupId>org.apache.kafka</groupId>
             <artifactId>kafka-clients</artifactId>
-            <version>3.9.1</version>
+            <version>3.7.1</version>
         </dependency>
         <!-- https://mvnrepository.com/artifact/com.amazonaws/aws-java-sdk-secretsmanager -->
         <dependency>

infrastructure/modules/backend/search/pom.xml

@@ -46,7 +46,7 @@ public Producer<SearchResultsKey, SearchResults> get() {
                     configuration.getCcCredentials(),
                     configuration.getSrURL(),
                     configuration.getSrCredentials(),
-                    false);
+                    true); // true to enable schema registry
 
             return new KafkaProducer<>(properties);
         } catch (Throwable e) {

infrastructure/modules/backend/search/src/main/java/io/confluent/pie/search/services/KafkaProducerSupplier.java

@@ -5,6 +5,7 @@
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
 import com.mongodb.client.model.search.FieldSearchPath;
+import com.mongodb.client.model.search.VectorSearchOptions;
 import io.confluent.pie.search.models.Credentials;
 import io.confluent.pie.search.models.SearchRequest;
 import io.confluent.pie.search.services.DBService;
@@ -19,7 +20,9 @@
 import static com.google.common.primitives.Doubles.asList;
 import static com.mongodb.client.model.Aggregates.project;
 import static com.mongodb.client.model.Aggregates.vectorSearch;
-import static com.mongodb.client.model.Projections.*;
+import static com.mongodb.client.model.Projections.exclude;
+import static com.mongodb.client.model.Projections.fields;
+import static com.mongodb.client.model.Projections.metaVectorSearchScore;
 import static com.mongodb.client.model.search.SearchPath.fieldPath;
 
 /**
@@ -81,12 +84,14 @@ public List<Document> find(SearchRequest request) {
         final List<Document> products = new ArrayList<>();
         final boolean hasScoreFilter = request.minScore() > 0;
 
+        final VectorSearchOptions options = VectorSearchOptions.approximateVectorSearchOptions((long) request.numberOfCandidate());
+        
         final List<Bson> pipeline = List.of(vectorSearch(
                         fieldSearchPath,
                         asList(request.embeddings()),
                         indexName,
-                        request.numberOfCandidate(),
-                        request.limit()),
+                        (long) request.limit(),
+                        options),
                 project(fields(
                         exclude(embeddingsFieldName),
                         metaVectorSearchScore("score"))));