refactor: add session id and mode

Author: kevinkim-ogp

packages/backend/src/graphql/__tests__/mutations/ai/generate-ai-steps.test.ts

@@ -74,6 +74,8 @@ const DEFAULT_MOCKED_OUTPUT = {
 const DEFAULT_INPUT = {
   prompt:
     '#### Start the workflow\nNew FormSG submission\n\n#### Actions\nSend a welcome email to the new user\nSend a welcome SMS to the new user',
+  isFormMode: true,
+  sessionId: '123',
 }
 
 describe('generateAiSteps mutation', () => {

packages/backend/src/graphql/mutations/ai/generate-ai-steps.ts

@@ -33,7 +33,7 @@ const generateAiSteps: MutationResolvers['generateAiSteps'] = async (
 
   try {
     const validatedInput = INPUT_SCHEMA.parse(params.input)
-    const { prompt: userPrompt } = validatedInput
+    const { prompt: userPrompt, isFormMode, sessionId } = validatedInput
 
     if (!context.currentUser) {
       throw new ForbiddenError('Not authorised!')
@@ -47,12 +47,23 @@ const generateAiSteps: MutationResolvers['generateAiSteps'] = async (
     const { prompt: systemPrompt } = prompt
 
     const result = await startActiveObservation(
-      'generate-ai-steps',
+      'generate-steps',
       async (trace) => {
+        const tags = ['ai-builder', 'generate-steps']
+        // TODO(kevinkim-ogp): we add the tags based on the input mode on the frontend
+        // whether it was from form or chat
+        // to be removed once A/B test is complete
+        if (isFormMode) {
+          tags.push('form')
+        } else {
+          tags.push('chat')
+        }
+
         trace.updateTrace({
           userId: context.currentUser.email,
           environment: appConfig.appEnv,
-          tags: ['graphql', 'ai-steps', 'generate-object'],
+          tags,
+          sessionId: sessionId ?? '',
         })
 
         trace.update({
@@ -62,7 +73,7 @@ const generateAiSteps: MutationResolvers['generateAiSteps'] = async (
         })
 
         const generation = trace.startObservation(
-          'ai-stream-generation',
+          'generate-steps',
           {
             model: MODEL_TYPE,
             input: [
@@ -85,7 +96,7 @@ const generateAiSteps: MutationResolvers['generateAiSteps'] = async (
           prompt: userPrompt,
           experimental_telemetry: {
             isEnabled: true,
-            functionId: 'generate-ai-steps',
+            functionId: 'generate-steps',
           },
         })
 

packages/backend/src/graphql/mutations/ai/schemas/input.zod.ts

@@ -6,4 +6,6 @@ export const INPUT_SCHEMA = z.object({
     .trim()
     .min(15, 'Prompt must be at least 15 characters')
     .max(3000, 'Prompt must not exceed 3000 characters'),
+  isFormMode: z.boolean(),
+  sessionId: z.string().nullish(),
 })

packages/backend/src/graphql/schema.graphql

@@ -128,6 +128,8 @@ type Mutation {
 
 input GenerateAiStepsInput {
   prompt: String!
+  isFormMode: Boolean!
+  sessionId: String
 }
 
 type AdminMutation {