diff --git a/src/frontend/src/modals/IOModal/components/chatView/index.tsx b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
index 8e7fb2cda..f183c830e 100644
--- a/src/frontend/src/modals/IOModal/components/chatView/index.tsx
+++ b/src/frontend/src/modals/IOModal/components/chatView/index.tsx
@@ -62,14 +62,11 @@ export default function ChatView({
     const chatMessages: ChatMessageType[] = chatOutputResponses
       .sort((a, b) => Date.parse(a.timestamp) - Date.parse(b.timestamp))
       //
-      .filter(
-        (output) =>
-          output?.data?.messages && output?.data?.messages?.length > 0,
-      )
+      .filter((output) => output.data.message)
       .map((output, index) => {
         try {
           const { sender, message, sender_name, stream_url, files } = output
-            ?.data?.messages[0] as ChatOutputType;
+            .data.message as ChatOutputType;
 
           const is_ai = sender === "Machine" || sender === null;
           return {
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 86646d15c..d966868a2 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -90,7 +90,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       }
       //update data results
       else {
-        newFlowPool[nodeId][index].data.messages[0] = data as
+        newFlowPool[nodeId][index].data.message = data as
           | ChatOutputType
           | ChatInputType;
       }
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index ca704a137..86dd13c87 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -167,6 +167,7 @@ export type VertexDataTypeAPI = {
   timedelta?: number;
   duration?: string;
   artifacts?: any;
+  message: ChatOutputType | ChatInputType;
 };
 
 export type CodeErrorDataTypeAPI = {