Skip to content

Commit b30d1ec

Browse files
kevinkim-ogpkevinkim-ogp
committed
chore: use dropdown for file selection
1 parent 3bf608a commit b30d1ec

File tree

11 files changed

+176
-127
lines changed

11 files changed

+176
-127
lines changed

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import { assert, beforeEach, describe, it } from 'vitest'
2+
3+
import {
4+
generalisedModelSchema,
5+
specificModelSchema,
6+
} from '../../common/schema'
7+
8+
describe('AISAY schema', () => {
9+
let generalisedModelPayload: Record<string, unknown>
10+
let specificModelPayload: Record<string, unknown>
11+
12+
beforeEach(() => {
13+
generalisedModelPayload = {
14+
file: 's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
15+
infoToExtract: [
16+
{
17+
infoToExtract: 'What is your name?',
18+
},
19+
{
20+
infoToExtract: 'What is your age?',
21+
},
22+
],
23+
}
24+
25+
specificModelPayload = {
26+
file: 's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
27+
documentType: 'BANK_STATEMENT',
28+
}
29+
})
30+
31+
describe('generalised model schema', () => {
32+
it('should validate generalised model schema', () => {
33+
const result = generalisedModelSchema.safeParse(generalisedModelPayload)
34+
assert(result.success === true)
35+
assert(
36+
result.data.file ===
37+
's3:plumber-test-bucket:123456/789abc/plumber-logo.jpg',
38+
)
39+
assert(Object.keys(result.data.infoToExtract).length === 2)
40+
assert(
41+
result.data.infoToExtract['additionalProp0'].description ===
42+
'Extract the What is your name?',
43+
)
44+
assert(
45+
result.data.infoToExtract['additionalProp1'].description ===
46+
'Extract the What is your age?',
47+
)
48+
})
49+
50+
it('should fail if file is not a valid S3 ID', () => {
51+
generalisedModelPayload.file = '123'
52+
const result = generalisedModelSchema.safeParse(generalisedModelPayload)
53+
assert(result.success === false)
54+
})
55+
})
56+
57+
describe('specific model schema', () => {
58+
it('should validate specific model schema', () => {
59+
const result = specificModelSchema.safeParse(specificModelPayload)
60+
assert(result.success === true)
61+
})
62+
63+
it.each(['BANK_STATEMENT', 'CHEQUE', 'INVOICE', 'PASSPORT', 'RECEIPT'])(
64+
`should validate for valid document type: %s`,
65+
(documentType) => {
66+
specificModelPayload.documentType = documentType
67+
const result = specificModelSchema.safeParse(specificModelPayload)
68+
assert(result.success === true)
69+
},
70+
)
71+
72+
it('should fail if document type is not a valid document type', () => {
73+
specificModelPayload.documentType = 'INVALID_DOCUMENT_TYPE'
74+
const result = specificModelSchema.safeParse(specificModelPayload)
75+
assert(result.success === false)
76+
})
77+
})
78+
})

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

Lines changed: 13 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
1-
import { IJSONArray, IRawAction } from '@plumber/types'
1+
import { IRawAction } from '@plumber/types'
22

33
import appConfig from '@/config/app'
44
import StepError from '@/errors/step'
5+
import logger from '@/helpers/logger'
56

67
import { getToken } from '../../auth/get-token'
78
import { parseError } from '../../common/error-parser'
8-
import { getAttachmentsFromS3, getValidationError } from '../../common/utils'
9+
import { generalisedModelSchema } from '../../common/schema'
10+
import { getAttachmentFromS3, getValidationError } from '../../common/utils'
911

1012
import getDataOutMetadata from './get-data-out-metadata'
11-
import { requestSchema } from './schema'
1213

