Feature/meta planner v2 (#68)

* Remove unused meta planner files

* Re-plan every time one step is done

* Stop meta-planner in case of user interrupt

* Adjust meta planner steps based on feedback

web/src/components/common/TaskUI.tsx

@@ -30,7 +30,6 @@ import { Thumbnails } from './Thumbnails'
 import { UserConfirmation } from './UserConfirmation'
 import { gdocReadSelected, gdocRead, gdocWrite, gdocImage, queryDOMSingle, readActiveSpreadsheet, getUserSelectedRange } from '../../app/rpc'
 import { forwardToTab } from '../../app/rpc'
-import { feelinLucky } from '../../app/lucky'
 import { getApp } from '../../helpers/app'
 import { JupyterNotebookState } from '../../../../apps/src/jupyter/helpers/DOMToState'
 import { querySelectorMap as jupyterQSMap } from '../../../../apps/src/jupyter/helpers/querySelectorMap'

web/src/helpers/LLM/remote.ts

@@ -52,7 +52,7 @@ export const getSuggestions = async(): Promise<string[]> => {
   return parsed.prompts;
 }
 
-export const getMetaPlan = async(text: string): Promise<string[]> => {
+export const getMetaPlan = async(text: string, steps: string[], messageHistory: string): Promise<string[]> => {
   const app = getApp()
 
   const llmSettings = {
@@ -69,27 +69,40 @@ export const getMetaPlan = async(text: string): Promise<string[]> => {
   You take the jupyter state and give a list of steps to perform to explore and analyze data.
   The steps will be taken by another agent and performed one by one. So give detailed steps.
 
-
   <JupyterAppState>
   {{ state }}
   </JupyterAppState>
 
+  <CurrentPendingSteps>
+  {{ steps }}
+  </CurrentPendingSteps>
+
+  <MessageHistory>
+  {{ messageHistory }}
+  </MessageHistory>
+
   - First, read the state of the notebook to figure out what data is being operated on
   - Then, use the JupyterAppState and the user's message to determine the goal of the user.
   - Then, give a detailed list of steps to reach the user's goal. Limit to under 7 steps always.
+  - If current pending steps are sufficient, you can return an empty list for steps.
+  - If you think that the current pending steps need to be adjusted, you can return a new list of steps.
   - There should always be a summary step at the end, with some actionable insights.
   - The output should be JSON formatted.
   
-  Sample output:
+  Sample outputs:
   If the dataframe has columns called prompt tokens, completion tokens, latency, and date, and if the user message is "I want to understand how tokens affect latency" the output could be:
   {"steps":  ["Plot the distribution of tokens",  "Plot the distribution of latency", "Plot the scatter plot of tokens vs latency", "Calculate the correlation between tokens and latency", "Plot the correlation between tokens and latency", "Perform a regression analysis on how the prompt and completion tokens affect latency", "Plot the 3d scatter plot and regression plane", "Summarize the results"]}
+
+  If current steps are sufficient, return an empty list.
+  { "steps": [] }
+
+  If current steps are not sufficient, return a new list of steps.
+  {"steps":  ["Plot the 3d & 2d scatter plot and regression plane", "Summarize the results"]}
   `
   const userMessage = text
 
   const appState = app.getState()
-  const finalSystemMessage = systemMessage.replaceAll("{{ state }}", JSON.stringify(appState))
-
-
+  const finalSystemMessage = systemMessage.replaceAll("{{ state }}", JSON.stringify(appState)).replaceAll("{{ steps }}", JSON.stringify(steps)).replaceAll("{{ messageHistory }}", messageHistory)
 
   const response = await getLLMResponse({
     messages: [{

web/src/planner/metaPlan.ts

@@ -1,33 +1,63 @@
 import chat from "../chat/chat";
 import { DefaultMessageContent } from '../state/chat/types'
-import { getState } from "../state/store"
+import { getState, RootState } from "../state/store"
 import { sleep } from "../helpers/utils"
-import _ from "lodash"
+import _, { isEmpty } from "lodash"
 import { getMetaPlan } from "../helpers/LLM/remote";
+import { ChatMessage } from "../state/chat/reducer";
+
+const getMessageHistory = () => {
+  const state = getState()
+  const thread = state.chat.activeThread
+  const messageHistory = JSON.stringify(
+    state.chat.threads[thread].messages.map((msg: ChatMessage) => ({
+      role: msg.role,
+      content: msg.content
+    }))
+  )
+  console.log('Message history is', messageHistory)
+  return messageHistory
+}
 
 export async function metaPlanner({text}: {text: string}) {
 
-  const steps = await getMetaPlan(text)
+  let steps = await getMetaPlan(text, [], getMessageHistory())
+  console.log('Initial steps are', steps)
 
-  console.log(steps, "steps")
-  for (const step of steps) {
+  while (!isEmpty(steps)) {
+    const step = steps.shift()
     const content: DefaultMessageContent = {
       type: "DEFAULT",
-      text: step,
+      //@ts-ignore
+      text: step, 
       images: []
     }
     chat.addUserMessage({content})
-    while (true){
-      await sleep(2000) // hack to avoid race condition
+    let shouldContinue = true
+    while (true) {
+      await sleep(500)
       const state = getState()
       const thread = state.chat.activeThread
       const threadStatus = state.chat.threads[thread].status
+      const isInterrupted = state.chat.threads[thread].interrupted
+      if (isInterrupted) {
+        shouldContinue = false
+        break;
+      }
       if (threadStatus === "FINISHED") {
         console.log("Thread finished!")
         break;
       }
       console.log(threadStatus)
-      await sleep(100)
     }
+    if (!shouldContinue) {
+      break;
+    }
+
+    const newSteps = await getMetaPlan(text, steps, getMessageHistory())
+    if (!isEmpty(newSteps)) {
+      steps = newSteps
+    }
+    console.log('New steps are', steps)
   }
 }

web/src/state/chat/reducer.ts

@@ -111,6 +111,7 @@ interface ChatThread {
   messages: Array<ChatMessage>
   status: ChatThreadStatus
   userConfirmation: UserConfirmationState
+  interrupted: boolean
 }
 
 interface ChatState {
@@ -130,6 +131,7 @@ const initialState: ChatState = {
     messages: [],
     status: 'FINISHED',
     userConfirmation: initialUserConfirmationState,
+    interrupted: false
   }],
   activeThread: 0,
 }
@@ -158,6 +160,7 @@ export const chatSlice = createSlice({
         createdAt: timestamp,
         updatedAt: timestamp,
       })
+      state.threads[state.activeThread].interrupted = false
     },
     deleteUserMessage: (
       state,
@@ -316,7 +319,8 @@ export const chatSlice = createSlice({
           show: false,
           content: '',
           userInput: 'NULL'
-        }
+        },
+        interrupted: false
       })
     },
     addReaction: (
@@ -376,11 +380,15 @@ export const chatSlice = createSlice({
       const userConfirmation = state.threads[state.activeThread].userConfirmation
       userConfirmation.userInput = action.payload
     },
+    abortPlan: (state) => {
+      const thread = state.activeThread
+      const activeThread = state.threads[thread]
+      activeThread.interrupted = true
+    }
   },
 })
 
-export const abortPlan = createAction('chat/abortPlan')
 // Action creators are generated for each case reducer function
-export const { addUserMessage, deleteUserMessage, addActionPlanMessage, startAction, finishAction, interruptPlan, startNewThread, addReaction, removeReaction, updateDebugChatIndex, setActiveThreadStatus, toggleUserConfirmation, setUserConfirmationInput, switchToThread } = chatSlice.actions
+export const { addUserMessage, deleteUserMessage, addActionPlanMessage, startAction, finishAction, interruptPlan, startNewThread, addReaction, removeReaction, updateDebugChatIndex, setActiveThreadStatus, toggleUserConfirmation, setUserConfirmationInput, switchToThread, abortPlan } = chatSlice.actions
 
 export default chatSlice.reducer