support for GPT generated state machines and lambda functions

mhlabs · Mar 29, 2023 · 8198b42 · 8198b42
1 parent cd3259d
commit 8198b42
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 29 deletions.
diff --git a/README.md b/README.md
@@ -58,7 +58,9 @@ Options:
 ![Demo](images/demo2.gif)
 
 ### sam-patterns generate
-Generates SAM resources based on a query to ChatGPT and merges them into your template. This is an experimental feature and requires a ChatGPT API key. You can get one [here](https://platform.openai.com/account/api-keys). Make sure to validate the output before deploying your template as it might contain errors or things that could incur cost 
+Generates SAM resources, CDK code, StepFunctions ASL and Lambda functions in any language based on a query to ChatGPT. If you ask for SAM resources, it will merges them into your existing template. 
+
+This is an experimental feature and requires a ChatGPT API key. You can get one [here](https://platform.openai.com/account/api-keys). Make sure to validate the output before deploying your template as it might contain errors or things that could incur cost 
 
 
 ```
@@ -67,14 +69,24 @@ Usage: sampat generate|g [options]
 Generates resources from a ChatGPT response
 
 Options:
-  -t, --template [template]  SAM template file (default: "template.yaml")
-  -q, --query [query]        Question to ask ChatGPT. I.e "a lambda function that's triggered by an S3 event"
-  -m, --model [model]        OpenAI model to use. Valid values are 'gpt-3.5-turbo' and 'gpt-4'. Note that gpt-3.5-turbo is fine for
-                             most use cases and that gpt-4 is slower and more expensive (default: "gpt-3.5-turbo")
-  -h, --help                 display help for command
+  -t, --template [template]         SAM template file (default: "template.yaml")
+  -q, --query [query]               Question to ask ChatGPT. I.e "a lambda function that's triggered by an S3 event"
+  -m, --model [model]               OpenAI model to use. Valid values are 'gpt-3.5-turbo' and 'gpt-4'. Note that gpt-3.5-turbo is fine for
+                                    most use cases and that gpt-4 is slower and more expensive (default: "gpt-3.5-turbo")
+  -o, --output [output]             Output feature. Valid values are 'SAM', 'CDK', 'lambda-<language>' or 'ASL'. If not 'SAM', set
+                                    --output-file (default: "SAM")
+  -of, --output-file [output-file]  Output file. Only applicable if --output is CDK
+  -h, --help                        display help for command
 ```
 ![Demo](images/demo-gpt.gif)
 
+#### Examples
+* To generate SAM resources for a Lambda function that reads off a DynamoDB table: `sam-patterns generate -q "a lambda function that reads off a dynamodb table"`
+* To generate a CDK stack for the same: `sam-patterns generate -q "a lambda function that reads off a dynamodb table" --output CDK --output-file cdk-stack.ts`
+* To generate a Lambda function in Rust that reads off a DynamoDB table: `sam-patterns generate -q "a lambda function that reads off a dynamodb table" --output lambda-rust --output-file lambda.py`
+* To generate a StepFunctions ASL definition that reads off a DynamoDB table: `sam-patterns generate -q "a lambda function that reads off a dynamodb table" --output ASL --output-file asl.yaml`
+
+Note that quality of results may vary and that you sometimes have to run the command a few times to get a good result.
 
 ### sam-patterns source
 Lets you add more sources. This could be public repositories, such as Jeremy Daly's [Serverless Reference Architectures](https://www.jeremydaly.com/serverless-reference-architectures/) or a private repository for company specific patterns.

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "sam-patterns-cli",
-  "version": "0.0.38",
+  "version": "0.0.39",
   "description": "Command line interface for quickly using patterns from https://github.com/aws-samples/serverless-patterns/",
   "main": "index.js",
   "scripts": {

diff --git a/src/commands/generate/generate.js b/src/commands/generate/generate.js
@@ -4,12 +4,14 @@ const { Configuration, OpenAIApi } = require("openai");
 const settingsUtil = require("../../shared/settingsUtil");
 const fs = require("fs-extra");
 var Spinner = require('cli-spinner').Spinner;
+const baseFile = require("../../shared/baseFile.json");
 
 var spinner = new Spinner('Waiting for ChatGPT... %s');
 spinner.setSpinnerString('|/-\\');
 async function run(cmd) {
   console.log("*** Note: This is an experimental feature and depends on the ChatGPT API. Make sure you review the output carefully before using it in production ***")
-  if (cmd.output.toLowerCase() === "cdk" && !cmd.outputFile) {
+  const output = cmd.output.toLowerCase();
+  if (output !== "sam" && !cmd.outputFile) {
     console.log(`You need to specify an output file with --output-file`);
     return;
   }
@@ -20,27 +22,52 @@ async function run(cmd) {
     );
     return;
   }
-  if (!fs.existsSync(cmd.template)) {
-    console.log(
-      `Can't find ${cmd.template}. Use -t option to specify template filename`
-    );
-    const create = await inputUtil.prompt(`Create ${cmd.template}?`);
-    if (create) {
-      fs.writeFileSync(cmd.template, parser.stringify("yaml", baseFile));
-    } else {
-      return;
+  let ownTemplate;
+  if (output === "sam") {
+    if (!fs.existsSync(cmd.template)) {
+      console.log(
+        `Can't find ${cmd.template}. Use -t option to specify template filename`
+      );
+      const create = await inputUtil.prompt(`Create ${cmd.template}?`);
+      if (create) {
+        fs.writeFileSync(cmd.template, parser.stringify("yaml", baseFile));
+      } else {
+        return;
+      }
     }
-  }
-
-  const ownTemplate = parser.parse("own", fs.readFileSync(cmd.template));
-  ownTemplate.Resources = ownTemplate.Resources || {};
 
+    ownTemplate = parser.parse("own", fs.readFileSync(cmd.template));
+    ownTemplate.Resources = ownTemplate.Resources || {};
+  }
   const configuration = new Configuration({
     apiKey,
   });
   const openai = new OpenAIApi(configuration);
-  const format = cmd.output.toLowerCase() === "sam" ? "SAM JSON" : "TypeScript CDK";
-  const prompt = `Generate this in AWS ${format}: ${cmd.query}`;
+  let outputString;
+  let language;
+  switch (output.substring(0, 3)) {
+    case "sam":
+      outputString = "SAM JSON";
+      break;
+    case "cdk":
+      outputString = "TypeScript CDK";
+      break;
+    case "lam":
+      outputString = "Lambda";
+      language = output.split("-")[1];
+      if (language) {
+        outputString += ` in ${language}`;
+      }
+
+      break;
+    case "asl":
+      outputString = "StepFunctions ASL YAML";
+      break;
+    default:
+      console.log(`Invalid output format ${output}. Valid values are 'SAM', 'CDK', 'lambda-<language>' or 'ASL'`);
+      return;
+  }
+  const prompt = `Generate this in AWS ${outputString}: ${cmd.query}. Only return code.`;
 
   const openAiRequest = {
     model: cmd.model,
@@ -56,10 +83,10 @@ async function run(cmd) {
   spinner.start();
   const response = await openai.createChatCompletion(openAiRequest);
   spinner.stop();
-  const text = response.data.choices[0].message.content;
+  let text = response.data.choices[0].message.content;
   // get the first JSON object in the text
   let obj;
-  if (cmd.output.toLowerCase() === "sam") {
+  if (output === "sam") {
     try {
       obj = JSON.parse(text.replace(/\n/g, '').replace(/```/g, '').match(/{.*}/)[0]);
     } catch (e) {
@@ -76,14 +103,23 @@ async function run(cmd) {
 
     console.log(`\n\nGenerated the following resources:\n\n${parser.stringify("yaml", obj)}`);
   } else {
-    console.log(`\n\nGenerated the following CDK code:\n\n${text}`);
+    console.log(text);
+    //check if text has a row starting with ``` followed by text. If so, remove that text
+    const match = text.match(/```.+\n/);
+    if (match && match[0]) {
+      text = text.replace(match[0], '```\n');
+    }
+    if (text.match(/```/g)) {
+      text = text.replace(/\n/g, "¶").match(/```.*```/)[0].split("¶").join('\n').replace(/```/g, '');
+    }
+    console.log(`\n\nGenerated the following ${language} code:\n\n${text}`);
   }
-  const cont = await inputUtil.prompt(cmd.output.toLowerCase() === "sam" ? `Add to template?` : `Add to ${cmd.outputFile}?`);
+  const cont = await inputUtil.prompt(output === "sam" ? `Add to template?` : `Add to ${cmd.outputFile}?`);
   if (!cont) {
     return;
   }
 
-  if (cmd.output.toLowerCase() === "cdk") {
+  if (output !== "sam") {
     if (!cmd.outputFile) {
       console.log(`You need to specify an output file with --output-file`);
       return;

diff --git a/src/commands/generate/index.js b/src/commands/generate/index.js
@@ -6,7 +6,7 @@ program
   .option("-t, --template [template]", "SAM template file", "template.yaml")
   .option("-q, --query [query]", "Question to ask ChatGPT. I.e \"a lambda function that's triggered by an S3 event\"")
   .option("-m, --model [model]", "OpenAI model to use. Valid values are 'gpt-3.5-turbo' and 'gpt-4'. Note that gpt-3.5-turbo is fine for most use cases and that gpt-4 is slower and more expensive", "gpt-3.5-turbo")
-  .option("-o, --output [output]", "Output language. Valid values SAM or CDK. If CDK, set --output-file", "SAM")
+  .option("-o, --output [output]", "Output feature. Valid values are 'SAM', 'CDK', 'lambda-<language>' or 'ASL'. If not 'SAM', set --output-file", "SAM")
   .option("-of, --output-file [output-file]", "Output file. Only applicable if --output is CDK")
   .description("Generates resources from a ChatGPT response")
   .action(async (cmd) => {