From 2add242a187713c3e3e802e86b0c12e7c6dc9802 Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Tue, 27 Feb 2024 16:19:37 -0300
Subject: [PATCH] Update FlowStoreType to include ChatOutputType and
 chatInputType

---
 src/frontend/src/stores/flowStore.ts         | 14 +++++++++++---
 src/frontend/src/types/zustand/flow/index.ts |  4 ++--
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 7f0f88a5f..351997aaf 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -19,7 +19,7 @@ import {
   sourceHandleType,
   targetHandleType,
 } from "../types/flow";
-import { FlowPoolObjectType, FlowStoreType } from "../types/zustand/flow";
+import { ChatOutputType, FlowPoolObjectType, FlowStoreType, chatInputType } from "../types/zustand/flow";
 import { buildVertices } from "../utils/buildUtils";
 import {
   cleanEdges,
@@ -59,7 +59,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     }
     get().setFlowPool(newFlowPool);
   },
-  updateFlowPool:(nodeId:string,data:FlowPoolObjectType,buildId?:string)=>{
+  updateFlowPool:(nodeId:string,data:FlowPoolObjectType| ChatOutputType | chatInputType,buildId?:string)=>{
     let newFlowPool = cloneDeep({ ...get().flowPool });
     if (!newFlowPool[nodeId]){
       return;
@@ -69,7 +69,15 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
       if(buildId){
         index = newFlowPool[nodeId].findIndex((flow)=>flow.id===buildId);
       }
-      newFlowPool[nodeId][index] = data;
+      //check if the data is a flowpool object
+      if((data as FlowPoolObjectType).data?.artifacts!==undefined){
+        newFlowPool[nodeId][index] = (data as FlowPoolObjectType);
+      }
+      //update data artifact
+      else
+      {
+        newFlowPool[nodeId][index].data.artifacts = data;
+      }
     }
     get().setFlowPool(newFlowPool);
   },
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 01dbbf009..8b6093b51 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -24,7 +24,7 @@ export type FlowPoolObjectType = {
   timestamp: string;
   valid: boolean;
   params: any;
-  data: { artifacts: any; results: any | ChatOutputType | chatInputType };
+  data: { artifacts: any | ChatOutputType | chatInputType; results: any | ChatOutputType | chatInputType };
   duration: string;
   progress: number;
   id: string;
@@ -92,5 +92,5 @@ export type FlowStoreType = {
   updateBuildStatus: (nodeId: string[], status: BuildStatus) => void;
   revertBuiltStatusFromBuilding: () => void;
   flowBuildStatus: { [key: string]: BuildStatus };
-  updateFlowPool: (nodeId:string, data:FlowPoolObjectType,buildId?:string) => void;
+  updateFlowPool: (nodeId:string, data:FlowPoolObjectType | ChatOutputType | chatInputType,buildId?:string) => void;
 };