From 6e5115985ea6f0f18d9e8ac1096f67b3530184d9 Mon Sep 17 00:00:00 2001
From: paulclindo <paulclindo@gmail.com>
Date: Fri, 4 Oct 2024 16:47:26 -0500
Subject: [PATCH] add playground flag

---
 .../src/components/chat/message-stream.tsx    |  1 +
 .../src/components/playground/baml-editor.tsx |  1 +
 .../components/playground/workflow-editor.tsx |  1 +
 libs/shinkai-message-ts/src/api/jobs/types.ts | 11 +++++++---
 .../src/v2/mutations/createJob/index.ts       | 22 ++++++++++++++++++-
 .../src/v2/mutations/createJob/types.ts       |  1 +
 6 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/apps/shinkai-desktop/src/components/chat/message-stream.tsx b/apps/shinkai-desktop/src/components/chat/message-stream.tsx
index d8673636d..68d3f80e0 100644
--- a/apps/shinkai-desktop/src/components/chat/message-stream.tsx
+++ b/apps/shinkai-desktop/src/components/chat/message-stream.tsx
@@ -138,6 +138,7 @@ export const useWebSocketMessage = ({
     if (lastMessage?.data) {
       try {
         const parseData: WsMessage = JSON.parse(lastMessage.data);
+        console.log(parseData, 'parseData');
         if (parseData.message_type !== 'Stream' || parseData.inbox !== inboxId)
           return;
         isStreamingFinished.current = false;
diff --git a/apps/shinkai-desktop/src/components/playground/baml-editor.tsx b/apps/shinkai-desktop/src/components/playground/baml-editor.tsx
index 3fffbe8be..8db5dd4b4 100644
--- a/apps/shinkai-desktop/src/components/playground/baml-editor.tsx
+++ b/apps/shinkai-desktop/src/components/playground/baml-editor.tsx
@@ -241,6 +241,7 @@ function BamlEditor() {
       content: escapedBamlInput,
       workflowCode,
       isHidden: true,
+      playground: true,
       chatConfig: {
         stream: false,
         custom_prompt: '',
diff --git a/apps/shinkai-desktop/src/components/playground/workflow-editor.tsx b/apps/shinkai-desktop/src/components/playground/workflow-editor.tsx
index 388db4e29..df8b464a5 100644
--- a/apps/shinkai-desktop/src/components/playground/workflow-editor.tsx
+++ b/apps/shinkai-desktop/src/components/playground/workflow-editor.tsx
@@ -232,6 +232,7 @@ function WorkflowEditor() {
       isHidden: true,
       selectedVRFiles,
       selectedVRFolders,
+      playground: true,
     });
   };
   const { mutateAsync: createWorkflow, isPending: isCreateWorkflowPending } =
diff --git a/libs/shinkai-message-ts/src/api/jobs/types.ts b/libs/shinkai-message-ts/src/api/jobs/types.ts
index 9e0353311..2beee6337 100644
--- a/libs/shinkai-message-ts/src/api/jobs/types.ts
+++ b/libs/shinkai-message-ts/src/api/jobs/types.ts
@@ -79,9 +79,14 @@ export type CreateJobRequest = {
       local_vrpack: [];
       local_vrkai: [];
     };
-    associated_ui: {
-      Sheet: string;
-    } | null;
+    associated_ui:
+      | null
+      | {
+          Sheet: string | null;
+        }
+      | {
+          Playground: '';
+        };
     is_hidden: boolean;
   };
 };
diff --git a/libs/shinkai-node-state/src/v2/mutations/createJob/index.ts b/libs/shinkai-node-state/src/v2/mutations/createJob/index.ts
index ecb063e2b..352fd2b4a 100644
--- a/libs/shinkai-node-state/src/v2/mutations/createJob/index.ts
+++ b/libs/shinkai-node-state/src/v2/mutations/createJob/index.ts
@@ -20,7 +20,27 @@ export const createJob = async ({
   selectedVRFiles,
   selectedVRFolders,
   chatConfig,
+  playground,
 }: CreateJobInput) => {
+  let associatedUI:
+    | null
+    | {
+        Sheet: string | null;
+      }
+    | {
+        Playground: '';
+      } = null;
+
+  if (playground) {
+    associatedUI = {
+      Playground: '',
+    };
+  } else if (sheetId) {
+    associatedUI = {
+      Sheet: sheetId,
+    };
+  }
+
   const { job_id: jobId } = await createJobApi(nodeAddress, token, {
     llm_provider: llmProvider,
     job_creation_info: {
@@ -31,7 +51,7 @@ export const createJob = async ({
         local_vrkai: [],
         network_folders: [],
       },
-      associated_ui: sheetId ? { Sheet: sheetId } : null,
+      associated_ui: associatedUI,
       is_hidden: isHidden,
     },
   });
diff --git a/libs/shinkai-node-state/src/v2/mutations/createJob/types.ts b/libs/shinkai-node-state/src/v2/mutations/createJob/types.ts
index ff793d2bb..4c61cca03 100644
--- a/libs/shinkai-node-state/src/v2/mutations/createJob/types.ts
+++ b/libs/shinkai-node-state/src/v2/mutations/createJob/types.ts
@@ -17,6 +17,7 @@ export type CreateJobInput = Token & {
   selectedVRFiles?: VectorFSItemScopeEntry[];
   selectedVRFolders?: VectorFSFolderScopeEntry[];
   chatConfig?: JobConfig;
+  playground?: boolean;
 };
 
 export type CreateJobOutput = {