diff --git a/docs/helpers/AI.md b/docs/helpers/AI.md
index 2ee6c6582..a3b6082ea 100644
--- a/docs/helpers/AI.md
+++ b/docs/helpers/AI.md
@@ -14,13 +14,13 @@ title: AI
 AI Helper for CodeceptJS.
 
 This helper class provides integration with the AI GPT-3.5 or 4 language model for generating responses to questions or prompts within the context of web pages. It allows you to interact with the GPT-3.5 model to obtain intelligent responses based on HTML fragments or general prompts.
-This helper should be enabled with any web helpers like Playwright or Puppeteer or WebDrvier to ensure the HTML context is available.
+This helper should be enabled with any web helpers like Playwright or Puppeteer or WebDriver to ensure the HTML context is available.
 
 Use it only in development mode. It is recommended to run it only inside pause() mode.
 
 ## Configuration
 
-This helper should be configured in codecept.json or codecept.conf.js
+This helper should be configured in codecept.conf.{js|ts}
 
 -   `chunkSize`:  - The maximum number of characters to send to the AI API at once. We split HTML fragments by 8000 chars to not exceed token limit. Increase this value if you use GPT-4.
 
diff --git a/lib/helper/AI.js b/lib/helper/AI.js
index afa88782a..d59b15531 100644
--- a/lib/helper/AI.js
+++ b/lib/helper/AI.js
@@ -10,17 +10,21 @@ const { beautify } = require('../utils')
 const output = require('../output')
 const { registerVariable } = require('../pause')
 
+const gtpRole = {
+  user: 'user',
+}
+
 /**
  * AI Helper for CodeceptJS.
  *
  * This helper class provides integration with the AI GPT-3.5 or 4 language model for generating responses to questions or prompts within the context of web pages. It allows you to interact with the GPT-3.5 model to obtain intelligent responses based on HTML fragments or general prompts.
- * This helper should be enabled with any web helpers like Playwright or Puppeteer or WebDrvier to ensure the HTML context is available.
+ * This helper should be enabled with any web helpers like Playwright or Puppeteer or WebDriver to ensure the HTML context is available.
  *
  * Use it only in development mode. It is recommended to run it only inside pause() mode.
  *
  * ## Configuration
  *
- * This helper should be configured in codecept.json or codecept.conf.js
+ * This helper should be configured in codecept.conf.{js|ts}
  *
  * * `chunkSize`: (optional, default: 80000) - The maximum number of characters to send to the AI API at once. We split HTML fragments by 8000 chars to not exceed token limit. Increase this value if you use GPT-4.
  */
@@ -33,6 +37,7 @@ class AI extends Helper {
       chunkSize: 80000,
     }
     this.options = { ...this.options, ...config }
+    this.aiAssistant.enable(this.config)
   }
 
   _beforeSuite() {
@@ -68,8 +73,8 @@ class AI extends Helper {
 
     for (const chunk of htmlChunks) {
       const messages = [
-        { role: 'user', content: prompt },
-        { role: 'user', content: `Within this HTML: ${minifyHtml(chunk)}` },
+        { role: gtpRole.user, content: prompt },
+        { role: gtpRole.user, content: `Within this HTML: ${minifyHtml(chunk)}` },
       ]
 
       if (htmlChunks.length > 1)
@@ -104,8 +109,8 @@ class AI extends Helper {
     const html = await this.helper.grabHTMLFrom(locator)
 
     const messages = [
-      { role: 'user', content: prompt },
-      { role: 'user', content: `Within this HTML: ${minifyHtml(html)}` },
+      { role: gtpRole.user, content: prompt },
+      { role: gtpRole.user, content: `Within this HTML: ${minifyHtml(html)}` },
     ]
 
     const response = await this._processAIRequest(messages)
@@ -121,7 +126,7 @@ class AI extends Helper {
    * @returns {Promise<string>} - A Promise that resolves to the generated response from the GPT model.
    */
   async askGptGeneralPrompt(prompt) {
-    const messages = [{ role: 'user', content: prompt }]
+    const messages = [{ role: gtpRole.user, content: prompt }]
 
     const response = await this._processAIRequest(messages)
 
@@ -156,40 +161,45 @@ class AI extends Helper {
    * @returns {Promise<Object>} A promise that resolves to the requested page object.
    */
   async askForPageObject(pageName, extraPrompt = null, locator = null) {
-    const html = locator ? await this.helper.grabHTMLFrom(locator) : await this.helper.grabSource()
-
     const spinner = ora(' Processing AI request...').start()
-    await this.aiAssistant.setHtmlContext(html)
-    const response = await this.aiAssistant.generatePageObject(extraPrompt, locator)
-    spinner.stop()
 
-    if (!response[0]) {
-      output.error('No response from AI')
-      return ''
-    }
+    try {
+      const html = locator ? await this.helper.grabHTMLFrom(locator) : await this.helper.grabSource()
+      await this.aiAssistant.setHtmlContext(html)
+      const response = await this.aiAssistant.generatePageObject(extraPrompt, locator)
+      spinner.stop()
+
+      if (!response[0]) {
+        output.error('No response from AI')
+        return ''
+      }
 
-    const code = beautify(response[0])
+      const code = beautify(response[0])
 
-    output.print('----- Generated PageObject ----')
-    output.print(code)
-    output.print('-------------------------------')
+      output.print('----- Generated PageObject ----')
+      output.print(code)
+      output.print('-------------------------------')
 
-    const fileName = path.join(output_dir, `${pageName}Page-${Date.now()}.js`)
+      const fileName = path.join(output_dir, `${pageName}Page-${Date.now()}.js`)
 
-    output.print(output.styles.bold(`Page object for ${pageName} is saved to ${output.styles.bold(fileName)}`))
-    fs.writeFileSync(fileName, code)
+      output.print(output.styles.bold(`Page object for ${pageName} is saved to ${output.styles.bold(fileName)}`))
+      fs.writeFileSync(fileName, code)
 
-    try {
-      registerVariable('page', require(fileName))
-      output.success('Page object registered for this session as `page` variable')
-      output.print('Use `=>page.methodName()` in shell to run methods of page object')
-      output.print('Use `click(page.locatorName)` to check locators of page object')
-    } catch (err) {
-      output.error('Error while registering page object')
-      output.error(err.message)
-    }
+      try {
+        registerVariable('page', require(fileName))
+        output.success('Page object registered for this session as `page` variable')
+        output.print('Use `=>page.methodName()` in shell to run methods of page object')
+        output.print('Use `click(page.locatorName)` to check locators of page object')
+      } catch (err) {
+        output.error('Error while registering page object')
+        output.error(err.message)
+      }
 
-    return code
+      return code
+    } catch (e) {
+      spinner.stop()
+      throw Error(`Something went wrong! ${e.message}`)
+    }
   }
 
   async _processAIRequest(messages) {