Add flowtest ai component in main process

Author: jsajal

packages/flowtest-electron/package.json

@@ -22,7 +22,9 @@
     "electron-store": "^8.1.0",
     "fs": "^0.0.1-security",
     "json-refs": "^3.0.15",
+    "langchain": "^0.1.28",
     "lodash": "^4.17.21",
+    "openai": "^4.29.1",
     "path": "^0.12.7",
     "uuid": "^9.0.1"
   }

packages/flowtest-electron/src/ipc/collection.js

@@ -15,8 +15,10 @@ const updateFile = require('../utils/filemanager/updatefile');
 const deleteFile = require('../utils/filemanager/deletefile');
 const { flowDataToReadableData, readableDataToFlowData } = require('../utils/parser');
 const readFile = require('../utils/filemanager/readfile');
+const FlowtestAI = require('../utils/flowtestai');
 
 const collectionStore = new Collections();
+const flowTestAI = new FlowtestAI();
 
 const timeout = 60000;
 
@@ -85,7 +87,7 @@ const registerRendererEventHandlers = (mainWindow, watcher) => {
         id: id,
         name: collectionName,
         pathname: pathname,
-        collection: spec,
+        openapi_spec: resolvedSpec.resolved,
         nodes: parsedNodes,
       };
 
@@ -247,6 +249,19 @@ const registerRendererEventHandlers = (mainWindow, watcher) => {
       }
     }
   });
+
+  ipcMain.handle('renderer:create-flowtest-ai', async (event, instruction, collectionId) => {
+    try {
+      const collection = collectionStore.getAll().find((c) => c.id === collectionId);
+      if (collection) {
+        return await flowTestAI.generate(collection.openapi_spec, instruction);
+      } else {
+        return Promise.reject(new Error('Collection not found'));
+      }
+    } catch (error) {
+      return Promise.reject(error);
+    }
+  });
 };
 
 module.exports = registerRendererEventHandlers;

packages/flowtest-electron/src/utils/flowtestai.js

