Merge pull request #113 from FlowTestAI/anthropic-model

feat: added backend support for anthropic claude via aws bedrock

Author: jsajal

.changeset/lovely-beers-promise.md

@@ -0,0 +1,5 @@
+---
+'flowtestai': minor
+---
+
+Add support for anthropic claude hosted on bedrock

packages/flowtest-electron/package.json

@@ -27,6 +27,15 @@
   },
   "dependencies": {
     "@apidevtools/swagger-parser": "^10.1.0",
+    "@aws-crypto/sha256-js": "^5.2.0",
+    "@aws-sdk/client-bedrock": "^3.583.0",
+    "@aws-sdk/client-bedrock-runtime": "^3.583.0",
+    "@aws-sdk/credential-provider-node": "^3.583.0",
+    "@aws-sdk/types": "^3.577.0",
+    "@smithy/eventstream-codec": "^3.0.0",
+    "@smithy/protocol-http": "^4.0.0",
+    "@smithy/signature-v4": "^3.0.0",
+    "@smithy/util-utf8": "^3.0.0",
     "axios": "^1.6.7",
     "chokidar": "^3.6.0",
     "dotenv": "^16.4.5",

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

@@ -0,0 +1,122 @@
+const BedrockClaudeGenerate = require('./models/bedrock_claude');
+const OpenAIGenerate = require('./models/openai');
+
+class FlowtestAI {
+  async generate(collection, user_instruction, model) {
+    if (model.name === 'OPENAI') {
+      const available_functions = await this.get_available_functions(collection);
+      const openai = new OpenAIGenerate();
+      const functions = await openai.filter_functions(available_functions, user_instruction, model.apiKey);
+      return await openai.process_user_instruction(functions, user_instruction, model.apiKey);
+    } else if (model.name === 'BEDROCK_CLAUDE') {
+      const available_functions = await this.get_available_functions(collection);
+      const bedrock_claude = new BedrockClaudeGenerate(model.apiKey);
+      const functions = await bedrock_claude.filter_functions(available_functions, user_instruction);
+      return await bedrock_claude.process_user_instruction(functions, user_instruction);
+    } else {
+      throw Error(`Model ${model.name} not supported`);
+    }
+  }
+
+  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,
+          };
+        }
+
+        const f = {
+          type: 'function',
+          function: { name: function_name, description: desc, parameters: schema },
+        };
+
+        if (this.isCyclic(f)) {
+          functions.push({
+            type: 'function',
+            function: { name: function_name, description: desc, parameters: {} },
+          });
+        } else {
+          functions.push(f);
+        }
+      });
+    });
+
+    return functions;
+  }
+
+  isCyclic(obj) {
+    var keys = [];
+    var stack = [];
+    var stackSet = new Set();
+    var detected = false;
+
+    function detect(obj, key) {
+      if (obj && typeof obj != 'object') {
+        return false;
+      }
+
+      if (stackSet.has(obj)) {
+        // it's cyclic! Print the object and its locations.
+        var oldindex = stack.indexOf(obj);
+        var l1 = keys.join('.') + '.' + key;
+        var l2 = keys.slice(0, oldindex + 1).join('.');
+        //console.log('CIRCULAR: ' + l1 + ' = ' + l2 + ' = ' + obj);
+        //console.log(obj);
+        detected = true;
+        return;
+      }
+
+      keys.push(key);
+      stack.push(obj);
+      stackSet.add(obj);
+      for (var k in obj) {
+        //dive on the object's children
+        if (Object.prototype.hasOwnProperty.call(obj, k)) {
+          detect(obj[k], k);
+        }
+      }
+
+      keys.pop();
+      stack.pop();
+      stackSet.delete(obj);
+      return;
+    }
+
+    detect(obj, 'obj', keys, stack, stackSet, detected);
+    return detected;
+  }
+}
+
+module.exports = FlowtestAI;

packages/flowtest-electron/src/ai/models/bedrock_claude.js

