diff --git a/src/frontend/src/modals/chatModal/index.tsx b/src/frontend/src/modals/chatModal/index.tsx
index c8a227ae8..bf2de6718 100644
--- a/src/frontend/src/modals/chatModal/index.tsx
+++ b/src/frontend/src/modals/chatModal/index.tsx
@@ -1,8 +1,5 @@
 import { Dialog, Transition } from "@headlessui/react";
-import {
-  ChatBubbleOvalLeftEllipsisIcon,
-  XMarkIcon,
-} from "@heroicons/react/24/outline";
+import { ChatBubbleOvalLeftEllipsisIcon } from "@heroicons/react/24/outline";
 import { Fragment, useContext, useEffect, useRef, useState } from "react";
 import { FlowType, NodeType } from "../../types/flow";
 import { alertContext } from "../../contexts/alertContext";
@@ -12,10 +9,10 @@ import ChatMessage from "./chatMessage";
 import { FaEraser } from "react-icons/fa";
 import { HiX } from "react-icons/hi";
 import { sendAllProps } from "../../types/api";
-import { ChatMessageType, ChatType } from "../../types/chat";
+import { ChatMessageType } from "../../types/chat";
 import ChatInput from "./chatInput";
 
-import _, { set } from "lodash";
+import _ from "lodash";
 
 export default function ChatModal({
   flow,
@@ -177,15 +174,23 @@ export default function ChatModal({
       updateLastMessage({ str: data.message });
     }
   }
+  function getWebSocketUrl(chatId, isDevelopment = false) {
+    const isSecureProtocol = window.location.protocol === "https:";
+    const webSocketProtocol = isSecureProtocol ? "wss" : "ws";
+    const host = isDevelopment ? "localhost:7860" : window.location.host;
+    const chatEndpoint = `/api/v1/chat/${chatId}`;
+
+    return `${
+      isDevelopment ? "ws" : webSocketProtocol
+    }://${host}${chatEndpoint}`;
+  }
 
   function connectWS() {
     try {
-      const urlWs =
+      const urlWs = getWebSocketUrl(
+        id.current,
         process.env.NODE_ENV === "development"
-          ? `ws://localhost:7860/api/v1/chat/${id.current}`
-          : `${window.location.protocol === "https:" ? "wss" : "ws"}://${
-              window.location.host
-            }api/v1/chat/${id.current}`;
+      );
       const newWs = new WebSocket(urlWs);
       newWs.onopen = () => {
         console.log("WebSocket connection established!");