ai.js

const debug = require('debug')('codeceptjs:ai');
const output = require('./output');
const event = require('./event');
const { removeNonInteractiveElements, minifyHtml, splitByChunks } = require('./html');

const defaultHtmlConfig = {
  maxLength: 50000,
  simplify: true,
  minify: true,
  html: {},
};

const defaultPrompts = {
  writeStep: (html, input) => [{
    role: 'user',
    content: `I am test engineer writing test in CodeceptJS
        I have opened web page and I want to use CodeceptJS to ${input} on this page
        Provide me valid CodeceptJS code to accomplish it
        Use only locators from this HTML: \n\n${html}`,
  },
  ],

  healStep: (html, { step, error, prevSteps }) => {
    return [{
      role: 'user',
      content: `As a test automation engineer I am testing web application using CodeceptJS.
        I want to heal a test that fails. Here is the list of executed steps: ${prevSteps.map(s => s.toString()).join(', ')}
        Propose how to adjust ${step.toCode()} step to fix the test.
        Use locators in order of preference: semantic locator by text, CSS, XPath. Use codeblocks marked with \`\`\`
        Here is the error message: ${error.message}
        Here is HTML code of a page where the failure has happened: \n\n${html}`,
    }];
  },
};

class AiAssistant {
  constructor() {
    this.totalTime = 0;
    this.numTokens = 0;

    this.reset();
    this.connectToEvents();
  }

  enable(config = {}) {
    debug('Enabling AI assistant');
    this.isEnabled = true;

    const { html, prompts, ...aiConfig } = config;

    this.config = Object.assign(this.config, aiConfig);
    this.htmlConfig = Object.assign(defaultHtmlConfig, html);
    this.prompts = Object.assign(defaultPrompts, prompts);

    debug('Config', this.config);
  }

  reset() {
    this.numTokens = 0;
    this.isEnabled = false;
    this.config = {
      maxTokens: 1000000,
      request: null,
      response: parseCodeBlocks,
      // lets limit token usage to 1M
    };
    this.minifiedHtml = null;
    this.response = null;
    this.totalTime = 0;
  }

  disable() {
    this.isEnabled = false;
  }

  connectToEvents() {
    event.dispatcher.on(event.all.result, () => {
      if (this.isEnabled && this.numTokens > 0) {
        const numTokensK = Math.ceil(this.numTokens / 1000);
        const maxTokensK = Math.ceil(this.config.maxTokens / 1000);
        output.print(`AI assistant took ${this.totalTime}s and used ~${numTokensK}K input tokens. Tokens limit: ${maxTokensK}K`);
      }
    });
  }

  checkRequestFn() {
    if (!this.isEnabled) {
      debug('AI assistant is disabled');
      return;
    }

    if (this.config.request) return;

    const noRequestErrorMessage = `
     No request function is set for AI assistant.  
     Please implement your own request function and set it in the config.    

     [!] AI request was decoupled from CodeceptJS. To connect to OpenAI or other AI service, please implement your own request function and set it in the config.

     Example (connect to OpenAI):

     ai: {
       request: async (messages) => {
         const OpenAI = require('openai');
         const openai = new OpenAI({ apiKey: process.env['OPENAI_API_KEY'] })
         const response = await openai.chat.completions.create({
          model: 'gpt-3.5-turbo-0125',
          messages,
         });
         return response?.data?.choices[0]?.message?.content;
       }
     }
    `.trim();

    throw new Error(noRequestErrorMessage);
  }

  async setHtmlContext(html) {
    let processedHTML = html;

    if (this.htmlConfig.simplify) {
      processedHTML = removeNonInteractiveElements(processedHTML, this.htmlConfig);
    }

    if (this.htmlConfig.minify) processedHTML = await minifyHtml(processedHTML);
    if (this.htmlConfig.maxLength) processedHTML = splitByChunks(processedHTML, this.htmlConfig.maxLength)[0];

    this.minifiedHtml = processedHTML;
  }

  getResponse() {
    return this.response || '';
  }

  async createCompletion(messages) {
    if (!this.isEnabled) return '';

    debug('Request', messages);

    this.checkRequestFn();

    this.response = null;

    this.calculateTokens(messages);

    try {
      const startTime = process.hrtime();
      this.response = await this.config.request(messages);
      const endTime = process.hrtime(startTime);
      const executionTimeInSeconds = endTime[0] + endTime[1] / 1e9;

      this.totalTime += Math.round(executionTimeInSeconds);
      debug('AI response time', executionTimeInSeconds);
      debug('Response', this.response);
      this.stopWhenReachingTokensLimit();
      return this.response;
    } catch (err) {
      debug(err.response);
      output.print('');
      output.error(`AI service error: ${err.message}`);
      if (err?.response?.data?.error?.code) output.error(err?.response?.data?.error?.code);
      if (err?.response?.data?.error?.message) output.error(err?.response?.data?.error?.message);
      this.stopWhenReachingTokensLimit();
      return '';
    }
  }

  async healFailedStep(failureContext) {
    if (!this.isEnabled) return [];
    if (!failureContext.html) throw new Error('No HTML context provided');

    await this.setHtmlContext(failureContext.html);

    if (!this.minifiedHtml) {
      debug('HTML context is empty after removing non-interactive elements & minification');
      return [];
    }

    const response = await this.createCompletion(this.prompts.healStep(this.minifiedHtml, failureContext));
    if (!response) return [];

    return this.config.response(response);
  }

  calculateTokens(messages) {
    // we implement naive approach for calculating tokens with no extra requests
    // this approach was tested via https://platform.openai.com/tokenizer
    // we need it to display current usage tokens usage so users could analyze effectiveness of AI

    const inputString = messages.map(m => m.content).join(' ').trim();
    const numWords = (inputString.match(/[^\s\-:=]+/g) || []).length;

    // 2.5 token is constant for average HTML input
    const tokens = numWords * 2.5;

    this.numTokens += tokens;

    return tokens;
  }

  stopWhenReachingTokensLimit() {
    if (this.numTokens < this.config.maxTokens) return;

    output.print(`AI assistant has reached the limit of ${this.config.maxTokens} tokens in this session. It will be disabled now`);
    this.disable();
  }

  async writeSteps(input) {
    if (!this.isEnabled) return;
    if (!this.minifiedHtml) throw new Error('No HTML context provided');

    const snippets = [];

    const response = await this.createCompletion(this.prompts.writeStep(this.minifiedHtml, input));
    if (!response) return;
    snippets.push(...this.config.response(response));

    debug(snippets[0]);

    return snippets[0];
  }
}

function parseCodeBlocks(response) {
  // Regular expression pattern to match code snippets
  const codeSnippetPattern = /```(?:javascript|js|typescript|ts)?\n([\s\S]+?)\n```/g;

  // Array to store extracted code snippets
  const codeSnippets = [];

  response = response.split('\n').map(line => line.trim()).join('\n');

  // Iterate over matches and extract code snippets
  let match;
  while ((match = codeSnippetPattern.exec(response)) !== null) {
    codeSnippets.push(match[1]);
  }

  // Remove "Scenario", "Feature", and "require()" lines
  const modifiedSnippets = codeSnippets.map(snippet => {
    const lines = snippet.split('\n');

    const filteredLines = lines.filter(line => !line.includes('I.amOnPage') && !line.startsWith('Scenario') && !line.startsWith('Feature') && !line.includes('= require('));

    return filteredLines.join('\n');
    // remove snippets that move from current url
  }); // .filter(snippet => !line.includes('I.amOnPage'));

  return modifiedSnippets.filter(snippet => !!snippet);
}

module.exports = new AiAssistant();