diff --git a/frontend/app/components/Kai/KaiChat.tsx b/frontend/app/components/Kai/KaiChat.tsx
index 904181460..0a53ab16d 100644
--- a/frontend/app/components/Kai/KaiChat.tsx
+++ b/frontend/app/components/Kai/KaiChat.tsx
@@ -1,16 +1,15 @@
 import React from 'react';
 import { useModal } from 'App/components/Modal';
-import { MessagesSquare, Trash } from 'lucide-react';
 import ChatHeader from './components/ChatHeader';
 import { PANEL_SIZES } from 'App/constants/panelSizes';
 import ChatLog from './components/ChatLog';
 import IntroSection from './components/IntroSection';
-import { useQuery } from '@tanstack/react-query';
 import { kaiService } from 'App/services';
 import { toast } from 'react-toastify';
 import { useStore } from 'App/mstore';
 import { observer } from 'mobx-react-lite';
 import { useHistory, useLocation } from 'react-router-dom';
+import ChatsModal from './components/ChatsModal';
 
 function KaiChat() {
   const { userStore, projectsStore } = useStore();
@@ -133,69 +132,4 @@ function KaiChat() {
   );
 }
 
-function ChatsModal({
-  onSelect,
-  projectId,
-}: {
-  onSelect: (threadId: string, title: string) => void;
-  projectId: string;
-}) {
-  const {
-    data = [],
-    isPending,
-    refetch,
-  } = useQuery({
-    queryKey: ['kai', 'chats', projectId],
-    queryFn: () => kaiService.getKaiChats(projectId),
-    staleTime: 1000 * 60,
-  });
-
-  const onDelete = async (id: string) => {
-    try {
-      await kaiService.deleteKaiChat(projectId, id);
-    } catch (e) {
-      toast.error("Something wen't wrong. Please try again later.");
-    }
-    refetch();
-  };
-  return (
-    <div className={'h-screen w-full flex flex-col gap-2 p-4'}>
-      <div className={'flex items-center font-semibold text-lg gap-2'}>
-        <MessagesSquare size={16} />
-        <span>Chats</span>
-      </div>
-      {isPending ? (
-        <div className="animate-pulse text-disabled-text">Loading chats...</div>
-      ) : (
-        <div className="flex flex-col overflow-y-auto -mx-4 px-4">
-          {data.map((chat) => (
-            <div
-              key={chat.thread_id}
-              className="flex items-center relative group min-h-8"
-            >
-              <div
-                style={{ width: 270 - 28 - 4 }}
-                className="rounded-l pl-2 h-full w-full hover:bg-active-blue flex items-center"
-              >
-                <div
-                  onClick={() => onSelect(chat.thread_id, chat.title)}
-                  className="cursor-pointer hover:underline truncate"
-                >
-                  {chat.title}
-                </div>
-              </div>
-              <div
-                onClick={() => onDelete(chat.thread_id)}
-                className="cursor-pointer opacity-0 group-hover:opacity-100 rounded-r h-full px-2 flex items-center group-hover:bg-active-blue"
-              >
-                <Trash size={14} className="text-disabled-text" />
-              </div>
-            </div>
-          ))}
-        </div>
-      )}
-    </div>
-  );
-}
-
 export default observer(KaiChat);
