PLU-445 - [AISAY-4] - Change file selector to Dropdown (#928)

### TL;DR

Refactored AISAY actions to use a single file input instead of multiple
attachments and centralized schema validation.

**Note to reviewer**
This initial rollout will only include the generalised model as there
are pending issues to iron out for the specific model.

### What changed?

- Changed the file input type from `multiselect` to `dropdown` in both
generalised and specific model actions
- Consolidated schema validation logic into a common `schema.ts` file
- Added support for specific variable types in `ControlledAutocomplete`
component
- Used logger to log errors

### How to test?

- [ ] Use the AISAY generalised model action with a single file input

### Why make this change?

The AISAY actions were designed to work with a single file, but the
actions were previously using the MultiSelect component.

### Screenshots

![Screenshot 2025-04-07 at
11.50.40 AM.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/Wrwhm8Mmhv1Z2GwlSb7V/8c7d2d6f-7db2-47ba-86ab-c56af1c374bc.png)

Author: kevinkim-ogp

packages/backend/src/apps/aisay/__tests__/common/schema.test.ts

@@ -0,0 +1,78 @@
+import { assert, beforeEach, describe, it } from 'vitest'
+
+import {
+  generalisedModelSchema,
+  specificModelSchema,
+} from '../../common/schema'
+
+describe('AISAY schema', () => {
+  let generalisedModelPayload: Record<string, unknown>
+  let specificModelPayload: Record<string, unknown>
+
+  beforeEach(() => {
+    generalisedModelPayload = {
+      file: 's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
+      prompts: [
+        {
+          prompt: 'What is your name?',
+        },
+        {
+          prompt: 'What is your age?',
+        },
+      ],
+    }
+
+    specificModelPayload = {
+      file: 's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
+      documentType: 'BANK_STATEMENT',
+    }
+  })
+
+  describe('generalised model schema', () => {
+    it('should validate generalised model schema', () => {
+      const result = generalisedModelSchema.safeParse(generalisedModelPayload)
+      assert(result.success === true)
+      assert(
+        result.data.file ===
+          's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
+      )
+      assert(Object.keys(result.data.prompts).length === 2)
+      assert(
+        result.data.prompts['additionalProp0'].description ===
+          'Extract the What is your name?',
+      )
+      assert(
+        result.data.prompts['additionalProp1'].description ===
+          'Extract the What is your age?',
+      )
+    })
+
+    it('should fail if file is not a valid S3 ID', () => {
+      generalisedModelPayload.file = '123'
+      const result = generalisedModelSchema.safeParse(generalisedModelPayload)
+      assert(result.success === false)
+    })
+  })
+
+  describe('specific model schema', () => {
+    it('should validate specific model schema', () => {
+      const result = specificModelSchema.safeParse(specificModelPayload)
+      assert(result.success === true)
+    })
+
+    it.each(['BANK_STATEMENT', 'CHEQUE', 'INVOICE', 'PASSPORT', 'RECEIPT'])(
+      `should validate for valid document type: %s`,
+      (documentType) => {
+        specificModelPayload.documentType = documentType
+        const result = specificModelSchema.safeParse(specificModelPayload)
+        assert(result.success === true)
+      },
+    )
+
+    it('should fail if document type is not a valid document type', () => {
+      specificModelPayload.documentType = 'INVALID_DOCUMENT_TYPE'
+      const result = specificModelSchema.safeParse(specificModelPayload)
+      assert(result.success === false)
+    })
+  })
+})

packages/backend/src/apps/aisay/actions/use-generalised-model/get-data-out-metadata.ts

@@ -1,6 +1,6 @@
 import { IDataOutMetadata, IExecutionStep } from '@plumber/types'
 
-import { getInfoToExtract } from '../../common/info-to-extract'
+import { getPrompts } from '../../common/get-prompts'
 
 async function getDataOutMetadata(
   executionStep: IExecutionStep,
@@ -11,23 +11,29 @@ async function getDataOutMetadata(
     return null
   }
 
-  const infoToExtract = await getInfoToExtract(stepId)
+  const prompts = await getPrompts(stepId)
 
   const fieldsMetadata: Record<string, IDataOutMetadata> = {}
   Object.keys(dataOut.fields).forEach((key) => {
     const index = parseInt(key.split('additionalProp')[1])
-    const fieldName = infoToExtract[index]
+    const fieldName = prompts[index]
     fieldsMetadata[key] = {
       label: fieldName,
     }
   })
 
+  const otherMetadata: Record<string, IDataOutMetadata> = {}
+  Object.keys(dataOut).forEach((key) => {
+    if (key !== 'fields') {
+      otherMetadata[key] = {
+        isHidden: true,
+      }
+    }
+  })
+
   return {
-    quota: {
-      label: 'Quota',
-      isHidden: true,
-    },
     fields: fieldsMetadata,
+    ...otherMetadata,
   }
 }
 

packages/backend/src/apps/aisay/actions/use-generalised-model/index.ts

@@ -1,14 +1,15 @@
-import { IJSONArray, IRawAction } from '@plumber/types'
+import { IRawAction } from '@plumber/types'
 
 import StepError from '@/errors/step'
+import logger from '@/helpers/logger'
 import Step from '@/models/step'
 
 import { getToken } from '../../auth/get-token'
 import { parseError } from '../../common/error-parser'
-import { getAttachmentsFromS3, getValidationError } from '../../common/utils'
+import { generalisedModelSchema } from '../../common/schema'
+import { getAttachmentFromS3, getValidationError } from '../../common/utils'
 
 import getDataOutMetadata from './get-data-out-metadata'
-import { requestSchema } from './schema'
 
 const action: IRawAction = {
   name: 'Extract data from all document types',
@@ -17,8 +18,8 @@ const action: IRawAction = {
   arguments: [
     {
       label: 'File',
-      key: 'attachments',
-      type: 'multiselect' as const,
+      key: 'file',
+      type: 'dropdown',
       required: true,
       variables: true,
       variableTypes: ['file'],
@@ -27,15 +28,15 @@ const action: IRawAction = {
       label: 'Prompts',
       description:
         'Enter prompts to specify how the data should be interpreted and extracted',
-      key: 'infoToExtract',
+      key: 'prompts',
       type: 'multirow' as const,
       required: true,
       variables: true,
       addRowButtonText: 'Add prompt',
       subFields: [
         {
           placeholder: 'E.g. Return the price of individual line items',
-          key: 'infoToExtract',
+          key: 'prompt',
           type: 'string' as const,
           required: true,
           variables: false,
@@ -44,17 +45,14 @@ const action: IRawAction = {
     },
   ],
   doesFileProcessing: (step: Step) => {
-    return (
-      step.parameters.attachments &&
-      (step.parameters.attachments as IJSONArray).length > 0
-    )
+    return step.parameters.file && step.parameters.file !== ''
   },
   getDataOutMetadata,
 
   async run($) {
-    const { attachments, infoToExtract } = $.step.parameters as {
-      attachments?: IJSONArray
-      infoToExtract: Array<{ infoToExtract: string }>
+    const { file, prompts } = $.step.parameters as {
+      file: string
+      prompts: Array<{ prompt: string }>
     }
 
     if (!$.auth.data?.clientId || !$.auth.data?.clientSecret) {
@@ -66,7 +64,7 @@ const action: IRawAction = {
       )
     }
 
-    const result = requestSchema.safeParse({ attachments, infoToExtract })
+    const result = generalisedModelSchema.safeParse({ file, prompts })
     if (!result.success) {
       const { stepErrorName, stepErrorSolution } = getValidationError(result)
 
@@ -79,26 +77,12 @@ const action: IRawAction = {
     }
 
     try {
-      /**
-       * FIXME (kevinkim-ogp): should only accept one attachment
-       * use a different selector on the frontend or update the
-       * multi-select to only allow one attachment
-       */
-      // Pre-call get attachments from S3 first
-      const attachmentFiles = await getAttachmentsFromS3(
-        result.data.attachments,
-        $.flow.id,
-      )
-      const attachment = attachmentFiles[0]
+      // get attachment from S3 first
+      const attachment = await getAttachmentFromS3(result.data.file, $.flow.id)
 
       // Step 1: get AWS Cognito access token
       const token = await getToken($)
 
-      /**
-       * TODO (kevinkim-ogp): first iteration of AISAY will only support synchronous calls
-       * - add a check to ensure that the attachment is less than 9 MB (7 MB to be safe)
-       * - add a check to ensure that the call to the model is less than 29 seconds
-       */
       // Step 2: Call the model to get the output
       const res = await $.http.request({
         url: `${$.app.baseUrl}/query`,
@@ -109,14 +93,14 @@ const action: IRawAction = {
           Authorization: `Bearer ${token}`,
         },
         data: {
-          gpt_query: result.data.infoToExtract,
+          gpt_query: result.data.prompts,
           image: attachment.data,
         },
       })
 
       $.setActionItem({ raw: { ...res.data } })
     } catch (err) {
-      console.error(err)
+      logger.error(err)
       const { stepErrorName, stepErrorSolution } = parseError(err)
       throw new StepError(
         stepErrorName,

packages/backend/src/apps/aisay/actions/use-specific-model/index.ts

@@ -1,15 +1,16 @@
-import { IJSONArray, IRawAction } from '@plumber/types'
+import { IRawAction } from '@plumber/types'
 
 import appConfig from '@/config/app'
 import StepError from '@/errors/step'
+import logger from '@/helpers/logger'
 import Step from '@/models/step'
 
 import { getToken } from '../../auth/get-token'
 import { parseError } from '../../common/error-parser'
-import { getAttachmentsFromS3, getValidationError } from '../../common/utils'
+import { specificModelSchema } from '../../common/schema'
+import { getAttachmentFromS3, getValidationError } from '../../common/utils'
 
 import getDataOutMetadata from './get-data-out-metadata'
-import { requestSchema } from './schema'
 
 const action: IRawAction = {
   name: 'Extract data from specific document types',
@@ -19,8 +20,8 @@ const action: IRawAction = {
   arguments: [
     {
       label: 'File',
-      key: 'attachments',
-      type: 'multiselect' as const,
+      key: 'file',
+      type: 'dropdown',
       required: true,
       variables: true,
       variableTypes: ['file'],
@@ -41,16 +42,13 @@ const action: IRawAction = {
     },
   ],
   doesFileProcessing: (step: Step) => {
-    return (
-      step.parameters.attachments &&
-      (step.parameters.attachments as IJSONArray).length > 0
-    )
+    return step.parameters.file && step.parameters.file !== ''
   },
   getDataOutMetadata,
 
   async run($) {
-    const { attachments, documentType } = $.step.parameters as {
-      attachments?: IJSONArray
+    const { file, documentType } = $.step.parameters as {
+      file: string
       documentType: string
     }
 
@@ -63,7 +61,7 @@ const action: IRawAction = {
       )
     }
 
-    const result = requestSchema.safeParse({ attachments, documentType })
+    const result = specificModelSchema.safeParse({ file, documentType })
     if (!result.success) {
       const { stepErrorName, stepErrorSolution } = getValidationError(result)
 
@@ -76,19 +74,12 @@ const action: IRawAction = {
     }
 
     try {
-      // Pre-call get attachments from S3 first
-      const attachmentFiles = await getAttachmentsFromS3(
-        result.data.attachments,
-        $.flow.id,
-      )
-      const attachment = attachmentFiles[0]
+      // get attachment from S3 first
+      const attachment = await getAttachmentFromS3(result.data.file, $.flow.id)
 
       // Step 1: get AWS Cognito access token
       const token = await getToken($)
 
-      // Assuming we do a synchronous call to the model
-      // which needs to be less than 29 seconds
-      // and with a document size of less than 9 MB
       // Step 2: Call the model to get the output
       const aisayRes = await $.http.request({
         url: `${appConfig.aisayApiUrl}/query`,
@@ -108,7 +99,7 @@ const action: IRawAction = {
         raw: { ...aisayRes.data, documentType: result.data.documentType },
       })
     } catch (err) {
-      console.error(err)
+      logger.error(err)
       const { stepErrorName, stepErrorSolution } = parseError(err)
 
       throw new StepError(

packages/backend/src/apps/aisay/common/constants.ts

@@ -1 +1,7 @@
-export const MAX_FILE_SIZE = 6 * 1024 * 1024 // 6 MB
+export const DOCUMENT_TYPES = [
+  'BANK_STATEMENT',
+  'CHEQUE',
+  'INVOICE',
+  'PASSPORT',
+  'RECEIPT',
+]

packages/backend/src/apps/aisay/common/get-prompts.ts

@@ -0,0 +1,9 @@
+import Step from '@/models/step'
+
+export const getPrompts = async (stepId: string) => {
+  const step = await Step.query().findById(stepId)
+  if (!Array.isArray(step.parameters.prompts)) {
+    return []
+  }
+  return step.parameters.prompts.map((i: { prompt: string }) => i.prompt)
+}