@@ -0,0 +1,153 @@
+const OpenAI = require('openai');
+const dotenv = require('dotenv');
+const { MemoryVectorStore } = require('langchain/vectorstores/memory');
+const { OpenAIEmbeddings } = require('@langchain/openai');
+
+dotenv.config();
+
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
+const SYSTEM_MESSAGE = `You are a helpful assistant. \ 
+        Respond to the following prompt by using function_call and then summarize actions. \ 
+        Ask for clarification if a user request is ambiguous.`;
+
+// Maximum number of function calls allowed to prevent infinite or lengthy loops
+const MAX_CALLS = 10;
+
+class FlowtestAI {
+  async generate(collection, user_instruction) {
+    const available_functions = await this.get_available_functions(collection);
+    const functions = await this.filter_functions(available_functions, user_instruction);
+    return await this.process_user_instruction(functions, user_instruction);
+  }
+
+  async get_available_functions(collection) {
+    let functions = [];
+    Object.entries(collection['paths']).map(([path, methods], index) => {
+      Object.entries(methods).map(([method, spec], index1) => {
+        const function_name = spec['operationId'];
+
+        const desc = spec['description'] || spec['summary'] || '';
+
+        let schema = { type: 'object', properties: {} };
+
+        let req_body = undefined;
+        if (spec['requestBody']) {
+          if (spec['requestBody']['content']) {
+            if (spec['requestBody']['content']['application/json']) {
+              if (spec['requestBody']['content']['application/json']['schema']) {
+                req_body = spec['requestBody']['content']['application/json']['schema'];
+              }
+            }
+          }
+        }
+
+        if (req_body != undefined) {
+          schema['properties']['requestBody'] = req_body;
+        }
+
+        const params = spec['parameters'] ? spec['parameters'] : [];
+        const param_properties = {};
+        if (params.length > 0) {
+          for (const param of params) {
+            if (param['schema']) {
+              param_properties[param['name']] = param['schema'];
+            }
+          }
+          schema['properties']['parameters'] = {
+            type: 'object',
+            properties: param_properties,
+          };
+        }
+
+        functions.push({ name: function_name, description: desc, parameters: schema });
+      });
+    });
+    // console.log(JSON.stringify(functions));
+    return functions;
+  }
+
+  async filter_functions(functions, instruction) {
+    const chunkSize = 32;
+    const chunks = [];
+
+    for (let i = 0; i < functions.length; i += chunkSize) {
+      const chunk = functions.slice(i, i + chunkSize);
+      chunks.push(chunk);
+    }
+
+    const documents = chunks.map((chunk) => JSON.stringify(chunk));
+
+    const vectorStore = await MemoryVectorStore.fromTexts(
+      documents,
+      [],
+      new OpenAIEmbeddings({
+        openAIApiKey: process.env.OPENAI_API_KEY,
+      }),
+    );
+
+    // 32 x 4 = 128 (max no of functions accepted by openAI function calling)
+    const retrievedDocuments = await vectorStore.similaritySearch(instruction, 4);
+    var selectedFunctions = [];
+    retrievedDocuments.forEach((document) => {
+      selectedFunctions = selectedFunctions.concat(JSON.parse(document.pageContent));
+    });
+
+    return selectedFunctions;
+  }
+
+  async get_openai_response(functions, messages) {
+    return await openai.chat.completions.create({
+      model: 'gpt-3.5-turbo-16k-0613',
+      functions: functions,
+      function_call: 'auto', // "auto" means the model can pick between generating a message or calling a function.
+      temperature: 0,
+      messages: messages,
+    });
+  }
+
+  async process_user_instruction(functions, instruction) {
+    let result = [];
+    let num_calls = 0;
+    const messages = [
+      { content: SYSTEM_MESSAGE, role: 'system' },
+      { content: instruction, role: 'user' },
+    ];
+
+    while (num_calls < MAX_CALLS) {
+      const response = await this.get_openai_response(functions, messages);
+      // console.log(response)
+      const message = response['choices'][0]['message'];
+
+      if (message['function_call']) {
+        console.log('Function call #: ', num_calls + 1);
+        console.log(message['function_call']);
+        messages.push(message);
+
+        // We'll simply add a message to simulate successful function call.
+        messages.push({
+          role: 'function',
+          content: 'success',
+          name: message['function_call']['name'],
+        });
+        result.push(message['function_call']);
+
+        num_calls += 1;
+      } else {
+        console.log('Message: ');
+        console.log(message['content']);
+        break;
+      }
+    }
+
+    if (num_calls >= MAX_CALLS) {
+      console.log('Reached max chained function calls: ', MAX_CALLS);
+    }
+
+    return result;
+  }
+}
+
+module.exports = FlowtestAI;

packages/flowtest-electron/tests/utils/flowtest-ai.test.js

@@ -0,0 +1,22 @@
+const fs = require('fs');
+const SwaggerParser = require('@apidevtools/swagger-parser');
+const JsonRefs = require('json-refs');
+const FlowtestAI = require('../../src/utils/flowtestai');
+
+describe('generate', () => {
+  it('should generate functions from openapi spec', async () => {
+    const f = new FlowtestAI();
+    const USER_INSTRUCTION =
+      'Instruction: Add a new pet to the store. \
+            Then get the created pet. \
+            Then get pet with status as available.';
+    //const testYaml = fs.readFileSync('tests/test.yaml', { encoding: 'utf8', flag: 'r' });
+    let api = await SwaggerParser.validate('tests/test.yaml');
+    console.log('API name: %s, Version: %s', api.info.title, api.info.version);
+    const resolvedSpec = (await JsonRefs.resolveRefs(api)).resolved;
+
+    let result = await f.generate(resolvedSpec, USER_INSTRUCTION);
+    const nodeNames = result.map((node) => node.name);
+    expect(nodeNames).toEqual(['addPet', 'getPetById', 'findPetsByStatus']);
+  }, 60000);
+});