@stack-spot/ai-chat-widget 0.11.0 → 1.0.0-dev.1768324407795

package/src/chat-interceptors/quick-commands.ts

@@ -1,40 +1,62 @@
 import { aiClient, CancelledError, FixedChatRequest, StackspotAPIError } from '@stack-spot/portal-network'
-import { QuickCommandFetchResponseResult, QuickCommandPromptResponse2, QuickCommandResponse } from '@stack-spot/portal-network/api/ai'
+import { QuickCommandStepResult, QuickCommandPromptResponse2, QuickCommandResponse, QuickCommandScriptExecutionResponse, QuickCommandStepFetchResponse, QuickCommandStepLlmResponse } from '@stack-spot/portal-network/api/ai'
 import { Dictionary, interpolate, translate } from '@stack-spot/portal-translate'
+import { isNil } from 'lodash'
 import type { editor } from 'monaco-editor'
 import { ulid } from 'ulid'
 import { AbortedError } from '../AbortedError'
+import { quickCommandRegex } from '../regex'
 import { ChatEntry } from '../state/ChatEntry'
 import { ChatState } from '../state/ChatState'
 import { LabeledWithImage } from '../state/types'
 import { WidgetState } from '../state/WidgetState'
 import { buildConversationContext } from '../utils/chat'
 import { getSizeOfString } from '../utils/string'
-import { CUSTOM_INPUTS_KEY, CustomInputs } from './CustomInputs'
+import { CustomInputs } from './CustomInputs'
 
-type SlugExecution = Record<string, string | QuickCommandFetchResponseResult | QuickCommandPromptResponse2>
+type SlugExecution = Record<string, QuickCommandStepResult>
 
 interface QCContext {
   qc: QuickCommandResponse,
   context: Required<FixedChatRequest>['context'],
   resultMap: SlugExecution,
+  customInputs: Record<string, string>,
   chat: ChatState,
   code?: string,
   executionId: string,
   signal: AbortSignal,
+  isRemote?: boolean,
 }
 
-const commandRegex = /^\/[\w\d-_]+$/
+/**
+ * While the quick command is running, if it takes more than `progressMessageDelayMS` a progress message is shown.
+ * This progress message is shown immediately if one of the steps is of type "LLM".
+ */
 const progressMessageDelayMS = 1000
 
