From 5de6d5de98ba6a5bc4d5ac3c9a9878f6f73e38a0 Mon Sep 17 00:00:00 2001
From: nick-delirium <nikita@openreplay.com>
Date: Mon, 19 May 2025 15:33:57 +0200
Subject: [PATCH] ui: visualization prep for kai

---
 frontend/app/components/Kai/KaiService.ts    |  4 ++-
 frontend/app/components/Kai/KaiStore.ts      | 30 +++++++++++++++++---
 frontend/app/components/Kai/SocketManager.ts | 21 ++++++--------
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/frontend/app/components/Kai/KaiService.ts b/frontend/app/components/Kai/KaiService.ts
index 045553b4e..1ca481e2d 100644
--- a/frontend/app/components/Kai/KaiService.ts
+++ b/frontend/app/components/Kai/KaiService.ts
@@ -31,8 +31,10 @@ export default class KaiService extends AiService {
       role: string;
       content: string;
       message_id: any;
-      duration?: number;
+      duration: number;
       feedback: boolean | null;
+      supports_visualization: boolean;
+      chart: 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 6c84d5a62..38e622ff1 100644
--- a/frontend/app/components/Kai/KaiStore.ts
+++ b/frontend/app/components/Kai/KaiStore.ts
@@ -1,12 +1,25 @@
 import { makeAutoObservable, runInAction } from 'mobx';
-import { BotChunk, ChatManager, Message } from './SocketManager';
+import { BotChunk, ChatManager } from './SocketManager';
 import { kaiService as aiService, kaiService } from 'App/services';
 import { toast } from 'react-toastify';
 
+export interface Message {
+  text: string;
+  isUser: boolean;
+  messageId: string;
+  chart: string;
+  supports_visualization: boolean;
+  feedback: boolean | null;
+  duration: number;
+}
+export interface SentMessage extends Omit<Message, 'duration' | 'feedback' | 'chart' | 'supports_visualization'> {
+  replace: boolean;
+}
+
 class KaiStore {
   chatManager: ChatManager | null = null;
   processingStage: BotChunk | null = null;
-  messages: Message[] = [];
+  messages: Array<Message> = [];
   queryText = '';
   loadingChat = false;
   replacing = false;
@@ -100,6 +113,8 @@ class KaiStore {
               messageId: m.message_id,
               duration: m.duration,
               feedback: m.feedback,
+              chart: m.chart,
+              supports_visualization: m.supports_visualization,
             };
           }),
         );
@@ -125,18 +140,21 @@ class KaiStore {
     this.chatManager = new ChatManager({ ...settings, token });
     this.chatManager.setOnMsgHook({
       msgCallback: (msg) => {
-        if ('state' in msg) {
+        if (msg.type === 'state') {
           if (msg.state === 'running') {
             this.setProcessingStage({
               content: 'Processing your request...',
               stage: 'chart',
               messageId: Date.now().toPrecision(),
               duration: msg.start_time ? Date.now() - msg.start_time : 0,
+              type: 'chunk',
+              supports_visualization: false,
             });
           } else {
             this.setProcessingStage(null);
           }
-        } else {
+        }
+        if (msg.type === 'chunk') {
           if (msg.stage === 'start') {
             this.setProcessingStage({
               ...msg,
@@ -153,6 +171,8 @@ class KaiStore {
               messageId: msg.messageId,
               duration: msg.duration,
               feedback: null,
+              chart: '',
+              supports_visualization: msg.supports_visualization,
             };
             this.addMessage(msgObj);
             this.setProcessingStage(null);
@@ -197,6 +217,8 @@ class KaiStore {
       messageId: Date.now().toString(),
       feedback: null,
       duration: 0,
+      supports_visualization: false,
+      chart: '',
     });
   };
 
diff --git a/frontend/app/components/Kai/SocketManager.ts b/frontend/app/components/Kai/SocketManager.ts
index e34ccc155..7ba184a33 100644
--- a/frontend/app/components/Kai/SocketManager.ts
+++ b/frontend/app/components/Kai/SocketManager.ts
@@ -74,12 +74,12 @@ export class ChatManager {
     titleCallback,
   }: {
     msgCallback: (
-      msg: BotChunk | { state: string; type: 'state'; start_time?: number },
+      msg: StateEvent | BotChunk,
     ) => void;
     titleCallback: (title: string) => void;
   }) => {
     this.socket.on('chunk', (msg: BotChunk) => {
-      msgCallback(msg);
+      msgCallback({ ...msg, type: 'chunk' });
     });
     this.socket.on('title', (msg: { content: string }) => {
       titleCallback(msg.content);
@@ -105,16 +105,13 @@ export interface BotChunk {
   stage: 'start' | 'chart' | 'final' | 'title';
   content: string;
   messageId: string;
-  duration?: number;
-}
-export interface Message {
-  text: string;
-  isUser: boolean;
-  messageId: string;
-  duration?: number;
-  feedback: boolean | null;
+  duration: number;
+  supports_visualization: boolean;
+  type: 'chunk'
 }
 
-export interface SentMessage extends Message {
-  replace: boolean;
+interface StateEvent {
+  state: string;
+  start_time?: number;
+  type: 'state';
 }