@@ -0,0 +1,132 @@
+const { BedrockChat } = require('@langchain/community/chat_models/bedrock');
+const { HumanMessage, SystemMessage, BaseMessage } = require('@langchain/core/messages');
+const { BedrockEmbeddings } = require('@langchain/community/embeddings/bedrock');
+const { MemoryVectorStore } = require('langchain/vectorstores/memory');
+
+class BedrockClaudeGenerate {
+  constructor(creds) {
+    this.model = new BedrockChat({
+      model: 'anthropic.claude-3-sonnet-20240229-v1:0',
+      region: 'us-west-2',
+      // endpointUrl: "custom.amazonaws.com",
+      credentials: creds,
+      modelKwargs: {
+        anthropic_version: 'bedrock-2023-05-31',
+      },
+    });
+
+    this.embeddings = new BedrockEmbeddings({
+      region: 'us-west-2',
+      credentials: creds,
+      model: 'amazon.titan-embed-text-v2:0', // Default value
+    });
+  }
+
+  async filter_functions(functions, instruction) {
+    const documents = functions.map((f) => {
+      const { parameters, ...fDescription } = f.function;
+      return JSON.stringify(fDescription);
+    });
+
+    const vectorStore = await MemoryVectorStore.fromTexts(documents, [], this.embeddings);
+    // 128 (max no of functions accepted by openAI function calling)
+    const retrievedDocuments = await vectorStore.similaritySearch(instruction, 10);
+    var selectedFunctions = [];
+    retrievedDocuments.forEach((document) => {
+      const pDocument = JSON.parse(document.pageContent);
+      const findF = functions.find(
+        (f) => f.function.name === pDocument.name && f.function.description === pDocument.description,
+      );
+      if (findF) {
+        selectedFunctions = selectedFunctions.concat(findF);
+      }
+    });
+    return selectedFunctions;
+  }
+
+  async process_user_instruction(functions, instruction) {
+    //console.log(functions.map((f) => f.function.name));
+    // Define the function call format
+    const fn = `{"name": "function_name"}`;
+
+    // Prepare the function string for the system prompt
+    const fnStr = functions.map((f) => JSON.stringify(f)).join('\n');
+
+    // Define the system prompt
+    const systemPrompt = `
+        You are a helpful assistant with access to the following functions:
+
+        ${fnStr}
+
+        To use these functions respond with, only output function names, ignore arguments needed by those functions:
+
+        <multiplefunctions>
+            <functioncall> ${fn} </functioncall>
+            <functioncall> ${fn} </functioncall>
+            ...
+        </multiplefunctions>
+
+        Edge cases you must handle:
+        - If there are multiple functions that can fullfill user request, list them all.
+        - If there are no functions that match the user request, you will respond politely that you cannot help.
+        - If the user has not provided all information to execute the function call, choose the best possible set of values. Only, respond with the information requested and nothing else.
+        - If asked something that cannot be determined with the user's request details, respond that it is not possible to fulfill the request and explain why.
+    `;
+
+    // Prepare the messages for the language model
+    const messages = [new SystemMessage({ content: systemPrompt }), new HumanMessage({ content: instruction })];
+
+    // Invoke the language model and get the completion
+    const completion = await this.model.invoke(messages);
+    const content = completion.content.trim();
+
+    // Extract function calls from the completion
+    const extractedFunctions = this.extractFunctionCalls(content);
+
+    console.log(extractedFunctions);
+
+    return extractedFunctions;
+  }
+
+  extractFunctionCalls(completion) {
+    let content = typeof completion === 'string' ? completion : completion.content;
+
+    // Multiple functions lookup
+    const mfnPattern = /<multiplefunctions>(.*?)<\/multiplefunctions>/s;
+    const mfnMatch = content.match(mfnPattern);
+
+    // Single function lookup
+    const singlePattern = /<functioncall>(.*?)<\/functioncall>/s;
+    const singleMatch = content.match(singlePattern);
+
+    let functions = [];
+
+    if (!mfnMatch && !singleMatch) {
+      // No function calls found
+      return null;
+    } else if (mfnMatch) {
+      // Multiple function calls found
+      const multiplefn = mfnMatch[1];
+      const fnMatches = [...multiplefn.matchAll(/<functioncall>(.*?)<\/functioncall>/gs)];
+      for (let fnMatch of fnMatches) {
+        const fnText = fnMatch[1].replace(/\\/g, '');
+        try {
+          functions.push(JSON.parse(fnText));
+        } catch {
+          // Ignore invalid JSON
+        }
+      }
+    } else {
+      // Single function call found
+      const fnText = singleMatch[1].replace(/\\/g, '');
+      try {
+        functions.push(JSON.parse(fnText));
+      } catch {
+        // Ignore invalid JSON
+      }
+    }
+    return functions;
+  }
+}
+
+module.exports = BedrockClaudeGenerate;