+/**
+ * This creates a chat interceptor for interpreting and running quick commands. We need a function to create the interceptor because, for
+ * quick commands, we need access to the full widget state and the editor, which are not provided to a chat interceptor.
+ * 
+ * The quick command interceptor intercepts all user messages (trimmed) that matches the quick command regex. It fetches the quick command
+ * from the backend, runs all of its steps and outputs the final response (in the chat or in the editor).
+ * 
+ * @param widget the full widget state.
+ * @param getEditor a function that returns the editor instance (monaco).
+ * @returns the quick commands chat interceptor.
+ */
 export function createQuickCommandInterceptor(widget: WidgetState, getEditor: () => typeof editor | undefined) {
+  /**
+   * This creates a CustomInputs instance in the current chat, which activates the `quickCommandQuestionsInterceptor`. In summary, this asks
+   * the first question, wait for all answers to be given and fills the result map in the context passed as parameter.
+   */
   async function computeCustomInputs(ctx: QCContext) {
     const { chat, qc: { custom_inputs: inputs, name, slug } } = ctx
     if (!inputs?.length) return
     chat.set('isLoading', false)
     const t = translate(dictionary)
     const customInputs = new CustomInputs(inputs, chat.get('agent'))
-    chat.interceptorMemory.set(CUSTOM_INPUTS_KEY, customInputs)
+    CustomInputs.setCustomInputsForChat(customInputs, chat)
     chat.pushMessage(new ChatEntry({
       agentType: 'bot',
       type: 'text',
@@ -43,17 +65,92 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     chat.pushMessage(customInputs.ask())
     const inputsValues = await customInputs.getValue()
     chat.set('isLoading', true)
-    chat.interceptorMemory.delete(CUSTOM_INPUTS_KEY)
-    ctx.resultMap = { ...ctx.resultMap, ...inputsValues }
+    CustomInputs.deleteCustomInputsFromChat(chat)
+    ctx.customInputs = { ...ctx.customInputs, ...inputsValues }
+    ctx.context.upload_ids = chat.getMessages()
+      .flatMap(message =>
+        (message.getValue().upload || []).map(upload => upload.id),
+      )
   }
 
-  async function runFetchStep(ctx: QCContext, stepIndex: number) {
-    const { qc: { slug, steps }, code, context, resultMap, executionId, signal } = ctx
+  /**
+   * Runs an Router step of a quick command.
+   */
+  async function runRouterStep(
+    ctx: QCContext,
+    stepIndex: number, iteration: Record<string, number>,
+    progress: { update: (index: number) => void, remove: () => void },
+  ) {
+    const { qc: { slug, steps }, code, resultMap, customInputs } = ctx
     const step = steps![stepIndex]
+    const inputData = Object.keys(customInputs).length > 0 && code ? { ...customInputs, [code]: code } : code ?? customInputs
+    try {
+      if (step.slug in iteration) {
+        iteration[step.slug] = iteration[step.slug] + 1
+      } else {
+        iteration[step.slug] = 1
+      }
+
+      const { next_step_slug } = await aiClient.calculateNextStep.mutate({
+        stepSlug: step.slug,
+        slug: slug,
+        quickCommandEvaluateStepRouterRequest: {
+          executions_count: iteration[step.slug],
+          input_data: inputData,
+          slugs_executions: resultMap,
+        },
+      })
+
+      if (next_step_slug === step.slug) {
+        return runStepsRecursively(stepIndex, progress, ctx, iteration)
+      }
+      const nextStepIndex = steps?.findIndex((step) => step.slug === next_step_slug)
+
+      if (isNil(nextStepIndex) || nextStepIndex === -1) return
+
+      return runStepsRecursively(nextStepIndex, progress, ctx, iteration)
+    }
+    catch (error: any) {
+      // eslint-disable-next-line no-console
+      console.error('Error executing QC step', error)
+    }
+  }
+
+  /**
+   * Runs a fetch step of a quick command and puts the result in the `resultMap` of the context passed as parameter.
+   */
+  async function runFetchStep(ctx: QCContext, stepIndex: number) {
+    const { qc: { slug, steps }, code, context, resultMap, customInputs, executionId, signal } = ctx
+    const step = steps![stepIndex] as QuickCommandStepFetchResponse
+    if (step.is_remote) {
+      ctx.isRemote = true
+
+      const { data } = await aiClient.fetchStepOfQuickCommandRemotely.mutate({
+        slug, stepSlug: step.slug,
+        quickCommandsExecutionRequest: {
+          code_selection: code, context, qc_execution_id: executionId,
+          slugs_executions: { ...resultMap, ...customInputs },
+        },
+      })
+
+      //data is the return of the request in the QC so we do not have full control over the response
+      //We handle the usual format with body, status_code and headers, but we might also handle other formats
+      const responseData = data as any
+      resultMap[step.slug] = {
+        status: responseData.status_code || 200,
+        data: JSON.stringify(responseData.body) ?? JSON.stringify(responseData),
+        headers: responseData.headers ?? {},
+      }
+      return
+    }
+
     const { headers, data, method, url } = await aiClient.fetchStepOfQuickCommand.mutate({
       slug,
       stepSlug: step.slug,
-      quickCommandsExecutionRequest: { input_data: code, context, slugs_executions: resultMap, qc_execution_id: executionId },
+      quickCommandsExecutionRequest: {
+        input_data: code, context, qc_execution_id: executionId,
+        slugs_executions: { ...resultMap, ...customInputs },
+      },
     }, signal)
     const body = ['get', 'head'].includes(method.toLowerCase()) ? undefined : data
     const response = await fetch(url, { headers: headers || undefined, body, method, signal })
@@ -65,71 +162,180 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     }
   }
 
+  /**
+   * Runs an LLM step of a quick command and puts the result in the `resultMap` of the context passed as parameter.
+   */
   async function runLLMStep(
-    { qc: { slug, steps }, code, context, executionId, resultMap, signal }: QCContext,
+    { qc: { slug, steps }, code, customInputs, context, executionId, resultMap, signal }: QCContext,
     stepIndex: number,
   ) {
-    const step = steps![stepIndex]
-    resultMap[step.slug] = await aiClient.llmStepOfQuickCommand.mutate({
+    const step = steps![stepIndex] as QuickCommandStepLlmResponse
+    let stepContext = context
+    if (!step.use_uploaded_files) {
+      const { upload_ids: _upload_ids, ...contextDataProps } = context
+      stepContext = { ...contextDataProps }
+    }
+
+    const stream = aiClient.streamLlmStepOfQuickCommand(
       slug,
-      stepSlug: step.slug,
-      quickCommandsExecutionRequest: { input_data: code, context, qc_execution_id: executionId, slugs_executions: resultMap },
-    }, signal)
+      step.slug,
+      {
+        input_data: code,
+        context: stepContext,
+        qc_execution_id: executionId,
+        slugs_executions: { ...resultMap, ...customInputs },
+      },
+    )
+
+    signal.addEventListener('abort', () => stream.cancel())
+
+    let finalValue: QuickCommandPromptResponse2
+    try {
+      finalValue = await stream.getValue()
+      resultMap[step.slug] = finalValue
+    } catch (error: any) {
+      // eslint-disable-next-line no-console
+      console.error('Error executing QC step', error)
+    }
   }
 
   function updateProgressMessageForStep(message: ChatEntry, qc: QuickCommandResponse, stepIndex: number) {
     const t = translate(dictionary)
     message.setValue({
       ...message.getValue(),
-      content: interpolate(t.progress, qc.name || qc.slug, stepIndex + 1, qc.steps?.[stepIndex]?.slug, qc.steps?.length),
-      updated: new Date().toISOString(),
+      content: interpolate(t.progress, qc.steps?.[stepIndex]?.slug, qc.name || qc.slug),
     })
   }
 
-  function showProgressMessage(ctx: QCContext, message: ChatEntry) {
-    const result = {
-      removeProgressMessage: () => clearTimeout(timeout),
-    }
-    const timeout = window.setTimeout(() => {
-      const chat = ctx.chat
-      chat.pushMessage(message)
-      result.removeProgressMessage = () => chat.popMessage()
-    }, progressMessageDelayMS)
-    return result
-  }
-
   function createProgressMessage(agent?: LabeledWithImage) {
     return new ChatEntry({
       agentType: 'bot',
       content: '',
       type: 'text',
       agent,
-      updated: new Date().toISOString(),
     })
   }
 
+  function showProgressMessage(ctx: QCContext) {
+    const showImmediately = ctx.qc.steps?.some(s => s.type === 'LLM')
+    const message = createProgressMessage(ctx.chat.get('agent'))
+    const controller = {
+      update: (index: number) => updateProgressMessageForStep(message, ctx.qc, index),
+      remove: () => clearTimeout(timeout),
+    }
+    const timeout = window.setTimeout(() => {
+      const chat = ctx.chat
+      chat.pushMessage(message)
+      controller.remove = () => chat.popMessage()
+    }, showImmediately ? 0 : progressMessageDelayMS)
+    return controller
+  }
+  async function getScriptStepStatus(
+    scriptExecutionId: string,
+    interval = 5000,
+    maxAttempts = 500,
+    currentAttempt = 0,
+  ): Promise<QuickCommandScriptExecutionResponse> {
+    if (currentAttempt >= maxAttempts) {
+      throw new Error('Max attempts reached in verify script status')
+    }
+    await aiClient.getStatusScriptStep.invalidate({ scriptExecutionId })
+    const response = await aiClient.getStatusScriptStep.query({ scriptExecutionId })
+    
+    if (response.status === 'success') {
+      return response
+    }
+
+    if (response.status === 'failure') {
+      throw response
+    }
+
+    await new Promise(resolve => {setTimeout(resolve, interval)})
+
+    return getScriptStepStatus(scriptExecutionId, interval, maxAttempts, currentAttempt + 1)
+  }
+
+  async function runScriptStep(
+    { qc: { slug, steps }, code, context, resultMap, customInputs, signal }: QCContext,
+    stepIndex: number,
+  ) {
+    const step = steps![stepIndex] as QuickCommandStepLlmResponse
+    let stepContext = context
+    if (!step.use_uploaded_files) {
+      const { upload_ids: _upload_ids, ...contextDataProps } = context
+      stepContext = { ...contextDataProps }
+    }
+
+    try {
+      const { script_execution_id } = await aiClient.startScriptStep.mutate({
+        stepSlug: step.slug,
+        slug: slug,
+        quickCommandStartScriptRequest: {
+          input_data: code,
+          custom_inputs: customInputs,
+          context: stepContext,
+          slugs_executions: resultMap,
+        },
+      }, signal)
+      const scriptResult = await getScriptStepStatus(script_execution_id)
+      resultMap[step.slug] = scriptResult
+    }
+    catch (error: any) {
+      throw new Error(`Failed to execute step "${step.slug}" of quick command "${slug}". Error ${error}.`)
+    }
+  }
+
+  async function runStepsRecursively(currentIndex: number, progress: { update: (index: number) => void, remove: () => void },
+    ctx: QCContext, iteration: Record<string, number>) {
+    const { qc } = ctx
+
+    if (!qc.steps || currentIndex >= qc.steps?.length) return
+    progress.update(currentIndex)
+
+    if (qc.steps[currentIndex].type === 'ROUTER') {
+      await runRouterStep(ctx, currentIndex, iteration, progress)
+      return
+    }
+
+    if (qc.steps[currentIndex].type === 'SCRIPT') {
+      await runScriptStep(ctx, currentIndex)
+    } else {
+      await (qc.steps[currentIndex].type === 'FETCH' ? runFetchStep(ctx, currentIndex) : runLLMStep(ctx, currentIndex))
+    }
+
+    const currentStep = qc.steps?.[currentIndex] as QuickCommandStepFetchResponse | QuickCommandStepLlmResponse
+    let nextIndex = currentIndex + 1
+
+    let nextStepSlug = currentStep.next_step_slug
+    const stepResult = ctx.resultMap[currentStep.slug]
+    if (stepResult && typeof stepResult !== 'string' && 'answer_status' in stepResult && !!stepResult.answer_status?.next_step_slug) {
+      nextStepSlug = stepResult.answer_status.next_step_slug
+    }
+
+    if (nextStepSlug) {
+      nextIndex = nextStepSlug === 'end' ?
+        qc.steps.length : qc.steps?.findIndex((step) => step.slug === nextStepSlug)
+    }
+    await runStepsRecursively(nextIndex, progress, ctx, iteration)
+  }
+
   async function runSteps(ctx: QCContext) {
-    const { qc, chat } = ctx
-    const progress = createProgressMessage(chat.get('agent'))
-    const { removeProgressMessage } = showProgressMessage(ctx, progress)
+    const progress = showProgressMessage(ctx)
     try {
-      for (let i = 0; i < (qc.steps?.length ?? 0); i++) {
-        updateProgressMessageForStep(progress, qc, i)
-        await (qc.steps?.[i].type === 'FETCH' ? runFetchStep(ctx, i) : runLLMStep(ctx, i))
-      }
+      await runStepsRecursively(0, progress, ctx, {})
     } finally {
-      removeProgressMessage()
+      progress.remove()
     }
   }
 
-  async function formatResult({ qc, code, executionId, context, resultMap, signal }: QCContext) {
+  async function formatResult({ qc, code, executionId, context, resultMap, customInputs, signal }: QCContext) {
     const formatted = await aiClient.formatResultOfQuickCommand.mutate({
       slug: qc.slug,
       quickCommandsExecutionRequest: {
         input_data: code,
         context,
         qc_execution_id: executionId,
-        slugs_executions: resultMap,
+        slugs_executions: { ...resultMap, ...customInputs },
       },
     }, signal)
     return formatted.result
@@ -146,6 +352,7 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
         widget.chatTabs.select(ctx.chat.id)
         ctx.context.conversation_id = ctx.chat.id
       } else {
+        // if we're running a command in the editor, we don't need to open a new tab, we just need to run it in a new conversation context.
         ctx.context.conversation_id = ulid()
       }
     }
@@ -161,8 +368,16 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     } catch (error) {
       throw error instanceof StackspotAPIError && error.status === 404 ? new Error(t.notFound) : error
     }
-    if (ctx.qc.use_selected_code && !code) {
-      throw new Error(t.requiresSelection)
+    
+    if (ctx.qc.use_selected_code && (!code && ctx.context.upload_ids?.length === 0)) {
+      widget.set('panel', 'editor')
+      ctx.chat.pushMessage(new ChatEntry({
+        agentType: 'bot',
+        type: 'text',
+        agent: ctx.chat.get('agent'),
+        content: t.requiresSelection,
+      }))
+      return 
     }
     manageConversationContext(ctx)
     await computeCustomInputs(ctx)
@@ -170,7 +385,10 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     return formatResult(ctx)
   }
 
-  async function registerAnalyticsEvent({ qc, executionId, code = '', context }: QCContext, status: string, start: number) {
+  /**
+   * This registers a quick command event in the backend (analytics).
+   */
+  async function registerAnalyticsEvent({ qc, isRemote, executionId, code = '', context }: QCContext, status: string, start: number) {
     const now = new Date().getTime()
     try {
       await aiClient.createEvent.mutate({
@@ -183,6 +401,8 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
             slug: qc.slug,
             qc_execution_id: executionId,
             id: qc.id,
+            //@ts-ignore
+            is_remote: isRemote,
           },
           code,
           context,
@@ -197,6 +417,9 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     }
   }
 
+  /**
+   * Writes the result in the chat or in the editor.
+   */
   function outputResult({ qc, chat, code }: QCContext, result: string) {
     if (qc.return_type === 'CHAT') {
       chat.pushMessage(new ChatEntry({
@@ -204,6 +427,7 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
         content: result,
         agent: chat.get('agent'),
         type: 'md',
+        updated: new Date().toISOString(),
       }))
     } else {
       const editor = getEditor()?.getEditors()[0]
@@ -215,26 +439,36 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
     }
   }
 
+  /**
+   * The chat interceptor itself.
+   * @param entry the chat entry.
+   * @param chat the chat state.
+   * @param signal a signal that can be canceled by the user.
+   * @returns false if, and only if, this entry (trimmed) matches the quick command regex. Undefined otherwise.
+   */
   async function quickCommandInterceptor(entry: ChatEntry, chat: ChatState, signal: AbortSignal) {
     const { agentType, content } = entry.getValue()
-    if (agentType !== 'user' || !commandRegex.test(content.trim())) return
+    if (!chat.get('features')?.quickCommands || agentType !== 'user' || !quickCommandRegex.test(content.trim())) return
     const t = translate(dictionary)
     const slug = content.trim().substring(1)
     const ctx: QCContext = {
       qc: { slug } as QuickCommandResponse,
       chat,
       code: chat.get('codeSelection'),
-      context: buildConversationContext(chat) ?? {},
+      context: buildConversationContext(chat, entry) ?? {},
       executionId: ulid(),
       resultMap: {},
+      customInputs: {},
       signal,
     }
     chat.set('isLoading', true)
     const start = new Date().getTime()
     try {
       const result = await runQuickCommand(ctx)
-      outputResult(ctx, result)
-      registerAnalyticsEvent(ctx, '200', start)
+      if (result) {
+        outputResult(ctx, result)
+        registerAnalyticsEvent(ctx, '200', start)
+      }
     } catch (error: any) {
       let message = error.message || `${error}`
       if (error instanceof AbortedError || error instanceof CancelledError) message = t.aborted
@@ -258,16 +492,16 @@ export function createQuickCommandInterceptor(widget: WidgetState, getEditor: ()
 
 const dictionary = {
   en: {
-    requiresSelection: 'This quick command requires some code to be selected in the editor.',
+    requiresSelection: 'Select the code in the editor next to it for this quick command to interact with it.',
     startQuestioning: 'To execute the Quick Command "$0", I\'ll need you to provide some information. Some may be mandatory, and others optional. Let\'s get started.',
-    progress: 'Running quick command "$0". Step $1 ($2) of $3.',
+    progress: 'Running step "$0" from Quick Command "$1".',
     aborted: 'The quick command execution aborted by the user.',
     notFound: 'There\'s no quick command with the provided name. If you don\'t wish to run a command, prefix the first "/" with a "\\".',
   },
   pt: {
-    requiresSelection: 'Este quick command precisa que algum código esteja selecionado no editor.',
+    requiresSelection: 'Selecione o código no editor ao lado para que este quick command possa interagir com ele.',
     startQuestioning: 'Para executar o Quick Command "$0", vou precisar que você providencie algumas explicações. Algumas são obrigatórias e outras opcionais. Vamos começar.',
-    progress: 'Executando quick command "$0". Passo $1 ($2) de $3.',
+    progress: 'Executando step "$0" do Quick Command "$1".',
     aborted: 'A execução do quick command foi abortada pelo usuário.',
     notFound: 'Não existe quick command com o nome providenciado. Se você não quiser rodar um comando, prefixe o primeiro "/" com um "\\".',
   },