feat: add api rate limit for azure openai

opengovsg · Mar 10, 2024 · 8decc71 · 8decc71
1 parent 9311c1f
commit 8decc71
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 82 deletions.
diff --git a/frontend/src/features/admin-form/assistance/AssistanceService.ts b/frontend/src/features/admin-form/assistance/AssistanceService.ts
@@ -8,19 +8,32 @@ export const generateQuestions = async (
   type: string,
   content: string,
 ): Promise<ChatResponseMessage> => {
-  return await ApiService.post(`${ADMIN_FORM_ENDPOINT}/assistance/questions`, {
-    type,
-    content,
-  }).then(({ data }) => data)
+  try {
+    const response = await ApiService.post(
+      `${ADMIN_FORM_ENDPOINT}/assistance/questions`,
+      {
+        type,
+        content,
+      },
+    )
+    return response.data
+  } catch (error) {
+    throw new Error('Error creating form')
+  }
 }
 
 export const generateFormFields = async (
   content: string,
 ): Promise<ChatResponseMessage> => {
-  return await ApiService.post(
-    `${ADMIN_FORM_ENDPOINT}/assistance/form-fields`,
-    {
-      content,
-    },
-  ).then(({ data }) => data)
+  try {
+    const response = await ApiService.post(
+      `${ADMIN_FORM_ENDPOINT}/assistance/form-fields`,
+      {
+        content,
+      },
+    )
+    return response.data
+  } catch (error) {
+    throw new Error('Error creating form')
+  }
 }
diff --git a/frontend/src/features/admin-form/assistance/mutations.ts b/frontend/src/features/admin-form/assistance/mutations.ts
@@ -1,6 +1,7 @@
 import { useContext } from 'react'
 import { useMutation, useQueryClient } from 'react-query'
 import { useParams } from 'react-router-dom'
+import { error } from 'console'
 
 import { ContentTypes } from '~shared/types/assistance'
 
@@ -30,82 +31,101 @@ export const useAssistanceMutations = () => {
   const onCloseContext = useContext(MagicFormBuilderModalOnCloseContext)
   const { onClose } = onCloseContext || {}
 
-  const createFieldsFromPromptMutation = useMutation((prompt: string) =>
-    generateQuestions(ContentTypes.PROMPT, prompt)
-      .then((questions) => {
-        if (!questions.content) {
-          throw new Error('No content in questions')
-        }
-        return generateFormFields(questions.content)
-      })
-      .then((data) => {
-        let formFields
-        if (data.content) {
-          try {
-            formFields = JSON.parse(parseModelOutput(data.content))
-          } catch (e) {
-            toast({
-              description: `Error creating form. Reason: ${e}`,
-              status: 'warning',
-            })
-            return
+  const createFieldsFromPromptMutation = useMutation(
+    (prompt: string) =>
+      generateQuestions(ContentTypes.PROMPT, prompt)
+        .then((questions) => {
+          if (!questions.content) {
+            throw new Error('No content in questions')
           }
-        }
-        return createFieldsMutation.mutate(formFields, {
-          onSuccess: () => {
-            queryClient.invalidateQueries(adminFormKeys.id(formId))
-            onClose()
-            toast({
-              description: 'Successfully created form',
-            })
-          },
-          onError: () => {
-            toast({
-              description: 'Error creating form.',
-              status: 'warning',
-            })
-          },
+          return generateFormFields(questions.content)
         })
-      }),
+        .then((data) => {
+          let formFields
+          if (data.content) {
+            try {
+              formFields = JSON.parse(parseModelOutput(data.content))
+            } catch (e) {
+              toast({
+                description: `Error creating form. Reason: ${e}`,
+                status: 'warning',
+              })
+              return
+            }
+          }
+          return createFieldsMutation.mutate(formFields, {
+            onSuccess: () => {
+              queryClient.invalidateQueries(adminFormKeys.id(formId))
+              onClose()
+              toast({
+                description: 'Successfully created form',
+              })
+            },
+            onError: () => {
+              toast({
+                description: 'Error creating form.',
+                status: 'warning',
+              })
+            },
+          })
+        }),
+    {
+      onError: (error: Error) => {
+        toast({
+          description: 'Too many forms generated! Please try again later.',
+          status: 'danger',
+        })
+      },
+    },
   )
 
-  const createFieldsFromPdfMutation = useMutation((pdfContent: string) =>
-    generateQuestions(ContentTypes.PDF, pdfContent)
-      .then((questions) => {
-        if (!questions.content) {
-          throw new Error('No content in questions')
-        }
-        return generateFormFields(questions.content)
-      })
-      .then((data) => {
-        let formFields
-        if (data.content) {
-          try {
-            formFields = JSON.parse(parseModelOutput(data.content))
-          } catch (e) {
-            toast({
-              description: `Error creating form. Reason: ${e}`,
-              status: 'warning',
-            })
-            return
+  const createFieldsFromPdfMutation = useMutation(
+    (pdfContent: string) =>
+      generateQuestions(ContentTypes.PDF, pdfContent)
+        .then((questions) => {
+          if (!questions.content) {
+            throw new Error('No content in questions')
+          }
+          return generateFormFields(questions.content)
+        })
+        .then((data) => {
+          let formFields
+          // Proceed with normal processing if there is an error with content generation
+          if (data.content) {
+            try {
+              formFields = JSON.parse(parseModelOutput(data.content))
+            } catch (e) {
+              toast({
+                description: `Error creating form. Reason: ${e}`,
+                status: 'warning',
+              })
+              return
+            }
           }
-        }
-        return createFieldsMutation.mutate(formFields, {
-          onSuccess: () => {
-            queryClient.invalidateQueries(adminFormKeys.id(formId))
-            onClose()
-            toast({
-              description: 'Successfully created form',
-            })
-          },
-          onError: () => {
-            toast({
-              description: 'Error creating form.',
-              status: 'warning',
-            })
-          },
+          return createFieldsMutation.mutate(formFields, {
+            onSuccess: () => {
+              queryClient.invalidateQueries(adminFormKeys.id(formId))
+              onClose()
+              toast({
+                description: 'Successfully created form',
+              })
+            },
+            onError: () => {
+              toast({
+                description: 'Error creating form.',
+                status: 'warning',
+              })
+            },
+          })
+        }),
+    {
+      onError: (error: Error) => {
+        toast({
+          description: 'Too many forms generated! Please try again later.',
+          status: 'danger',
         })
-      }),
+      },
+    },
   )
 
   return {

diff --git a/src/app/routes/api/v3/admin/assistance/assistance.routes.ts b/src/app/routes/api/v3/admin/assistance/assistance.routes.ts
@@ -1,11 +1,20 @@
 import { Router } from 'express'
 
+import { limitRate } from '../../../../../../app/utils/limit-rate'
 import {
   handleGenerateFormFields,
   handleGenerateQuestions,
 } from '../../../../../modules/form/admin-form/admin-form.assistance.controller'
 
 export const AssistanceRouter = Router()
 
-AssistanceRouter.post('/questions', handleGenerateQuestions)
-AssistanceRouter.post('/form-fields', handleGenerateFormFields)
+AssistanceRouter.post(
+  '/questions',
+  limitRate({ windowMs: 60 * 1000 * 60, max: 20 }),
+  handleGenerateQuestions,
+)
+AssistanceRouter.post(
+  '/form-fields',
+  limitRate({ windowMs: 60 * 1000 * 60, max: 20 }),
+  handleGenerateFormFields,
+)