1314
const action: IRawAction = {
1415
name: 'Use generalised model',
@@ -18,9 +19,9 @@ const action: IRawAction = {
1819
arguments: [
1920
{
2021
label: 'File',
21-
key: 'attachments',
22-
description: 'Use files from previous steps',
23-
type: 'multiselect' as const,
22+
key: 'file',
23+
description: 'Us files from previous steps',
24+
type: 'dropdown',
2425
required: true,
2526
variables: true,
2627
variableTypes: ['file'],
@@ -47,8 +48,8 @@ const action: IRawAction = {
4748
getDataOutMetadata,
4849

4950
async run($) {
50-
const { attachments, infoToExtract } = $.step.parameters as {
51-
attachments?: IJSONArray
51+
const { file, infoToExtract } = $.step.parameters as {
52+
file: string
5253
infoToExtract: Array<{ infoToExtract: string }>
5354
}
5455

@@ -61,7 +62,7 @@ const action: IRawAction = {
6162
)
6263
}
6364

64-
const result = requestSchema.safeParse({ attachments, infoToExtract })
65+
const result = generalisedModelSchema.safeParse({ file, infoToExtract })
6566
if (!result.success) {
6667
const { stepErrorName, stepErrorSolution } = getValidationError(result)
6768

@@ -74,26 +75,12 @@ const action: IRawAction = {
7475
}
7576

7677
try {
77-
/**
78-
* FIXME (kevinkim-ogp): should only accept one attachment
79-
* use a different selector on the frontend or update the
80-
* multi-select to only allow one attachment
81-
*/
82-
// Pre-call get attachments from S3 first
83-
const attachmentFiles = await getAttachmentsFromS3(
84-
result.data.attachments,
85-
$.flow.id,
86-
)
87-
const attachment = attachmentFiles[0]
78+
// get attachment from S3 first
79+
const attachment = await getAttachmentFromS3(result.data.file, $.flow.id)
8880

8981
// Step 1: get AWS Cognito access token
9082
const token = await getToken($)
9183

92-
/**
93-
* TODO (kevinkim-ogp): first iteration of AISAY will only support synchronous calls
94-
* - add a check to ensure that the attachment is less than 9 MB (7 MB to be safe)
95-
* - add a check to ensure that the call to the model is less than 29 seconds
96-
*/
9784
// Step 2: Call the model to get the output
9885
const res = await $.http.request({
9986
url: `${appConfig.aisayApiUrl}/query`,
@@ -111,7 +98,7 @@ const action: IRawAction = {
11198

11299
$.setActionItem({ raw: { ...res.data } })
113100
} catch (err) {
114-
console.error(err)
101+
logger.error(err)
115102
const { stepErrorName, stepErrorSolution } = parseError(err)
116103
throw new StepError(
117104
stepErrorName,

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

Lines changed: 0 additions & 36 deletions
This file was deleted.

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

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
1-
import { IJSONArray, IRawAction } from '@plumber/types'
1+
import { IRawAction } from '@plumber/types'
22

33
import appConfig from '@/config/app'
44
import StepError from '@/errors/step'
5+
import logger from '@/helpers/logger'
56

67
import { getToken } from '../../auth/get-token'
78
import { parseError } from '../../common/error-parser'
8-
import { getAttachmentsFromS3, getValidationError } from '../../common/utils'
9+
import { specificModelSchema } from '../../common/schema'
10+
import { getAttachmentFromS3, getValidationError } from '../../common/utils'
911

1012
import getDataOutMetadata from './get-data-out-metadata'
11-
import { requestSchema } from './schema'
1213

1314
const action: IRawAction = {
1415
name: 'Use specific model',
@@ -18,9 +19,9 @@ const action: IRawAction = {
1819
arguments: [
1920
{
2021
label: 'File',
21-
key: 'attachments',
22+
key: 'file',
2223
description: 'Use file from previous steps',
23-
type: 'multiselect' as const,
24+
type: 'dropdown',
2425
required: true,
2526
variables: true,
2627
variableTypes: ['file'],
@@ -43,8 +44,8 @@ const action: IRawAction = {
4344
getDataOutMetadata,
4445

4546
async run($) {
46-
const { attachments, documentType } = $.step.parameters as {
47-
attachments?: IJSONArray
47+
const { file, documentType } = $.step.parameters as {
48+
file: string
4849
documentType: string
4950
}
5051

@@ -57,7 +58,7 @@ const action: IRawAction = {
5758
)
5859
}
5960

60-
const result = requestSchema.safeParse({ attachments, documentType })
61+
const result = specificModelSchema.safeParse({ file, documentType })
6162
if (!result.success) {
6263
const { stepErrorName, stepErrorSolution } = getValidationError(result)
6364

@@ -70,19 +71,12 @@ const action: IRawAction = {
7071
}
7172

7273
try {
73-
// Pre-call get attachments from S3 first
74-
const attachmentFiles = await getAttachmentsFromS3(
75-
result.data.attachments,
76-
$.flow.id,
77-
)
78-
const attachment = attachmentFiles[0]
74+
// get attachment from S3 first
75+
const attachment = await getAttachmentFromS3(result.data.file, $.flow.id)
7976

8077
// Step 1: get AWS Cognito access token
8178
const token = await getToken($)
8279

83-
// Assuming we do a synchronous call to the model
84-
// which needs to be less than 29 seconds
85-
// and with a document size of less than 9 MB
8680
// Step 2: Call the model to get the output
8781
const aisayRes = await $.http.request({
8882
url: `${appConfig.aisayApiUrl}/query`,
@@ -102,7 +96,7 @@ const action: IRawAction = {
10296
raw: { ...aisayRes.data, documentType: result.data.documentType },
10397
})
10498
} catch (err) {
105-
console.error(err)
99+
logger.error(err)
106100
const { stepErrorName, stepErrorSolution } = parseError(err)
107101

108102
throw new StepError(

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

Lines changed: 0 additions & 25 deletions
This file was deleted.

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,8 @@
11
export const MAX_FILE_SIZE = 6 * 1024 * 1024 // 6 MB
2+
export const DOCUMENT_TYPES = [
3+
'BANK_STATEMENT',
4+
'CHEQUE',
5+
'INVOICE',
6+
'PASSPORT',
7+
'RECEIPT',
8+
]

‎packages/backend/src/apps/aisay/common/schema.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { z } from 'zod'
2+
3+
import { parseS3Id } from '@/helpers/s3'
4+
5+
import { DOCUMENT_TYPES } from './constants'
6+
7+
export const fileSchema = z.string().transform((value, context) => {
8+
if (!parseS3Id(value)) {
9+
context.addIssue({
10+
code: z.ZodIssueCode.custom,
11+
message: `${value} is not a S3 ID.`,
12+
})
13+
return z.NEVER
14+
}
15+
return value
16+
})
17+
18+
export const generalisedModelSchema = z.object({
19+
file: fileSchema,
20+
infoToExtract: z
21+
.array(z.object({ infoToExtract: z.string() }))
22+
.transform((array) => {
23+
const result: Record<string, { description: string; type: string }> = {}
24+
array.forEach((a, index) => {
25+
result[`additionalProp${index}`] = {
26+
description: `Extract the ${a.infoToExtract}`,
27+
type: 'string',
28+
}
29+
})
30+
return result
31+
}),
32+
})
33+
34+
export const documentTypeEnum = z.enum(DOCUMENT_TYPES as [string, ...string[]])
35+
36+
export const specificModelSchema = z.object({
37+
file: fileSchema,
38+
documentType: documentTypeEnum,
39+
})

‎packages/backend/src/apps/aisay/common/utils.ts‎

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,12 @@ export const uint8ArrayToBase64 = (uint8Array: Uint8Array) => {
77
return Buffer.from(uint8Array).toString('base64')
88
}
99

10-
/**
11-
* TODO (kevinkim-ogp):
12-
* - add a check to ensure that the attachment is less than 9 MB (7 MB to be safe)
13-
* - update to fetch a single file
14-
* - validate flow id
15-
*/
16-
export const getAttachmentsFromS3 = async (
17-
attachments: string[],
10+
export const getAttachmentFromS3 = async (
11+
attachment: string,
1812
flowId: string,
1913
) => {
20-
const attachmentFiles = await Promise.all(
21-
attachments.map(async (attachment) => {
22-
const obj = await getObjectFromS3Id(attachment, { flowId })
23-
return { fileName: obj.name, data: uint8ArrayToBase64(obj.data) }
24-
}),
25-
)
26-
27-
return attachmentFiles
14+
const obj = await getObjectFromS3Id(attachment, { flowId })
15+
return { fileName: obj.name, data: uint8ArrayToBase64(obj.data) }
2816
}
2917

3018
export const getValidationError = (

0 commit comments

Comments
 (0)