packages/flowtest-electron/src/ai/models/openai.js

@@ -0,0 +1,101 @@
+const OpenAI = require('openai');
+const { MemoryVectorStore } = require('langchain/vectorstores/memory');
+const { OpenAIEmbeddings } = require('@langchain/openai');
+
+const SYSTEM_MESSAGE = `You are a helpful assistant. \ 
+        Respond to the following prompt by using function_call and then summarize actions. \ 
+        If a user request is ambiguous, choose the best response possible.`;
+
+// Maximum number of function calls allowed to prevent infinite or lengthy loops
+const MAX_CALLS = 10;
+
+class OpenAIGenerate {
+  async filter_functions(functions, instruction, apiKey) {
+    const documents = functions.map((f) => {
+      const { parameters, ...fDescription } = f.function;
+      return JSON.stringify(fDescription);
+    });
+
+    const vectorStore = await MemoryVectorStore.fromTexts(
+      documents,
+      [],
+      new OpenAIEmbeddings({
+        openAIApiKey: apiKey,
+      }),
+    );
+
+    // 128 (max no of functions accepted by openAI function calling)
+    const retrievedDocuments = await vectorStore.similaritySearch(instruction, 10);
+    var selectedFunctions = [];
+    retrievedDocuments.forEach((document) => {
+      const pDocument = JSON.parse(document.pageContent);
+      const findF = functions.find(
+        (f) => f.function.name === pDocument.name && f.function.description === pDocument.description,
+      );
+      if (findF) {
+        selectedFunctions = selectedFunctions.concat(findF);
+      }
+    });
+
+    return selectedFunctions;
+  }
+
+  async get_openai_response(functions, messages, apiKey) {
+    const openai = new OpenAI({
+      apiKey,
+    });
+
+    return await openai.chat.completions.create({
+      model: 'gpt-4', //gpt-3.5-turbo-16k-0613
+      tools: functions,
+      tool_choice: '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, apiKey) {
+    //console.log(functions.map((f) => f.function.name));
+    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, apiKey);
+      const message = response['choices'][0]['message'];
+
+      if (message.tool_calls) {
+        messages.push(message);
+        message.tool_calls.map((tool_call) => {
+          console.log('Function call #: ', num_calls + 1);
+          console.log(JSON.stringify(tool_call));
+
+          // We'll simply add a message to simulate successful function call.
+          messages.push({
+            role: 'tool',
+            content: 'success',
+            tool_call_id: tool_call.id,
+          });
+          result.push(tool_call.function);
+
+          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 = OpenAIGenerate;

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

@@ -14,7 +14,7 @@ const createFile = require('../utils/filemanager/createfile');
 const updateFile = require('../utils/filemanager/updatefile');
 const deleteFile = require('../utils/filemanager/deletefile');
 const readFile = require('../utils/filemanager/readfile');
-const FlowtestAI = require('../utils/flowtestai');
+const FlowtestAI = require('../ai/flowtestai');
 const { stringify, parse } = require('flatted');
 const { deserialize, serialize } = require('../utils/flowparser/parser');