diff --git a/frontend/app/components/Kai/KaiService.ts b/frontend/app/components/Kai/KaiService.ts
index 50b1ccf9a..468b42550 100644
--- a/frontend/app/components/Kai/KaiService.ts
+++ b/frontend/app/components/Kai/KaiService.ts
@@ -3,7 +3,7 @@ import AiService from '@/services/AiService';
 export default class KaiService extends AiService {
   getKaiChats = async (
     projectId: string,
-  ): Promise<{ title: string; thread_id: string }[]> => {
+  ): Promise<{ title: string; thread_id: string, datetime: string }[]> => {
     const r = await this.client.get(`/kai/${projectId}/chats`);
     if (!r.ok) {
       throw new Error('Failed to fetch chats');
@@ -80,4 +80,25 @@ export default class KaiService extends AiService {
     const data = await r.json();
     return data;
   };
+
+  getMsgChart = async (messageId: string, projectId: string): Promise<{ filters: any[], chart: string, eventsOrder: string }> => {
+    const r = await this.client.get(`/kai/${projectId}/chats/data/${messageId}`);
+    if (!r.ok) {
+      throw new Error('Failed to fetch chart data');
+    }
+    const data = await r.json();
+    return data;
+  }
+
+  saveChartData = async (messageId: string, projectId: string, chartData: any) => {
+    const r = await this.client.post(`/kai/${projectId}/chats/data/${messageId}`, {
+      chart_data: JSON.stringify(chartData),
+    });
+    if (!r.ok) {
+      throw new Error('Failed to save chart data');
+    }
+
+    const data = await r.json();
+    return data;
+  }
 }
diff --git a/frontend/app/components/Kai/KaiStore.ts b/frontend/app/components/Kai/KaiStore.ts
index 3208a403c..b5660faf0 100644
--- a/frontend/app/components/Kai/KaiStore.ts
+++ b/frontend/app/components/Kai/KaiStore.ts
@@ -2,6 +2,7 @@ import { makeAutoObservable, runInAction } from 'mobx';
 import { BotChunk, ChatManager } from './SocketManager';
 import { kaiService as aiService, kaiService } from 'App/services';
 import { toast } from 'react-toastify';
+import Widget from 'App/mstore/types/widget';
 
 export interface Message {
   text: string;
@@ -15,7 +16,11 @@ export interface Message {
   feedback: boolean | null;
   duration: number;
 }
-export interface SentMessage extends Omit<Message, 'duration' | 'feedback' | 'chart' | 'supports_visualization'> {
+export interface SentMessage
+  extends Omit<
+    Message,
+    'duration' | 'feedback' | 'chart' | 'supports_visualization'
+  > {
   replace: boolean;
 }
 
@@ -25,7 +30,7 @@ class KaiStore {
   messages: Array<Message> = [];
   queryText = '';
   loadingChat = false;
-  replacing = false;
+  replacing: string | null = null;
 
   constructor() {
     makeAutoObservable(this);
@@ -83,9 +88,9 @@ class KaiStore {
     this.messages.push(message);
   };
 
-  editMessage = (text: string) => {
+  editMessage = (text: string, messageId: string) => {
     this.setQueryText(text);
-    this.setReplacing(true);
+    this.setReplacing(messageId);
   };
 
   replaceAtIndex = (message: Message, index: number) => {
@@ -192,13 +197,13 @@ class KaiStore {
     }
   };
 
-  setReplacing = (replacing: boolean) => {
+  setReplacing = (replacing: string | null) => {
     this.replacing = replacing;
   };
 
   sendMessage = (message: string) => {
     if (this.chatManager) {
-      this.chatManager.sendMessage(message, this.replacing);
+      this.chatManager.sendMessage(message, !!this.replacing);
     }
     if (this.replacing) {
       console.log(
@@ -280,7 +285,50 @@ class KaiStore {
     }
   };
 
-  getMessageChart = (msgId: string, projectId: string) => { }
+  getMessageChart = async (msgId: string, projectId: string) => {
+    this.setProcessingStage({
+      content: 'Generating visualization...',
+      stage: 'chart',
+      messageId: msgId,
+      duration: 0,
+      type: 'chunk',
+      supports_visualization: false,
+    });
+    try {
+      const filters = await kaiService.getMsgChart(msgId, projectId);
+      const data = {
+        metricId: undefined,
+        dashboardId: undefined,
+        widgetId: undefined,
+        metricOf: undefined,
+        metricType: undefined,
+        metricFormat: undefined,
+        viewType: undefined,
+        name: 'Kai Visualization',
+        series: [
+          {
+            name: 'Kai Visualization',
+            filter: {
+              eventsOrder: filters.eventsOrder,
+              filters: filters.filters,
+            },
+          },
+        ],
+      };
+      const metric = new Widget().fromJson(data);
+      return metric;
+    } catch (e) {
+      console.error(e);
+      throw new Error('Failed to generate visualization');
+    } finally {
+      this.setProcessingStage(null);
+    }
+  };
+
+  getParsedChart = (data: string) => {
+    const parsedData = JSON.parse(data);
+    return new Widget().fromJson(parsedData);
+  };
 }
 
 export const kaiStore = new KaiStore();
diff --git a/frontend/app/components/Kai/components/ChatInput.tsx b/frontend/app/components/Kai/components/ChatInput.tsx
index 94bfeb89f..49c8739f6 100644
--- a/frontend/app/components/Kai/components/ChatInput.tsx
+++ b/frontend/app/components/Kai/components/ChatInput.tsx
@@ -1,55 +1,96 @@
-import React from 'react'
-import { Button, Input } from "antd";
-import { SendHorizonal, OctagonX } from "lucide-react";
-import { kaiStore } from "../KaiStore";
-import { observer } from "mobx-react-lite";
+import React from 'react';
+import { Button, Input, Tooltip } from 'antd';
+import { SendHorizonal, OctagonX } from 'lucide-react';
+import { kaiStore } from '../KaiStore';
+import { observer } from 'mobx-react-lite';
 
-function ChatInput({ isLoading, onSubmit, threadId }: { isLoading?: boolean, onSubmit: (str: string) => void, threadId: string }) {
+function ChatInput({
+  isLoading,
+  onSubmit,
+  threadId,
+}: {
+  isLoading?: boolean;
+  onSubmit: (str: string) => void;
+  threadId: string;
+}) {
   const inputRef = React.useRef<Input>(null);
   const inputValue = kaiStore.queryText;
-  const isProcessing = kaiStore.processingStage !== null
+  const isProcessing = kaiStore.processingStage !== null;
   const setInputValue = (text: string) => {
-    kaiStore.setQueryText(text)
-  }
+    kaiStore.setQueryText(text);
+  };
 
   const submit = () => {
     if (isProcessing) {
-      const settings = { projectId: '2325', userId: '0', threadId, };
-      void kaiStore.cancelGeneration(settings)
+      const settings = { projectId: '2325', userId: '0', threadId };
+      void kaiStore.cancelGeneration(settings);
     } else {
       if (inputValue.length > 0) {
-        onSubmit(inputValue)
-        setInputValue('')
+        onSubmit(inputValue);
+        setInputValue('');
       }
     }
-  }
+  };
+
+  const cancelReplace = () => {
+    setInputValue('');
+    kaiStore.setReplacing(null);
+  };
 
   React.useEffect(() => {
     if (inputRef.current) {
-      inputRef.current.focus()
+      inputRef.current.focus();
     }
-  }, [inputValue])
+  }, [inputValue]);
+
+  const isReplacing = kaiStore.replacing !== null;
 
   return (
     <Input
       onPressEnter={submit}
+      onKeyDown={(e) => {
+        if (e.key === 'Escape') {
+          cancelReplace();
+        }
+      }}
       ref={inputRef}
       placeholder={'Ask anything about your product and users...'}
       size={'large'}
       value={inputValue}
       onChange={(e) => setInputValue(e.target.value)}
       suffix={
-        <Button
-          loading={isLoading}
-          onClick={submit}
-          icon={isProcessing ? <OctagonX size={16} /> : <SendHorizonal size={16} />}
-          type={'text'}
-          size={'small'}
-          shape={'circle'}
-        />
+        <>
+          {isReplacing ? (
+            <Tooltip title={'Cancel replacing'}>
+              <Button
+                onClick={cancelReplace}
+                icon={<OctagonX size={16} />}
+                type={'text'}
+                size={'small'}
+                shape={'circle'}
+              />
+            </Tooltip>
+          ) : null}
+          <Tooltip title={'Send message'}>
+            <Button
+              loading={isLoading}
+              onClick={submit}
+              icon={
+                isProcessing ? (
+                  <OctagonX size={16} />
+                ) : (
+                  <SendHorizonal size={16} />
+                )
+              }
+              type={'text'}
+              size={'small'}
+              shape={'circle'}
+            />
+          </Tooltip>
+        </>
       }
     />
-  )
+  );
 }
 
-export default observer(ChatInput)
+export default observer(ChatInput);
diff --git a/frontend/app/components/Kai/components/ChatLog.tsx b/frontend/app/components/Kai/components/ChatLog.tsx
index 50708dad7..9b015ad22 100644
--- a/frontend/app/components/Kai/components/ChatLog.tsx
+++ b/frontend/app/components/Kai/components/ChatLog.tsx
@@ -1,6 +1,6 @@
 import React from 'react';
 import ChatInput from './ChatInput';
-import { ChatMsg, ChatNotice } from './ChatMsg';
+import ChatMsg, { ChatNotice } from './ChatMsg';
 import { Loader } from 'UI';
 import { kaiStore } from '../KaiStore';
 import { observer } from 'mobx-react-lite';
diff --git a/frontend/app/components/Kai/components/ChatMsg.tsx b/frontend/app/components/Kai/components/ChatMsg.tsx
index 4c5100337..04406840b 100644
--- a/frontend/app/components/Kai/components/ChatMsg.tsx
+++ b/frontend/app/components/Kai/components/ChatMsg.tsx
@@ -1,5 +1,6 @@
 import React from 'react';
 import { Icon, CopyButton } from 'UI';
+import { observer } from 'mobx-react-lite';
 import cn from 'classnames';
 import Markdown from 'react-markdown';
 import remarkGfm from 'remark-gfm';
@@ -16,8 +17,11 @@ import { Button, Tooltip } from 'antd';
 import { kaiStore, Message } from '../KaiStore';
 import { toast } from 'react-toastify';
 import { durationFormatted } from 'App/date';
+import WidgetChart from '@/components/Dashboard/components/WidgetChart';
+import Widget from 'App/mstore/types/widget';
+import cn from 'classnames';
 
-export function ChatMsg({
+function ChatMsg({
   userName,
   siteId,
   canEdit,
@@ -28,12 +32,22 @@ export function ChatMsg({
   canEdit?: boolean;
   siteId: string;
 }) {
+  const [metric, setMetric] = React.useState<Widget | null>(null);
   const [loadingChart, setLoadingChart] = React.useState(false);
-  const { text, isUser, messageId, duration, feedback, supports_visualization, chart_data } = message;
+  const {
+    text,
+    isUser,
+    messageId,
+    duration,
+    feedback,
+    supports_visualization,
+    chart_data,
+  } = message;
+  const isEditing = kaiStore.replacing && messageId === kaiStore.replacing;
   const [isProcessing, setIsProcessing] = React.useState(false);
   const bodyRef = React.useRef<HTMLDivElement>(null);
   const onRetry = () => {
-    kaiStore.editMessage(text);
+    kaiStore.editMessage(text, messageId);
   };
   const onFeedback = (feedback: 'like' | 'dislike', messageId: string) => {
     kaiStore.sendMsgFeedback(feedback, messageId, siteId);
@@ -70,11 +84,24 @@ export function ChatMsg({
       });
   };
 
-  const getChart = () => {
-    setLoadingChart(true);
-    kaiStore.getMessageChart(messageId, siteId)
-    setLoadingChart(false);
-  }
+  React.useEffect(() => {
+    if (chart_data) {
+      const metric = kaiStore.getParsedChart(chart_data);
+      setMetric(metric);
+    }
+  }, [chart_data]);
+
+  const getChart = async () => {
+    try {
+      setLoadingChart(true);
+      const metric = await kaiStore.getMessageChart(messageId, siteId);
+      setMetric(metric);
+    } catch (e) {
+      toast.error(e.message);
+    } finally {
+      setLoadingChart(false);
+    }
+  };
   return (
     <div
       className={cn(
@@ -85,7 +112,7 @@ export function ChatMsg({
       {isUser ? (
         <div
           className={
-            'rounded-full bg-main text-white min-w-8 min-h-8 flex items-center justify-center sticky top-0'
+            'rounded-full bg-main text-white min-w-8 min-h-8 flex items-center justify-center sticky top-0 shadow'
           }
         >
           <span className={'font-semibold'}>{userName}</span>
@@ -93,16 +120,28 @@ export function ChatMsg({
       ) : (
         <div
           className={
-            'rounded-full bg-white shadow min-w-8 min-h-8 flex items-center justify-center sticky top-0'
+            'rounded-full bg-gray-lightest shadow min-w-8 min-h-8 flex items-center justify-center sticky top-0'
           }
         >
           <Icon name={'kai_colored'} size={18} />
         </div>
       )}
       <div className={'mt-1 flex flex-col'}>
-        <div className="markdown-body" data-openreplay-obscured ref={bodyRef}>
+        <div
+          className={cn(
+            'markdown-body',
+            isEditing ? 'border-l border-l-main pl-2' : '',
+          )}
+          data-openreplay-obscured
+          ref={bodyRef}
+        >
           <Markdown remarkPlugins={[remarkGfm]}>{text}</Markdown>
         </div>
+        {metric ? (
+          <div className="p-2 border-gray-light rounded-lg shadow">
+            <WidgetChart metric={metric} isPreview />
+          </div>
+        ) : null}
         {isUser ? (
           canEdit ? (
             <div
@@ -133,13 +172,15 @@ export function ChatMsg({
             >
               <ThumbsDown size={16} />
             </IconButton>
-            <IconButton
-              tooltip="Visualize this answer"
-              onClick={getChart}
-              processing={loadingChart}
-            >
-              <ChartLine size={16}/>
-            </IconButton>
+            {supports_visualization ? (
+              <IconButton
+                tooltip="Visualize this answer"
+                onClick={getChart}
+                processing={loadingChart}
+              >
+                <ChartLine size={16} />
+              </IconButton>
+            ) : null}
             <CopyButton
               getHtml={() => bodyRef.current?.innerHTML}
               content={text}
@@ -223,3 +264,5 @@ function MsgDuration({ duration }: { duration: number }) {
     </div>
   );
 }
+
+export default observer(ChatMsg);
diff --git a/frontend/app/components/Kai/components/ChatsModal.tsx b/frontend/app/components/Kai/components/ChatsModal.tsx
new file mode 100644
index 000000000..b6175a978
--- /dev/null
+++ b/frontend/app/components/Kai/components/ChatsModal.tsx
@@ -0,0 +1,121 @@
+import React from 'react';
+import { splitByDate } from '../utils';
+import { useQuery } from '@tanstack/react-query';
+import { MessagesSquare, Trash } from 'lucide-react';
+import { kaiService } from 'App/services';
+import { toast } from 'react-toastify';
+import { useTranslation } from 'react-i18next';
+
+function ChatsModal({
+  onSelect,
+  projectId,
+}: {
+  onSelect: (threadId: string, title: string) => void;
+  projectId: string;
+}) {
+  const { t } = useTranslation();
+  const {
+    data = [],
+    isPending,
+    refetch,
+  } = useQuery({
+    queryKey: ['kai', 'chats', projectId],
+    queryFn: () => kaiService.getKaiChats(projectId),
+    staleTime: 1000 * 60,
+  });
+
+  const datedCollections = React.useMemo(() => {
+    return data.length ? splitByDate(data) : [];
+  }, [data.length]);
+
+  const onDelete = async (id: string) => {
+    try {
+      await kaiService.deleteKaiChat(projectId, id);
+    } catch (e) {
+      toast.error("Something wen't wrong. Please try again later.");
+    }
+    refetch();
+  };
+  return (
+    <div className={'h-screen w-full flex flex-col gap-2 p-4'}>
+      <div className={'flex items-center font-semibold text-lg gap-2'}>
+        <MessagesSquare size={16} />
+        <span>{t('Chats')}</span>
+      </div>
+      {isPending ? (
+        <div className="animate-pulse text-disabled-text">{t('Loading chats')}...</div>
+      ) : (
+        <div className="overflow-y-auto flex flex-col gap-2">
+          {datedCollections.map((col) => (
+            <ChatCollection
+              data={col.entries}
+              date={col.date}
+              onSelect={onSelect}
+              onDelete={onDelete}
+            />
+          ))}
+        </div>
+      )}
+    </div>
+  );
+}
+
+function ChatCollection({
+  data,
+  onSelect,
+  onDelete,
+  date,
+}: {
+  data: { title: string; thread_id: string }[];
+  onSelect: (threadId: string, title: string) => void;
+  onDelete: (threadId: string) => void;
+  date: string;
+}) {
+  return (
+    <div>
+      <div className="text-disabled-text">{date}</div>
+      <ChatsList data={data} onSelect={onSelect} onDelete={onDelete} />
+    </div>
+  );
+}
+
+function ChatsList({
+  data,
+  onSelect,
+  onDelete,
+}: {
+  data: { title: string; thread_id: string }[];
+  onSelect: (threadId: string, title: string) => void;
+  onDelete: (threadId: string) => void;
+}) {
+  return (
+    <div className="flex flex-col gap-1 -mx-4 px-4">
+      {data.map((chat) => (
+        <div
+          key={chat.thread_id}
+          className="flex items-center relative group min-h-7"
+        >
+          <div
+            style={{ width: 270 - 28 - 4 }}
+            className="rounded-l pl-2 min-h-7 h-full w-full hover:bg-active-blue flex items-center"
+          >
+            <div
+              onClick={() => onSelect(chat.thread_id, chat.title)}
+              className="cursor-pointer hover:underline truncate"
+            >
+              {chat.title}
+            </div>
+          </div>
+          <div
+            onClick={() => onDelete(chat.thread_id)}
+            className="cursor-pointer opacity-0 group-hover:opacity-100 rounded-r min-h-7 h-full px-2 flex items-center group-hover:bg-active-blue"
+          >
+            <Trash size={14} className="text-disabled-text" />
+          </div>
+        </div>
+      ))}
+    </div>
+  );
+}
+
+export default ChatsModal;
diff --git a/frontend/app/components/Kai/utils.ts b/frontend/app/components/Kai/utils.ts
new file mode 100644
index 000000000..691ad785f
--- /dev/null
+++ b/frontend/app/components/Kai/utils.ts
@@ -0,0 +1,36 @@
+import { DateTime } from 'luxon';
+
+type DatedEntry = {
+  date: string;
+  entries: { datetime: string }[];
+}
+
+export function splitByDate(entries: { datetime: string }[]) {
+  const today = DateTime.now().startOf('day');
+  const yesterday = today.minus({ days: 1 });
+
+  const result: DatedEntry[] = [
+    { date: 'Today', entries: [] },
+    { date: 'Yesterday', entries: [] },
+  ];
+
+  entries.forEach((ent) => {
+    const entryDate = DateTime.fromISO(ent.datetime).startOf('day');
+
+    if (entryDate.toMillis() === today.toMillis()) {
+      result[0].entries.push(ent);
+    } else if (entryDate.toMillis() === yesterday.toMillis()) {
+      result[1].entries.push(ent);
+    } else {
+      const date = entryDate.toFormat('dd LLL, yyyy')
+      const existingEntry = result.find((r) => r.date === date);
+      if (existingEntry) {
+        existingEntry.entries.push(ent);
+      } else {
+        result.push({ entries: [ent], date });
+      }
+    }
+  });
+
+  return result;
+}