diff --git a/src/frontend/src/components/chatComponent/index.tsx b/src/frontend/src/components/chatComponent/index.tsx
index 0b251ee16..8108f3f36 100644
--- a/src/frontend/src/components/chatComponent/index.tsx
+++ b/src/frontend/src/components/chatComponent/index.tsx
@@ -26,7 +26,7 @@ export default function Chat({ flow }: ChatType) {
 		}, []);
 	return (
 		<>
-			<ChatModal flow={flow} open={open} setOpen={setOpen} />
+			<ChatModal key={flow.id} flow={flow} open={open} setOpen={setOpen} />
 			<ChatTrigger open={open} setOpen={setOpen} flow={flow} />
 		</>
 	);
diff --git a/src/frontend/src/modals/chatModal/chatMessage/index.tsx b/src/frontend/src/modals/chatModal/chatMessage/index.tsx
index d49d5244d..d960a149b 100644
--- a/src/frontend/src/modals/chatModal/chatMessage/index.tsx
+++ b/src/frontend/src/modals/chatModal/chatMessage/index.tsx
@@ -1,13 +1,11 @@
 import {
-	ChatBubbleLeftEllipsisIcon,
 	ChatBubbleOvalLeftEllipsisIcon,
-	PlusSmallIcon,
 } from "@heroicons/react/24/outline";
 import { useState } from "react";
 import { ChatMessageType } from "../../../types/chat";
 import { classNames } from "../../../utils";
+import AiIcon from  "../../../assets/Gooey Ring-5s-271px.svg"
 import { UserIcon } from "@heroicons/react/24/solid";
-import {AiFillRobot} from "react-icons/ai"
 var Convert = require("ansi-to-html");
 var convert = new Convert({ newline: true });
 
@@ -22,11 +20,11 @@ export default function ChatMessage({ chat }: { chat: ChatMessageType }) {
 		>
 			<div
 				className={classNames(
-					"rounded-full w-8 h-8 flex items-center my-3 justify-center",chat.isSend?"bg-black":"bg-black"
+					"rounded-full w-9 h-9 flex items-center my-3 justify-center",chat.isSend?"bg-gray-200":"bg-gray-200"
 				)}
 			>
-				{!chat.isSend && <object data={"../"}></object>}
-				{chat.isSend && <img src={"../../../assets/Gooey Ring-5s-271px.svg"}/>}
+				{!chat.isSend && <img className="scale-150" src={AiIcon}/>}
+				{chat.isSend && <UserIcon/>}
 			</div>
 			{!chat.isSend ? (
 				<div className="w-full text-start flex items-center">
diff --git a/src/frontend/src/modals/chatModal/index.tsx b/src/frontend/src/modals/chatModal/index.tsx
index c1a0645ad..f0783a744 100644
--- a/src/frontend/src/modals/chatModal/index.tsx
+++ b/src/frontend/src/modals/chatModal/index.tsx
@@ -58,7 +58,7 @@ export default function ChatModal({ flow, open, setOpen }:{open:boolean,setOpen:
 	};
 
 	useEffect(() => {
-		const newWs = new WebSocket(`ws://backend:7860/chat/${flow.id}`);
+		const newWs = new WebSocket(`ws://localhost:7860/chat/${flow.id}`);
 		newWs.onopen = () => {
 		  console.log('WebSocket connection established!');
 		};
@@ -133,7 +133,6 @@ export default function ChatModal({ flow, open, setOpen }:{open:boolean,setOpen:
 	}
 
 	function validateNodes() {
-		console.log(reactFlowInstance);
 		return reactFlowInstance
 			.getNodes()
 			.flatMap((n: NodeType) => validateNode(n));
diff --git a/src/frontend/src/svg.d.ts b/src/frontend/src/svg.d.ts
new file mode 100644
index 000000000..a6e109e02
--- /dev/null
+++ b/src/frontend/src/svg.d.ts
@@ -0,0 +1,5 @@
+declare module '*.svg' {
+    const content: any;
+    export default content;
+  }
+  
\ No newline at end of file