From b35a9af8a6c495eb61f09922a7e0011872bd79c0 Mon Sep 17 00:00:00 2001
From: nick-delirium <nikita@openreplay.com>
Date: Tue, 20 May 2025 14:49:21 +0200
Subject: [PATCH] ui: chart btn

---
 frontend/app/components/Kai/KaiService.ts     |  1 +
 frontend/app/components/Kai/KaiStore.ts       |  8 +++++
 .../app/components/Kai/components/ChatLog.tsx | 19 +++++------
 .../app/components/Kai/components/ChatMsg.tsx | 32 ++++++++++++-------
 4 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/frontend/app/components/Kai/KaiService.ts b/frontend/app/components/Kai/KaiService.ts
index 1ca481e2d..50b1ccf9a 100644
--- a/frontend/app/components/Kai/KaiService.ts
+++ b/frontend/app/components/Kai/KaiService.ts
@@ -35,6 +35,7 @@ export default class KaiService extends AiService {
       feedback: boolean | null;
       supports_visualization: boolean;
       chart: string;
+      chart_data: string;
     }[]
   > => {
     const r = await this.client.get(`/kai/${projectId}/chats/${threadId}`);
diff --git a/frontend/app/components/Kai/KaiStore.ts b/frontend/app/components/Kai/KaiStore.ts
index 38e622ff1..3208a403c 100644
--- a/frontend/app/components/Kai/KaiStore.ts
+++ b/frontend/app/components/Kai/KaiStore.ts
@@ -7,7 +7,10 @@ export interface Message {
   text: string;
   isUser: boolean;
   messageId: string;
+  /** filters to get chart */
   chart: string;
+  /** chart data */
+  chart_data: string;
   supports_visualization: boolean;
   feedback: boolean | null;
   duration: number;
@@ -115,6 +118,7 @@ class KaiStore {
               feedback: m.feedback,
               chart: m.chart,
               supports_visualization: m.supports_visualization,
+              chart_data: m.chart_data,
             };
           }),
         );
@@ -173,6 +177,7 @@ class KaiStore {
               feedback: null,
               chart: '',
               supports_visualization: msg.supports_visualization,
+              chart_data: '',
             };
             this.addMessage(msgObj);
             this.setProcessingStage(null);
@@ -219,6 +224,7 @@ class KaiStore {
       duration: 0,
       supports_visualization: false,
       chart: '',
+      chart_data: '',
     });
   };
 
@@ -273,6 +279,8 @@ class KaiStore {
       this.chatManager = null;
     }
   };
+
+  getMessageChart = (msgId: string, projectId: string) => { }
 }
 
 export const kaiStore = new KaiStore();
diff --git a/frontend/app/components/Kai/components/ChatLog.tsx b/frontend/app/components/Kai/components/ChatLog.tsx
index 9a8bd4f90..50708dad7 100644
--- a/frontend/app/components/Kai/components/ChatLog.tsx
+++ b/frontend/app/components/Kai/components/ChatLog.tsx
@@ -61,17 +61,14 @@ function ChatLog({
       >
         <div className={'flex flex-col gap-4 w-2/3 min-h-max'}>
           {messages.map((msg, index) => (
-            <ChatMsg
-              key={index}
-              text={msg.text}
-              isUser={msg.isUser}
-              userName={userLetter}
-              messageId={msg.messageId}
-              duration={msg.duration}
-              feedback={msg.feedback}
-              siteId={projectId}
-              canEdit={processingStage === null && msg.isUser && index === lastHumanMsgInd}
-            />
+            <React.Fragment key={msg.messageId ?? index}>
+              <ChatMsg
+                userName={userLetter}
+                siteId={projectId}
+                message={msg}
+                canEdit={processingStage === null && msg.isUser && index === lastHumanMsgInd}
+              />
+            </React.Fragment>
           ))}
           {processingStage ? (
             <ChatNotice
diff --git a/frontend/app/components/Kai/components/ChatMsg.tsx b/frontend/app/components/Kai/components/ChatMsg.tsx
index 89174f368..4c5100337 100644
--- a/frontend/app/components/Kai/components/ChatMsg.tsx
+++ b/frontend/app/components/Kai/components/ChatMsg.tsx
@@ -10,31 +10,26 @@ import {
   ListRestart,
   FileDown,
   Clock,
+  ChartLine,
 } from 'lucide-react';
 import { Button, Tooltip } from 'antd';
-import { kaiStore } from '../KaiStore';
+import { kaiStore, Message } from '../KaiStore';
 import { toast } from 'react-toastify';
 import { durationFormatted } from 'App/date';
 
 export function ChatMsg({
-  text,
-  isUser,
   userName,
-  messageId,
-  duration,
-  feedback,
   siteId,
   canEdit,
+  message,
 }: {
-  text: string;
-  isUser: boolean;
-  messageId: string;
+  message: Message;
   userName?: string;
-  duration?: number;
-  feedback: boolean | null;
-  siteId: string;
   canEdit?: boolean;
+  siteId: string;
 }) {
+  const [loadingChart, setLoadingChart] = React.useState(false);
+  const { text, isUser, messageId, duration, feedback, supports_visualization, chart_data } = message;
   const [isProcessing, setIsProcessing] = React.useState(false);
   const bodyRef = React.useRef<HTMLDivElement>(null);
   const onRetry = () => {
@@ -74,6 +69,12 @@ export function ChatMsg({
         setIsProcessing(false);
       });
   };
+
+  const getChart = () => {
+    setLoadingChart(true);
+    kaiStore.getMessageChart(messageId, siteId)
+    setLoadingChart(false);
+  }
   return (
     <div
       className={cn(
@@ -132,6 +133,13 @@ export function ChatMsg({
             >
               <ThumbsDown size={16} />
             </IconButton>
+            <IconButton
+              tooltip="Visualize this answer"
+              onClick={getChart}
+              processing={loadingChart}
+            >
+              <ChartLine size={16}/>
+            </IconButton>
             <CopyButton
               getHtml={() => bodyRef.current?.innerHTML}
               content={text}