diff --git a/src/Document/document.schema.ts b/src/Document/document.schema.ts
index ec3382c..77c9bce 100644
--- a/src/Document/document.schema.ts
+++ b/src/Document/document.schema.ts
@@ -3,6 +3,7 @@ import { Document } from '@prisma/client'
 import { Pothos, PothosRef, PothosSchema, SchemaBuilderToken } from '@smatch-corp/nestjs-pothos'
 import Delta from 'quill-delta'
 import { MinioService } from 'src/Minio/minio.service'
+import { PromptType } from 'src/OpenAI/openai.service'
 import { Builder, SchemaContext } from '../Graphql/graphql.builder'
 import { PrismaService } from '../Prisma/prisma.service'
 import { DocumentEvent } from './document.event'
@@ -199,10 +200,18 @@ export class DocumentSchema extends PothosSchema {
 
       testCheckGrammar: t.field({
         type: 'Boolean',
-        resolve: async (_query, _args, ctx: SchemaContext) => {
-          const documentId = '0a07abe1-4d21-4530-85b9-ed128bfd8f9e'
-          const pageId = 0
-          await this.documentService.checkGrammarForPage(documentId, pageId)
+        args: {
+          documentId: t.arg({ type: 'String', required: true }),
+          pageId: t.arg({ type: 'Int', required: true }),
+          promptType: t.arg({
+            type: this.builder.enumType('PromptType', {
+              values: ['CHECK_GRAMMAR', 'REWRITE_TEXT', 'SUMMARIZE', 'TRANSLATE', 'EXPAND_CONTENT'] as const,
+            }),
+            required: true,
+          }),
+        },
+        resolve: async (_query, args, ctx: SchemaContext) => {
+          await this.documentService.checkGrammarForPage(args.documentId, args.pageId, args.promptType as PromptType)
           return true
         },
       }),
diff --git a/src/Document/document.service.ts b/src/Document/document.service.ts
index 3b1462d..6d22e55 100644
--- a/src/Document/document.service.ts
+++ b/src/Document/document.service.ts
@@ -67,7 +67,7 @@ export class DocumentService {
   ) {}
 
   // check grammar for a page by parallely send each sentence to OpenAI service as text and get the result, after that return the result as delta and publish the result to the document
-  async checkGrammarForPage(documentId: string, pageId: number): Promise<void> {
+  async checkGrammarForPage(documentId: string, pageId: number , promptType: PromptType): Promise<void> {
     const content = await this.minio.getDocumentContent(documentId, pageId)
     content.ops.forEach(async (op) => {
       if (typeof op.insert !== 'string') {
@@ -81,20 +81,23 @@ export class DocumentService {
         return
       }
       const originalDelta = new Delta().push(op)
-      const grammarCheckResult = await this.openai.processText(op.insert, PromptType.CHECK_GRAMMAR)
+      Logger.log(op.insert, 'op.insert')
+      Logger.log(promptType, 'promptType')
+      const grammarCheckResult = await this.openai.processText(op.insert, promptType)
       if (!grammarCheckResult) {
         return
       }
+      Logger.log(grammarCheckResult, 'grammarCheckResult')
       // create new delta and maintain the original delta attributes
       const newDelta = new Delta().push(op)
       newDelta.ops[0].attributes = originalDelta.ops[0].attributes
       newDelta.ops[0].insert = grammarCheckResult
       // compose the original delta with the grammarCheckResult
-      const correctedDelta = originalDelta.compose(newDelta)
+      // const correctedDelta = originalDelta.ops[0].insert = grammarCheckResult
       // calculate where to insert the correctedDelta
-      const index = content.ops.findIndex((op) => op.insert === op.insert)
-      content.ops.splice(index, 0, correctedDelta.ops[0])
-      Logger.log(JSON.stringify(content), 'content')
+      // const index = content.ops.findIndex((op) => op.insert === op.insert)
+      // content.ops.splice(index, 0, newDelta)
+      Logger.log(JSON.stringify(newDelta), 'newDelta')
       
       
       
@@ -105,10 +108,10 @@ export class DocumentService {
         documentId,
         pageIndex: pageId,
         eventType: DocumentEvent.AI_SUGGESTION,
-        delta: content,
+        delta: newDelta,
         senderId: 'system',
       }
-      this.pubSub.publish(`${DocumentEvent.AI_SUGGESTION}.${documentId}.${pageId}`, payload)
+      await this.pubSub.publish(`${DocumentEvent.AI_SUGGESTION}.${documentId}`, payload)
     })
   }
 
diff --git a/src/OpenAI/openai.service.ts b/src/OpenAI/openai.service.ts
index 7b2ea1b..912f076 100644
--- a/src/OpenAI/openai.service.ts
+++ b/src/OpenAI/openai.service.ts
@@ -20,12 +20,14 @@ const prompts = {
   REWRITE_TEXT: `You will be provided with text delimited by triple quotes.
   Rewrite the text to improve clarity, conciseness, and overall readability.
   Maintain the original meaning and tone of the text.
-  If the text is already well-written, return it as is.
+  Return the text within triple quotes.
+  If the text is already well-written, return back the text delimited by triple quotes.
   \"\"\"{text}\"\"\"
   `,
   SUMMARIZE: `You will be provided with text delimited by triple quotes.
   Create a concise summary that captures the key points and main ideas of the text.
   The summary should be clear, objective, and no longer than 3-4 sentences.
+  If the text is already well-written, return back the text delimited by triple quotes.
   \"\"\"{text}\"\"\"
   `,
   TRANSLATE: `You will be provided with text delimited by triple quotes and a target language.